From b3cc8289956b70e36383d614bb1f98b5eca5be69 Mon Sep 17 00:00:00 2001
From: Giles Payne <gilespayne@gmail.com>
Date: Mon, 20 Sep 2021 22:44:13 +0900
Subject: [PATCH] Fix put/get functions for non-contiguous Mat on iOS/macOS

---
 modules/core/misc/objc/common/Mat.mm      | 95 ++++++++++++++---------
 modules/core/misc/objc/test/MatTest.swift | 22 +++++-
 2 files changed, 78 insertions(+), 39 deletions(-)

diff --git a/modules/core/misc/objc/common/Mat.mm b/modules/core/misc/objc/common/Mat.mm
index 045bd8393e..dc2316c87b 100644
--- a/modules/core/misc/objc/common/Mat.mm
+++ b/modules/core/misc/objc/common/Mat.mm
@@ -13,25 +13,28 @@
 #import "CvType.h"
 #import "CVObjcUtil.h"
 
-// return true if we have reached the final index
-static bool incIdx(cv::Mat* mat, std::vector<int>& indices) {
-    for (int dim = mat->dims-1; dim>=0; dim--) {
-        indices[dim] = (indices[dim] + 1) % mat->size[dim];
-        if (indices[dim] != 0) {
-            return false;
-        }
+static int idx2Offset(cv::Mat* mat, std::vector<int>& indices) {
+    int offset = indices[0];
+    for (int dim=1; dim < mat->dims; dim++) {
+        offset = offset*mat->size[dim] + indices[dim];
     }
-    return true;
+    return offset;
 }
 
-// returns true if final index was reached
-static bool updateIdx(cv::Mat* mat, std::vector<int>& indices, int inc) {
-    for (int index = 0; index < inc; index++) {
-        if (incIdx(mat, indices)) {
-            return true;
-        }
+static void offset2Idx(cv::Mat* mat, size_t offset, std::vector<int>& indices) {
+    for (int dim=mat->dims-1; dim>=0; dim--) {
+        indices[dim] = offset % mat->size[dim];
+        offset = (offset - indices[dim]) / mat->size[dim];
     }
-    return false;
+}
+
+// returns true if final index was reached
+static bool updateIdx(cv::Mat* mat, std::vector<int>& indices, size_t inc) {
+    size_t currentOffset = idx2Offset(mat, indices);
+    size_t newOffset = currentOffset + inc;
+    bool reachedEnd = newOffset>=(size_t)mat->total();
+    offset2Idx(mat, reachedEnd?0:newOffset, indices);
+    return reachedEnd;
 }
 
 @implementation Mat {
@@ -724,22 +727,31 @@ template<typename T> int getData(NSArray<NSNumber*>* indices, cv::Mat* mat, int
     }
 
     int arrayAvailable = count;
+    size_t countBytes = count * sizeof(T);
+    size_t remainingBytes = (size_t)(mat->total() - idx2Offset(mat, tempIndices))*mat->elemSize();
+    countBytes = (countBytes>remainingBytes)?remainingBytes:countBytes;
+    int result = (int)countBytes;
     int matAvailable = getMatAvailable(mat, tempIndices);
     int available = MIN(arrayAvailable, matAvailable);
-    int result = (int)(available * mat->elemSize() / mat->channels());
     if (mat->isContinuous()) {
         memcpy(tBuffer, mat->ptr(tempIndices.data()), available * sizeof(T));
     } else {
-        int copyOffset = 0;
-        int copyCount = MIN((mat->size[mat->dims - 1] - tempIndices[mat->dims - 1]) * mat->channels(), available);
-        while (available > 0) {
-            memcpy(tBuffer + copyOffset, mat->ptr(tempIndices.data()), copyCount * sizeof(T));
-            if (updateIdx(mat, tempIndices, copyCount / mat->channels())) {
-                break;
-            }
-            available -= copyCount;
-            copyOffset += copyCount * sizeof(T);
-            copyCount = MIN(mat->size[mat->dims-1] * mat->channels(), available);
+        char* buff = (char*)tBuffer;
+        size_t blockSize = mat->size[mat->dims-1] * mat->elemSize();
+        size_t firstPartialBlockSize = (mat->size[mat->dims-1] - tempIndices[mat->dims-1]) * mat->step[mat->dims-1];
+        for (int dim=mat->dims-2; dim>=0 && blockSize == mat->step[dim]; dim--) {
+            blockSize *= mat->size[dim];
+            firstPartialBlockSize += (mat->size[dim] - (tempIndices[dim]+1)) * mat->step[dim];
+        }
+        size_t copyCount = (countBytes<firstPartialBlockSize)?countBytes:firstPartialBlockSize;
+        uchar* data = mat->ptr(tempIndices.data());
+        while(countBytes>0) {
+            memcpy(buff, data, copyCount);
+            updateIdx(mat, tempIndices, copyCount / mat->elemSize());
+            countBytes -= copyCount;
+            buff += copyCount;
+            copyCount = countBytes<blockSize?countBytes:blockSize;
+            data = mat->ptr(tempIndices.data());
         }
     }
     return result;
@@ -817,22 +829,31 @@ template<typename T> int putData(NSArray<NSNumber*>* indices, cv::Mat* mat, int
     }
 
     int arrayAvailable = count;
+    size_t countBytes = count * sizeof(T);
+    size_t remainingBytes = (size_t)(mat->total() - idx2Offset(mat, tempIndices))*mat->elemSize();
+    countBytes = (countBytes>remainingBytes)?remainingBytes:countBytes;
+    int result = (int)countBytes;
     int matAvailable = getMatAvailable(mat, tempIndices);
     int available = MIN(arrayAvailable, matAvailable);
-    int result = (int)(available * mat->elemSize() / mat->channels());
     if (mat->isContinuous()) {
         memcpy(mat->ptr(tempIndices.data()), tBuffer, available * sizeof(T));
     } else {
-        int copyOffset = 0;
-        int copyCount = MIN((mat->size[mat->dims - 1] - tempIndices[mat->dims - 1]) * mat->channels(), available);
-        while (available > 0) {
-            memcpy(mat->ptr(tempIndices.data()), tBuffer + copyOffset, copyCount * sizeof(T));
-            if (updateIdx(mat, tempIndices, copyCount / mat->channels())) {
-                break;
-            }
-            available -= copyCount;
-            copyOffset += copyCount * sizeof(T);
-            copyCount = MIN(mat->size[mat->dims-1] * (int)mat->channels(), available);
+        char* buff = (char*)tBuffer;
+        size_t blockSize = mat->size[mat->dims-1] * mat->elemSize();
+        size_t firstPartialBlockSize = (mat->size[mat->dims-1] - tempIndices[mat->dims-1]) * mat->step[mat->dims-1];
+        for (int dim=mat->dims-2; dim>=0 && blockSize == mat->step[dim]; dim--) {
+            blockSize *= mat->size[dim];
+            firstPartialBlockSize += (mat->size[dim] - (tempIndices[dim]+1)) * mat->step[dim];
+        }
+        size_t copyCount = (countBytes<firstPartialBlockSize)?countBytes:firstPartialBlockSize;
+        uchar* data = mat->ptr(tempIndices.data());
+        while(countBytes>0){
+            memcpy(data, buff, copyCount);
+            updateIdx(mat, tempIndices, copyCount / mat->elemSize());
+            countBytes -= copyCount;
+            buff += copyCount;
+            copyCount = countBytes<blockSize?countBytes:blockSize;
+            data = mat->ptr(tempIndices.data());
         }
     }
     return result;
diff --git a/modules/core/misc/objc/test/MatTest.swift b/modules/core/misc/objc/test/MatTest.swift
index 8a513505cc..2dcfedf41f 100644
--- a/modules/core/misc/objc/test/MatTest.swift
+++ b/modules/core/misc/objc/test/MatTest.swift
@@ -431,7 +431,7 @@ class MatTests: OpenCVTestCase {
         // whole Mat
         var bytesNum = try m.get(row: 1, col: 1, data: &buff)
 
-        XCTAssertEqual(12, bytesNum);
+        XCTAssertEqual(12, bytesNum)
         XCTAssert(buff == [110, 111, 120, 121, 130, 131])
 
         // sub-Mat
@@ -442,8 +442,26 @@ class MatTests: OpenCVTestCase {
         XCTAssert(buff00 == [230, 231, 240, 241])
         var buff11 = [Int16](repeating: 0, count: 4)
         bytesNum = try sm.get(row: 1, col: 1, data: &buff11)
-        XCTAssertEqual(4, bytesNum);
+        XCTAssertEqual(4, bytesNum)
         XCTAssert(buff11 == [340, 341, 0, 0])
+
+        let m2 = Mat(sizes: [5, 6, 8], type: CvType.CV_16S)
+        let data:[Int16] = (0..<m2.total()).map { Int16($0) }
+        try m2.put(indices: [0, 0, 0], data: data)
+        let matNonContinuous = m2.submat(ranges:[Range(start:1, end:4), Range(start:2, end:5), Range(start:3, end:6)])
+        let matContinuous = matNonContinuous.clone()
+        var outNonContinuous = [Int16](repeating:0, count: matNonContinuous.total())
+        try matNonContinuous.get(indices: [0, 0, 0], data: &outNonContinuous)
+        var outContinuous = [Int16](repeating: 0, count:matNonContinuous.total())
+        try matContinuous.get(indices: [0, 0, 0], data: &outContinuous)
+        XCTAssertEqual(outNonContinuous, outContinuous)
+        let subMat2 = m2.submat(ranges:[Range(start:1, end:4), Range(start:1, end:5), Range(start:0, end:8)])
+        let subMatClone2 = subMat2.clone()
+        var outNonContinuous2 = [Int16](repeating:0, count: subMat2.total())
+        try subMat2.get(indices: [0, 1, 1], data: &outNonContinuous2)
+        var outContinuous2 = [Int16](repeating:0, count:subMat2.total())
+        try subMatClone2.get(indices: [0, 1, 1], data: &outContinuous2)
+        XCTAssertEqual(outNonContinuous2, outContinuous2)
     }
 
     func testGetIntIntUInt16Array() throws {