[memory-ownership] Automate memory management of LowLevelViews

src/cpyrt/CPPOverload.cxx

@@ -7,6 +7,7 @@ using namespace cppjit;
 #include "CPPInstance.h"
 #include "CPPOverload.h"
 #include "CallContext.h"
+#include "LowLevelViews.h"
 #include "PyStrings.h"
 #include "Utility.h"
 #include "structmember.h" // from Python
@@ -162,6 +163,10 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
     // if this method creates new objects, always take ownership
     if (allocates || IsCreator(pymeth->fMethodInfo->fFlags)) {
 
+      // Note: for CPPInstance returning functions, if IsCreator(..) returns
+      // true it will be assumed allocation is done with `new`, and for
+      // LowLevelView objects, it will be assumed allocator is `new[]`
+
       // either be a constructor with a fresh object proxy self ...
       if (IsConstructor(pymeth->fMethodInfo->fFlags)) {
         if (im_self)
@@ -190,6 +195,24 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
         default:
           break;
         }
+      } else if (LowLevelView_Check(result)) {
+        LowLevelView* LL = (LowLevelView*)result;
+        auto& flags = (intptr_t&)LL->fBufInfo.internal;
+        flags |= LowLevelView::kIsOwner;
+        switch (AT) {
+        case interop::AllocType::Malloc:
+          flags &= ~LowLevelView::kIsCppArray;
+          break;
+        case interop::AllocType::New:
+        case interop::AllocType::OperatorNew:
+          flags |= LowLevelView::kIsNew;
+          flags &= ~LowLevelView::kIsCppArray;
+          break;
+        // ll_new assumes LL is allocated with newarr default, so no need to
+        // touch
+        default:
+          break;
+        }
       }
     }
 
@@ -230,9 +253,15 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
 
     // the array holds the positional values first, then the keyword ones
     for (Py_ssize_t i = 0; i < nargs + numKwds; i++) {
+      const bool isOwnedCppIns =
+          CPPInstance_Check(args[i]) &&
+          (((CPPInstance*)args[i])->fFlags & CPPInstance::kIsOwner);
+      const bool isOwnedLL =
+          (LowLevelView_Check(args[i]) &&
+           ((intptr_t)((LowLevelView*)args[i])->fBufInfo.internal &
+            LowLevelView::kIsOwner));
       // only a C++ object can hold memory that the callee releases
-      if (!CPPInstance_Check(args[i]) ||
-          !(((CPPInstance*)args[i])->fFlags & CPPInstance::kIsOwner))
+      if (!(isOwnedCppIns || isOwnedLL))
         continue;
 
       Py_ssize_t slot = -1; // no slot of its own
@@ -258,7 +287,11 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
       if (DT == interop::DeallocType::None ||
           DT == interop::DeallocType::Opaque)
         continue;
-      ((CPPInstance*)args[i])->CppOwns();
+      if (isOwnedCppIns)
+        ((CPPInstance*)args[i])->CppOwns();
+      else
+        (intptr_t&)((LowLevelView*)args[i])->fBufInfo.internal &=
+            ~LowLevelView::kIsOwner;
     }
     Py_XDECREF(argnames);
 

src/cpyrt/LowLevelViews.cxx

@@ -11,6 +11,7 @@ using namespace cppjit;
 #include <assert.h>
 #include <limits.h>
 #include <map>
+#include <new>
 #include <string.h>
 
 //= memoryview-like object ===================================================
