[memory-ownership] Honour ownership attributes when taking and releasing memory

src/cpyrt/CPPClassMethod.cxx

@@ -32,15 +32,11 @@ PyObject* cpyrt::CPPClassMethod::Call(CPPInstance*& self, cpyrt_PyArgs_t args,
   // first argument matches the enclosing scope of the class method and it isn't
   // needed for the call, then assume that the method was (incorrectly) bound
   // and so drop that instance from args.
-  int nargs = (int)cpyrt_PyArgs_GET_SIZE(args, nargsf);
-  if ((!self || (PyObject*)self == Py_None) && nargs) {
-    PyObject* arg0 = cpyrt_PyArgs_GET_ITEM(args, 0);
-    if (CPPInstance_Check(arg0) && fArgsRequired <= nargs - 1 &&
-        interop::IsSubclass(reinterpret_cast<CPPInstance*>(arg0)->ObjectIsA(),
-                            GetScope())) {
-      args += 1; // drops first argument
-      nargsf -= 1;
-    }
+  Py_ssize_t nargs = cpyrt_PyArgs_GET_SIZE(args, nargsf);
+  const bool isBound = self && (PyObject*)self != Py_None;
+  if (ShouldDropSelf(isBound, args, nargs)) {
+    args += 1; // drops first argument
+    nargsf -= 1;
   }
 
   if (!this->ConvertAndSetArgs(args, nargsf, ctxt))
@@ -50,6 +46,21 @@ PyObject* cpyrt::CPPClassMethod::Call(CPPInstance*& self, cpyrt_PyArgs_t args,
   return this->Execute(nullptr, 0, ctxt);
 }
 
+//----------------------------------------------------------------------------
+bool cpyrt::CPPClassMethod::ShouldDropSelf(bool isBound, PyObject* const* args,
+                                           Py_ssize_t nargs) {
+  // A bound call keeps self outside of args
+  // Also if there is no arg in call, there is no self to drop
+  if (isBound || nargs == 0)
+    return false;
+
+  if (CPPInstance_Check(args[0]) && fArgsRequired <= nargs - 1 &&
+      interop::IsSubclass(((CPPInstance*)args[0])->ObjectIsA(), GetScope()))
+    return true;
+
+  return false;
+}
+
 //----------------------------------------------------------------------------
 PyObject* cpyrt::CPPClassMethod::GetTypeName() {
   PyObject* cppname =

src/cpyrt/CPPClassMethod.h

@@ -14,6 +14,8 @@ class CPPClassMethod : public CPPMethod {
   PyObject* GetTypeName() override;
 
 public:
+  bool ShouldDropSelf(bool isBound, PyObject* const* args, Py_ssize_t nargs);
+
   PyCallable* Clone() override { return new CPPClassMethod(*this); }
   PyObject* Call(CPPInstance*& self, cpyrt_PyArgs_t args, size_t nargsf,
                  PyObject* kwds, CallContext* ctxt = nullptr) override;

src/cpyrt/CPPInstance.cxx

@@ -17,6 +17,7 @@ using namespace cppjit;
 
 // Standard
 #include <algorithm>
+#include <cstdlib>
 #include <sstream>
 
 //- data _____________________________________________________________________
@@ -227,7 +228,17 @@ void cpyrt::op_dealloc_nofree(CPPInstance* pyobj) {
     if (pyobj->fFlags & CPPInstance::kIsValue) {
       interop::CallDestructor(klass, cppobj);
       interop::Deallocate(klass, cppobj);
-    } else
+    } else if (pyobj->fFlags & CPPInstance::kIsMalloc)
+      std::free(cppobj);
+    else if (pyobj->fFlags & CPPInstance::kIsNoConstruct) {
+      if (pyobj->fFlags & CPPInstance::kIsArrayAlloc)
+        ::operator delete[](cppobj);
+      else
+        ::operator delete(cppobj);
+    } else if (pyobj->fFlags & CPPInstance::kIsArrayAlloc)
+      interop::Destruct(klass, cppobj, 1);
+    // Default case: just kIsOwner set in all of memory-ownership flags
+    else
       interop::Destruct(klass, cppobj);
   }
   cppobj = nullptr;
@@ -938,10 +949,30 @@ static int op_setownership(CPPInstance* pyobj, PyObject* value, void*) {
   return 0;
 }
 
+// Added for testing purposes
+//-----------------------------------------------------------------------------
+static PyObject* op_get_array_alloc(CPPInstance* pyobj, void*) {
+  return PyBool_FromLong((long)(pyobj->fFlags & CPPInstance::kIsArrayAlloc));
+}
+//-----------------------------------------------------------------------------
+static PyObject* op_get_no_construct(CPPInstance* pyobj, void*) {
+  return PyBool_FromLong((long)(pyobj->fFlags & CPPInstance::kIsNoConstruct));
+}
+//-----------------------------------------------------------------------------
+static PyObject* op_get_malloc(CPPInstance* pyobj, void*) {
+  return PyBool_FromLong((long)(pyobj->fFlags & CPPInstance::kIsMalloc));
+}
 //-----------------------------------------------------------------------------
 static PyGetSetDef op_getset[] = {
     {(char*)"__python_owns__", (getter)op_getownership, (setter)op_setownership,
      (char*)"If true, python manages the life time of this object", nullptr},
+    {(char*)"__is_array_alloc__", (getter)op_get_array_alloc, nullptr,
+     (char*)"If true, the object was allocated with new[]/operator new[]",
+     nullptr},
+    {(char*)"__is_no_construct__", (getter)op_get_no_construct, nullptr,
+     (char*)"If true, the memory is raw: no constructor was run", nullptr},
+    {(char*)"__is_malloc__", (getter)op_get_malloc, nullptr,
+     (char*)"If true, the object was allocated with malloc", nullptr},
     {(char*)nullptr, nullptr, nullptr, nullptr, nullptr}};
 
 //= cpyrt type number stubs to allow dynamic overrides =====================

src/cpyrt/CPPInstance.h

@@ -29,20 +29,24 @@ class CPPInstance {
 public:
   enum EFlags {
     kDefault = 0x0000,
-    kNoWrapConv = 0x0001,  // use type as-is (eg. no smart ptr wrap)
-    kIsOwner = 0x0002,     // Python instance owns C++ object/memory
-    kIsExtended = 0x0004,  // has extended data
-    kIsValue = 0x0008,     // was created from a by-value return
-    kIsReference = 0x0010, // represents one indirection
-    kIsArray = 0x0020,     // represents an array of objects
-    kIsSmartPtr = 0x0040,  // is or embeds a smart pointer
-    kIsPtrPtr = 0x0080,    // represents two indirections
-    kIsRValue = 0x0100,    // can be used as an r-value
-    kIsLValue = 0x0200,    // can be used as an l-value
-    kNoMemReg = 0x0400,    // do not register with memory regulator
-    kIsRegulated = 0x0800, // is registered with memory regulator
-    kIsActual = 0x1000,    // has been downcasted to actual type
-    kHasLifeLine = 0x2000, // has a life line set
+    kNoWrapConv = 0x0001,   // use type as-is (eg. no smart ptr wrap)
+    kIsOwner = 0x0002,      // Python instance owns C++ object/memory
+    kIsExtended = 0x0004,   // has extended data
+    kIsValue = 0x0008,      // was created from a by-value return
+    kIsReference = 0x0010,  // represents one indirection
+    kIsArray = 0x0020,      // represents an array of objects
+    kIsSmartPtr = 0x0040,   // is or embeds a smart pointer
+    kIsPtrPtr = 0x0080,     // represents two indirections
+    kIsRValue = 0x0100,     // can be used as an r-value
+    kIsLValue = 0x0200,     // can be used as an l-value
+    kNoMemReg = 0x0400,     // do not register with memory regulator
+    kIsRegulated = 0x0800,  // is registered with memory regulator
+    kIsActual = 0x1000,     // has been downcasted to actual type
+    kHasLifeLine = 0x2000,  // has a life line set
+    kIsArrayAlloc = 0x4000, // represents a heap allocated array of objects
+    kIsNoConstruct =
+        0x8000, // represents constructor is not called in the allocation
+    kIsMalloc = 0x10000, // is allocated with malloc
   };
 
 public: // public, as the python C-API works with C structs

src/cpyrt/CPPMethod.cxx

@@ -85,6 +85,8 @@ inline void cpyrt::CPPMethod::Copy_(const CPPMethod& /* other */) {
   fExecutor = nullptr;
   fArgIndices = nullptr;
   fArgsRequired = -1;
+  fAllocType.reset();
+  fDeallocType.reset();
 }
 
 //----------------------------------------------------------------------------
@@ -703,6 +705,37 @@ PyObject* cpyrt::CPPMethod::GetArgDefault(int iarg, bool silent) {
 
 bool cpyrt::CPPMethod::IsConst() { return interop::IsConstMethod(GetMethod()); }
 
+//----------------------------------------------------------------------------
+interop::AllocType cpyrt::CPPMethod::GetAllocBehaviour() {
+  if (fAllocType.has_value())
+    return *fAllocType;
+  interop::TCppMethod_t method = GetMethod();
+  if (!method) {
+    fAllocType = interop::AllocType::None;
+    return interop::AllocType::None;
+  }
+  if (!IsPointerType(interop::GetMethodReturnType(method))) {
+    fAllocType = interop::AllocType::None;
+    return interop::AllocType::None;
+  }
+  interop::AllocType attrResult = interop::IsAllocator(method);
+  fAllocType = attrResult;
+  return attrResult;
+}
+
+//----------------------------------------------------------------------------
+interop::DeallocType cpyrt::CPPMethod::GetDeallocBehaviour(size_t index) {
+  if (!fDeallocType) {
+    fDeallocType.emplace();
+    if (interop::TCppMethod_t method = GetMethod())
+      interop::IsDeallocator(method, *fDeallocType);
+  }
+  if (index < fDeallocType->size()) {
+    return (*fDeallocType)[index];
+  }
+  return interop::DeallocType::Opaque;
+}
+
 //----------------------------------------------------------------------------
 PyObject* cpyrt::CPPMethod::GetScopeProxy() {
   // Get or build the scope of this method.

src/cpyrt/CPPMethod.h

@@ -5,6 +5,7 @@
 #include "PyCallable.h"
 
 // Standard
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -62,6 +63,8 @@ class CPPMethod : public PyCallable {
   PyObject* GetCoVarNames() override;
   PyObject* GetArgDefault(int iarg, bool silent = true) override;
   bool IsConst() override;
+  interop::AllocType GetAllocBehaviour() override;
+  interop::DeallocType GetDeallocBehaviour(size_t index) override;
 
   PyObject* GetScopeProxy() override;
   interop::TCppFuncAddr_t GetFunctionAddress() override;
@@ -121,6 +124,8 @@ class CPPMethod : public PyCallable {
 protected:
   // cached value that doubles as initialized flag (uninitialized if -1)
   int fArgsRequired;
+  std::optional<interop::AllocType> fAllocType;
+  std::optional<std::vector<interop::DeallocType>> fDeallocType;
 };
 
 } // namespace cppjit::cpyrt

src/cpyrt/CPPOverload.cxx

@@ -2,6 +2,8 @@
 #include "cpyrt.h"
 
 using namespace cppjit;
+#include "CPPClassMethod.h"
+#include "CPPFunction.h"
 #include "CPPInstance.h"
 #include "CPPOverload.h"
 #include "CallContext.h"
@@ -142,14 +144,23 @@ static inline void ResetCallState(CPPInstance* descr_self,
 
 // helper to factor out return logic of mp_call / mp_vectorcall
 static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
-                                     PyObject* result) {
+                                     PyObject* result, PyCallable* method,
+                                     PyObject* const* args, size_t nargsf,
+                                     PyObject* kwds) {
   // special case for python exceptions, propagated through C++ layer
   if (result) {
     CPPInstance* cppres =
         (CPPInstance*)(CPPInstance_Check(result) ? result : nullptr);
 
+    interop::AllocType AT = method->GetAllocBehaviour();
+    // the overloads of a group share fMethodInfo, so the answer of this one
+    // stays in a local instead of going into the shared flags
+    const bool allocates = AT != interop::AllocType::None &&
+                           AT != interop::AllocType::Null &&
+                           AT != interop::AllocType::Unknown;
+
     // if this method creates new objects, always take ownership
-    if (IsCreator(pymeth->fMethodInfo->fFlags)) {
+    if (allocates || IsCreator(pymeth->fMethodInfo->fFlags)) {
 
       // either be a constructor with a fresh object proxy self ...
       if (IsConstructor(pymeth->fMethodInfo->fFlags)) {
@@ -158,9 +169,98 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
       }
 
       // ... or be a regular method with an object proxy return value
-      else if (cppres)
+      else if (cppres) {
         cppres->PythonOwns();
+        // After giving ownership, set proper flags to indicate allocation
+        // method/func
+        switch (AT) {
+        case interop::AllocType::Malloc:
+          cppres->fFlags |= CPPInstance::kIsMalloc;
+          break;
+        case interop::AllocType::NewArr:
+          cppres->fFlags |= CPPInstance::kIsArrayAlloc;
+          break;
+        case interop::AllocType::OperatorNew:
+          cppres->fFlags |= CPPInstance::kIsNoConstruct;
+          break;
+        case interop::AllocType::OperatorNewArr:
+          cppres->fFlags |= CPPInstance::kIsArrayAlloc;
+          cppres->fFlags |= CPPInstance::kIsNoConstruct;
+          break;
+        default:
+          break;
+        }
+      }
+    }
+
+    // an argument that the called function releases belongs to C++ from now
+    // on, else Python frees it a second time
+    const Py_ssize_t nargs = cpyrt_PyArgs_GET_SIZE(args, nargsf);
+    const Py_ssize_t numKwds =
+        kwds && PyTuple_CheckExact(kwds) ? PyTuple_GET_SIZE(kwds) : 0;
+    // the keyword values follow the positional ones in call order, so their
+    // names are what tell which parameter each one belongs to
+    PyObject* argnames = numKwds ? method->GetSignatureNames() : nullptr;
+
+    const bool isBound = pymeth->fSelf != nullptr;
+    // Tells if there is an implicit `this`
+    Py_ssize_t paramShift = 0;
+    // Tells if self parameter should be dropped (may be bound function or
+    // static with self in) or added
+    Py_ssize_t argShift = 0;
+    if (dynamic_cast<CPPFunction*>(method)) {
+      // For a plain func, only case there is a self is when it is bound
+      argShift = isBound ? 1 : 0;
+      // There is no implicit `this` in plain functions
+      paramShift = 0;
+    } else if (auto* CPPCM = dynamic_cast<CPPClassMethod*>(method)) {
+      // Check ShouldDropSelf
+      argShift = CPPCM->ShouldDropSelf(isBound, args, nargs) ? -1 : 0;
+      // Static methods do not carry implicit `this`
+      paramShift = 0;
+    }
+    // Only remain possibility is CPPMethod and its other children which acts as
+    // CPPMethod
+    else {
+      // If method is bound, self is carried outside of args
+      argShift = isBound ? 1 : 0;
+      // A non-static method always keeps implicit `this`
+      paramShift = 1;
+    }
+
+    // the array holds the positional values first, then the keyword ones
+    for (Py_ssize_t i = 0; i < nargs + numKwds; i++) {
+      // only a C++ object can hold memory that the callee releases
+      if (!CPPInstance_Check(args[i]) ||
+          !(((CPPInstance*)args[i])->fFlags & CPPInstance::kIsOwner))
+        continue;
+
+      Py_ssize_t slot = -1; // no slot of its own
+      if (i < nargs)
+        slot = i + argShift;
+      else {
+        // The place of a keyword value means nothing, its name picks the
+        // parameter. kwds has the names in the same order, so kwds[i - nargs]
+        // names args[i].
+        for (Py_ssize_t iarg = 0; iarg < PyTuple_GET_SIZE(argnames); iarg++) {
+          if (PyUnicode_Compare(PyTuple_GET_ITEM(argnames, iarg),
+                                PyTuple_GET_ITEM(kwds, i - nargs)) == 0) {
+            slot = iarg + paramShift;
+            break;
+          }
+        }
+      }
+      // the callee drops this argument, or no name of it matches a parameter
+      if (slot < 0)
+        continue;
+
+      const interop::DeallocType DT = method->GetDeallocBehaviour((size_t)slot);
+      if (DT == interop::DeallocType::None ||
+          DT == interop::DeallocType::Opaque)
+        continue;
+      ((CPPInstance*)args[i])->CppOwns();
     }
+    Py_XDECREF(argnames);
 
     // if this new object falls inside self, make sure its lifetime is proper
     if (!(pymeth->fMethodInfo->fFlags & CallContext::kNeverLifeLine)) {
@@ -576,7 +676,8 @@ static PyObject* mp_vectorcall(CPPOverload* pymeth, PyObject* const* args,
     if (!NoImplicit(&ctxt))
       ctxt.fFlags |= CallContext::kAllowImplicit; // no two rounds needed
     PyObject* result = methods[0]->Call(im_self, args, nargsf, kwds, &ctxt);
-    return HandleReturn(pymeth, im_self, result);
+    return HandleReturn(pymeth, im_self, result, methods[0], args, nargsf,
+                        kwds);
   }
 
   // otherwise, handle overloading
@@ -599,7 +700,8 @@ static PyObject* mp_vectorcall(CPPOverload* pymeth, PyObject* const* args,
       ctxt.fFlags |= CallContext::kAllowImplicit;
     PyObject* result = memoized_pc->Call(im_self, args, nargsf, kwds, &ctxt);
     if (result)
-      return HandleReturn(pymeth, im_self, result);
+      return HandleReturn(pymeth, im_self, result, memoized_pc, args, nargsf,
+                          kwds);
 
     // fall through: python is dynamic, and so, the hashing isn't infallible
     ctxt.fFlags &= ~(CallContext::kAllowImplicit | CallContext::kPyException |
@@ -656,7 +758,8 @@ static PyObject* mp_vectorcall(CPPOverload* pymeth, PyObject* const* args,
           }
         }
 
-        return HandleReturn(pymeth, im_self, result);
+        return HandleReturn(pymeth, im_self, result, methods[i], args, nargsf,
+                            kwds);
       }
 
       // else failure ..

src/cpyrt/PyCallable.h

@@ -37,6 +37,12 @@ class PyCallable {
   virtual PyObject* GetCoVarNames() = 0;
   virtual PyObject* GetArgDefault(int /* iarg */, bool silent = true) = 0;
   virtual bool IsConst() { return false; }
+  virtual interop::AllocType GetAllocBehaviour() {
+    return interop::AllocType::None;
+  }
+  virtual interop::DeallocType GetDeallocBehaviour(size_t /* index */) {
+    return interop::DeallocType::Opaque;
+  }
 
   virtual PyObject* GetScopeProxy() = 0;
   virtual interop::TCppFuncAddr_t GetFunctionAddress() = 0;