From a9a45da2f37fae7445841ee0627e74726122dd47 Mon Sep 17 00:00:00 2001
From: Austin Kearns <59812315+Lillenne@users.noreply.github.com>
Date: Thu, 8 Aug 2024 22:00:27 -0700
Subject: [PATCH] Track managed array pinning in Mat

---
 src/OpenCvSharp/Modules/core/Mat/Mat.cs       | 14 ++--
 .../Modules/core/Mat/MatPinning.cs            | 70 +++++++++++++++++++
 2 files changed, 79 insertions(+), 5 deletions(-)
 create mode 100644 src/OpenCvSharp/Modules/core/Mat/MatPinning.cs

diff --git a/src/OpenCvSharp/Modules/core/Mat/Mat.cs b/src/OpenCvSharp/Modules/core/Mat/Mat.cs
index a2c21bd1c..544acc363 100644
--- a/src/OpenCvSharp/Modules/core/Mat/Mat.cs
+++ b/src/OpenCvSharp/Modules/core/Mat/Mat.cs
@@ -115,6 +115,7 @@ protected Mat(Mat m)
 
         NativeMethods.HandleException(
             NativeMethods.core_Mat_new12(m.ptr, out ptr));
+        pinLifetime = m.pinLifetime?.Ref();
         if (ptr == IntPtr.Zero)
             throw new OpenCvSharpException("imread failed.");
     }
@@ -312,10 +313,10 @@ public static Mat FromPixelData(int rows, int cols, MatType type, IntPtr data, l
     /// If the parameter is missing (set to AUTO_STEP ), no padding is assumed and the actual step is calculated as cols*elemSize() .</param>
     protected Mat(int rows, int cols, MatType type, Array data, long step = 0)
     {
-        var handle = AllocGCHandle(data);
+        pinLifetime = new ArrayPinningLifetime(data);
         NativeMethods.HandleException(
             NativeMethods.core_Mat_new8(rows, cols, type,
-                handle.AddrOfPinnedObject(), new IntPtr(step), out ptr));
+                pinLifetime.DataPtr, new IntPtr(step), out ptr));
     }
 
     /// <summary>
@@ -390,7 +391,7 @@ protected Mat(IEnumerable<int> sizes, MatType type, Array data, IEnumerable<long
         if (data is null)
             throw new ArgumentNullException(nameof(data));
 
-        var handle = AllocGCHandle(data);
+        pinLifetime = new ArrayPinningLifetime(data);
 #pragma warning disable CA1508
         var sizesArray = sizes as int[] ?? sizes.ToArray();
 #pragma warning restore CA1508
@@ -398,14 +399,14 @@ protected Mat(IEnumerable<int> sizes, MatType type, Array data, IEnumerable<long
         {
             NativeMethods.HandleException(
                 NativeMethods.core_Mat_new9(sizesArray.Length, sizesArray,
-                    type, handle.AddrOfPinnedObject(), IntPtr.Zero, out ptr));
+                    type, pinLifetime.DataPtr, IntPtr.Zero, out ptr));
         }
         else
         {
             var stepsArray = steps.Select(s => new IntPtr(s)).ToArray();
             NativeMethods.HandleException(
                 NativeMethods.core_Mat_new9(sizesArray.Length, sizesArray,
-                    type, handle.AddrOfPinnedObject(), stepsArray, out ptr));
+                    type, pinLifetime.DataPtr, stepsArray, out ptr));
         }
     }
 
@@ -2472,6 +2473,9 @@ public Mat SubMat(int rowStart, int rowEnd, int colStart, int colEnd)
             NativeMethods.core_Mat_subMat1(ptr, rowStart, rowEnd, colStart, colEnd, out var ret));
         GC.KeepAlive(this);
         var retVal = new Mat(ret);
+
+        // If this is a managed array, keep the array pinned as long as the Mat is alive
+        retVal.pinLifetime = pinLifetime?.Ref();
         return retVal;
     }
 
diff --git a/src/OpenCvSharp/Modules/core/Mat/MatPinning.cs b/src/OpenCvSharp/Modules/core/Mat/MatPinning.cs
new file mode 100644
index 000000000..e43ce07fd
--- /dev/null
+++ b/src/OpenCvSharp/Modules/core/Mat/MatPinning.cs
@@ -0,0 +1,70 @@
+using System.Runtime.InteropServices;
+
+namespace OpenCvSharp;
+
+public partial class Mat
+{
+    private ArrayPinningLifetime? pinLifetime;
+
+    /// <inheritdoc/>
+    protected override void DisposeManaged()
+    {
+        pinLifetime?.Dispose();
+        base.DisposeManaged();
+    }
+
+    /// <summary>
+    /// Pins an array and unpins it when all references are released.
+    /// </summary>
+    private class ArrayPinningLifetime : IDisposable
+    {
+        private GCHandle handle;
+        private int refCount;
+
+        /// <summary>
+        /// Creates an instance of the <see cref="ArrayPinningLifetime"/> class.
+        /// </summary>
+        /// <param name="array">The array to be pinned.</param>
+        public ArrayPinningLifetime(Array array)
+        {
+            handle = GCHandle.Alloc(array, GCHandleType.Pinned);
+        }
+
+        /// <summary>
+        /// Gets the data pointer of the pinned <see cref="Array"/>.
+        /// </summary>
+        /// <exception cref="ObjectDisposedException">Thrown when the handle has been deallocated.</exception>
+        public IntPtr DataPtr => handle.IsAllocated ? handle.AddrOfPinnedObject() : throw new ObjectDisposedException(nameof(ArrayPinningLifetime));
+
+        /// <summary>Increments the reference count of the pinned array.</summary>
+        /// <returns>Returns a reference to this <see cref="ArrayPinningLifetime"/> instance.</returns>
+        public ArrayPinningLifetime Ref()
+        {
+            Interlocked.Increment(ref refCount);
+            return this;
+        }
+
+        /// <summary>
+        /// Decrements the reference count of the pinned array. If the reference count reaches zero, the array will be
+        /// unpinned.
+        /// </summary>
+        public void Dispose()
+        {
+            if (Interlocked.Decrement(ref refCount) != 0 || !handle.IsAllocated)
+            {
+                return;
+            }
+
+            handle.Free();
+            GC.SuppressFinalize(this);
+        }
+
+        ~ArrayPinningLifetime()
+        {
+            if (handle.IsAllocated)
+            {
+                handle.Free();
+            }
+        }
+    }
+}