@@ -66,9 +67,11 @@ static void ll_dealloc(cpyrt::LowLevelView* pyobj) {
   PyMem_Free(pyobj->fBufInfo.strides);
   if ((intptr_t)pyobj->fBufInfo.internal & cpyrt::LowLevelView::kIsOwner) {
     if ((intptr_t)pyobj->fBufInfo.internal & cpyrt::LowLevelView::kIsCppArray)
-      delete[] pyobj->fBuf;
+      ::operator delete[](pyobj->get_buf());
+    else if ((intptr_t)pyobj->fBufInfo.internal & cpyrt::LowLevelView::kIsNew)
+      ::operator delete(pyobj->get_buf());
     else
-      free(pyobj->fBuf);
+      free(pyobj->get_buf());
   }
 
   if (pyobj->fElemCnv != pyobj->fConverter && pyobj->fElemCnv &&
@@ -107,6 +110,7 @@ static void ll_dealloc(cpyrt::LowLevelView* pyobj) {
 // clang-format off
 CPYRT_LL_FLAG_GETSET(ownership, cpyrt::LowLevelView::kIsOwner,    __python_owns__)
 CPYRT_LL_FLAG_GETSET(cpparray,  cpyrt::LowLevelView::kIsCppArray, __cpp_array__)
+CPYRT_LL_FLAG_GETSET(isnew, cpyrt::LowLevelView::kIsNew, __is_new__)
 // clang-format on
 
 //---------------------------------------------------------------------------
@@ -916,6 +920,8 @@ static PyGetSetDef ll_getset[] = {
     {(char*)"format", (getter)ll_typecode, nullptr, nullptr, nullptr},
     {(char*)"typecode", (getter)ll_typecode, nullptr, nullptr, nullptr},
     {(char*)"shape", (getter)ll_shape, (setter)ll_setshape, nullptr, nullptr},
+    {(char*)"__is_new__", (getter)ll_getisnew, (setter)ll_setisnew,
+     (char*)"If true, this buffer is allocated with new", nullptr},
     {(char*)nullptr, nullptr, nullptr, nullptr, nullptr}};
 
 namespace cppjit::cpyrt {

src/cpyrt/LowLevelViews.h

@@ -19,8 +19,9 @@ class LowLevelView {
     kDefault = 0x0000,
     kIsCppArray = 0x0001, // allocated with new[]
     kIsFixed = 0x0002,    // fixed size array (assumed flat)
-    kIsOwner = 0x0004
-  }; // Python owns
+    kIsOwner = 0x0004,    // Python owns
+    kIsNew = 0x0008       // memory is allocated with new
+  };
 
 public:
   PyObject_HEAD Py_buffer fBufInfo;

test/cpp/memory_analysis.cxx

@@ -71,4 +71,26 @@ void memOwner::releaseOverloaded(memOwn* keep, memOwn* obj) { delete obj; }
 
 void releaseFree(memOwner* self, memOwn* obj) { delete obj; }
 
+int* allocIntNewArr(int count) { return new int[count]; }
+
+int* allocIntOperatorNewArr(size_t count) {
+  return (int*)::operator new[](count * sizeof(int));
+}
+
+int* allocIntNew() { return new int; }
+
+int* allocIntOperatorNew() { return (int*)::operator new(sizeof(int)); }
+
+int* allocIntMalloc(size_t count) { return (int*)malloc(count * sizeof(int)); }
+
+int* intNoAttrAlloc() { return new int[5]; }
+
+void releaseBufferArr(int* buf) { delete[] buf; }
+
+void releaseBuffer(int* buf) { delete buf; }
+
+void releaseBufferMalloc(int* buf) { free(buf); }
+
+void keepBuffer(int* buf) { return; }
+
 } // namespace memory

test/cpp/memory_analysis.h

@@ -90,6 +90,36 @@ struct memOwner {
 
 [[clang::annotate("cppDeallocDelete", 2)]]
 void releaseFree(memOwner* self, memOwn* obj);
-} // namespace memory
 
+// A raw pointer to a builtin type binds to a LowLevelView, not an instance
+[[clang::annotate("cppAllocNewArr")]]
+int* allocIntNewArr(int count);
+
+[[clang::annotate("cppAllocOperatorNewArr")]]
+int* allocIntOperatorNewArr(size_t count);
+
+[[clang::annotate("cppAllocNew")]]
+int* allocIntNew();
+
+[[clang::annotate("cppAllocOperatorNew")]]
+int* allocIntOperatorNew();
+
+[[clang::annotate("cppAllocMalloc")]]
+int* allocIntMalloc(size_t count);
+
+// No ownership attribute anywhere
+int* intNoAttrAlloc();
+
+[[clang::annotate("cppDeallocDeleteArr", 1)]]
+void releaseBufferArr(int* buf);
+
+[[clang::annotate("cppDeallocDelete", 1)]]
+void releaseBuffer(int* buf);
+
+[[clang::annotate("cppDeallocFree", 1)]]
+void releaseBufferMalloc(int* buf);
+
+[[clang::annotate("cppDeallocNone", 1)]]
+void keepBuffer(int* buf);
+} // namespace memory
 #endif // MEMORY_ANALYSIS_H

test/test_memoryanalysis.py

@@ -477,6 +477,156 @@ def test38_bind_object_keeps_the_allocation_kind(self):
         gc.collect()
 
 
+class TestLOWLEVELVIEWMEMORY:
+    def setup_class(cls):
+        cls.test_dct = test_dct
+        import cppjit
+
+        cppjit.add_include_path(str(currpath.join("cpp", "MemoryOwnership")))
+        cppjit.include("../memory_analysis.h")
+        cls.memory_analysis = cppjit.load_library(cls.test_dct + ".so")
+
+    def test01_newarr_attr(self):
+        import cppjit
+        import cppjit.types
+
+        view = cppjit.gbl.memory.allocIntNewArr(4)
+        assert type(view) == cppjit.types.LowLevelView
+        assert view.__python_owns__
+        assert view.__cpp_array__
+        assert not view.__is_new__
+
+        # the view frees the buffer with ::operator delete[]
+        del view
+        gc.collect()
+
+    def test02_operator_newarr_attr(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntOperatorNewArr(4)
+        assert view.__python_owns__
+        assert view.__cpp_array__
+        assert not view.__is_new__
+
+        del view
+        gc.collect()
+
+    def test03_new_attr(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntNew()
+        assert view.__python_owns__
+        assert view.__is_new__
+        assert not view.__cpp_array__
+
+        # the view frees the buffer with ::operator delete
+        del view
+        gc.collect()
+
+    def test04_operator_new_attr(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntOperatorNew()
+        assert view.__python_owns__
+        assert view.__is_new__
+        assert not view.__cpp_array__
+
+        del view
+        gc.collect()
+
+    def test05_malloc_attr(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntMalloc(4)
+        assert view.__python_owns__
+        # malloc has no flag of its own: no allocation flag means free()
+        assert not view.__cpp_array__
+        assert not view.__is_new__
+
+        del view
+        gc.collect()
+
+    def test06_no_attr(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.intNoAttrAlloc()
+        assert not view.__python_owns__
+
+        # function has no attribute, so the flags stay as ll_new wrote them:
+        # the user takes ownership, therefore the default delete[] is called
+        view.__python_owns__ = True
+        assert view.__cpp_array__
+        assert not view.__is_new__
+
+        del view
+        gc.collect()
+
+    def test07_release_newarr_buffer(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntNewArr(4)
+        assert view.__python_owns__
+
+        cppjit.gbl.memory.releaseBufferArr(view)
+        assert not view.__python_owns__
+
+        del view
+        gc.collect()
+
+    def test08_release_newarr_buffer_keyword(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntNewArr(4)
+
+        cppjit.gbl.memory.releaseBufferArr(buf=view)
+        assert not view.__python_owns__
+
+    def test09_release_new_buffer(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntNew()
+
+        cppjit.gbl.memory.releaseBuffer(view)
+        assert not view.__python_owns__
+
+        del view
+        gc.collect()
+
+    def test10_release_malloc_buffer(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntMalloc(4)
+
+        cppjit.gbl.memory.releaseBufferMalloc(view)
+        assert not view.__python_owns__
+
+        del view
+        gc.collect()
+
+    def test11_keep_buffer(self):
+        import cppjit
+
+        view = cppjit.gbl.memory.allocIntNewArr(4)
+
+        # a cppDeallocNone parameter leaves the ownership where it was
+        cppjit.gbl.memory.keepBuffer(view)
+        assert view.__python_owns__
+
+        del view
+        gc.collect()
+
+    def test12_release_with_free(self):
+        import cppjit
+
+        view = cppjit.gbl.malloc(16)
+
+        cppjit.gbl.free(view)
+        assert not view.__python_owns__
+
+        del view
+        gc.collect()
+
+
 @mark.skipif(not IN_CHILD, reason="needs " + FLAGS)
 class TestMEMORYANALYSIS_APINOTES:
     def setup_class(cls):