[WIP] On demand analyze for allocation

CMakeLists.txt

@@ -11,8 +11,8 @@ include(GNUInstallDirs)
 # This option won't make a lot of sense since we only ship the shared library in site-packages
 # Perhaps this should permanently be OFF and users can build their own CppInterOp if they want to run the tests?
 option(CPPJIT_ENABLE_CPPINTEROP_TESTS "enable CppInterOp tests" OFF)
-set(CPPINTEROP_GIT_REPOSITORY "https://github.com/compiler-research/CppInterOp.git" CACHE STRING "")
-set(CPPINTEROP_GIT_TAG "8d624c621a4b95e36ff73ac708c85a768287478f" CACHE STRING "")
+set(CPPINTEROP_GIT_REPOSITORY "https://github.com/keremsahn/CppInterOp.git" CACHE STRING "")
+set(CPPINTEROP_GIT_TAG "attr-design" CACHE STRING "")
 set(CPPINTEROP_SOURCE_DIR "" CACHE PATH
     "Override default CppInterOp built by ExternalProject_Add, with a path to local CppInterOp source")
 

python/cppjit/__init__.py

@@ -48,6 +48,7 @@
     "add_library_path",  # add a path to search for libraries
     "add_autoload_map",  # explicitly include an autoload map
     "set_debug",  # enable/disable debug output
+    "use_alloc_analyzer",  # enable/disable memory ownership analyzer
 ]
 
 import ctypes
@@ -397,6 +398,11 @@ def set_debug(enable=True):
     gbl.Cpp.EnableDebugOutput(enable)
 
 
+def use_alloc_analyzer(enable=True):
+    """Enable/disable memory ownership analyzer"""
+    _backend.UseAllocAnalyzer(enable)
+
+
 def _get_name(tt):
     if isinstance(tt, str):
         return tt

src/cpyrt/CPPMethod.cxx

@@ -34,6 +34,7 @@ extern PyObject* gBusException;
 extern PyObject* gSegvException;
 extern PyObject* gIllException;
 extern PyObject* gAbrtException;
+extern bool gUseAllocAnalyzer;
 } // namespace cppjit::cpyrt
 
 //- public helper ------------------------------------------------------------
@@ -752,6 +753,19 @@ 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::AllocType attrResult = interop::IsAllocator(GetMethod());
+  if (attrResult == interop::AllocType::Unknown && gUseAllocAnalyzer) {
+    interop::AllocType analyzeResult = interop::GetAllocType(GetMethod());
+    fAllocType = analyzeResult;
+    return analyzeResult;
+  }
+  fAllocType = attrResult;
+  return attrResult;
+}
 //----------------------------------------------------------------------------
 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,7 @@ class CPPMethod : public PyCallable {
   PyObject* GetCoVarNames() override;
   PyObject* GetArgDefault(int iarg, bool silent = true) override;
   bool IsConst() override;
+  cppjit::interop::AllocType GetAllocBehaviour() override;
 
   PyObject* GetScopeProxy() override;
   interop::TCppFuncAddr_t GetFunctionAddress() override;
@@ -116,6 +118,7 @@ class CPPMethod : public PyCallable {
 protected:
   // cached value that doubles as initialized flag (uninitialized if -1)
   int fArgsRequired;
+  std::optional<cppjit::interop::AllocType> fAllocType;
 };
 
 } // namespace cppjit::cpyrt

src/cpyrt/CPPOverload.cxx

@@ -155,6 +155,12 @@ static inline PyObject* HandleReturn(CPPOverload* pymeth, CPPInstance* im_self,
     CPPInstance* cppres =
         (CPPInstance*)(CPPInstance_Check(result) ? result : nullptr);
 
+    interop::AllocType AT =
+        pymeth->fMethodInfo->fMethods[0]->GetAllocBehaviour();
+    if (AT != interop::AllocType::None && AT != interop::AllocType::Null &&
+        AT != interop::AllocType::Unknown)
+      pymeth->fMethodInfo->fFlags |= CallContext::kIsCreator;
+
     // if this method creates new objects, always take ownership
     if (IsCreator(pymeth->fMethodInfo->fFlags)) {
 

src/cpyrt/PyCallable.h

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

src/cpyrt/cpyrtModule.cxx

@@ -279,6 +279,7 @@ PyObject* gAbrtException = nullptr;
 std::unordered_set<interop::TCppScope_t> gPinnedTypes;
 std::ostringstream gCapturedError;
 std::streambuf* gOldErrorBuffer = nullptr;
+bool gUseAllocAnalyzer = false;
 
 std::unordered_map<std::string, std::vector<PyObject*>>& pythonizations() {
   static std::unordered_map<std::string, std::vector<PyObject*>> pyzMap;
@@ -1012,6 +1013,20 @@ static PyObject* EndCaptureStderr(PyObject*, PyObject*) {
 
   return Py_BuildValue("s", capturedError.c_str());
 }
+
+//----------------------------------------------------------------------------
+static PyObject* UseAllocAnalyzer(PyObject*, PyObject* args) {
+  // Set allocation-analyzer policy, disabled by default
+  // Usage: enabling ->SetUseAllocAnalyzer(True) / SetUseAllocAnalyzer(1)
+  // disabling ->SetUseAllocAnalyzer(False) / SetUseAllocAnalyzer(0)
+  int enable = 0;
+  if (!PyArg_ParseTuple(args, const_cast<char*>("p"), &enable))
+    return nullptr;
+
+  gUseAllocAnalyzer = enable;
+
+  Py_RETURN_NONE;
+}
 } // unnamed namespace
 
 //- data -----------------------------------------------------------------------
@@ -1061,6 +1076,8 @@ static PyMethodDef gcpyrtMethods[] = {
      METH_NOARGS, (char*)"Begin capturing stderr to a in memory buffer."},
     {(char*)"_end_capture_stderr", (PyCFunction)EndCaptureStderr, METH_NOARGS,
      (char*)"End capturing stderr and returns the captured buffer."},
+    {(char*)"UseAllocAnalyzer", (PyCFunction)UseAllocAnalyzer, METH_VARARGS,
+     (char*)"Enable/disable memory-allocation analyzer."},
     {nullptr, nullptr, 0, nullptr}};
 
 struct module_state {

src/interop/cppjit_interop.h

@@ -44,6 +44,7 @@ typedef Cpp::FuncRef TCppMethod_t;
 typedef Cpp::InterpRef TInterp_t;
 typedef size_t TCppIndex_t;
 typedef void* TCppFuncAddr_t;
+typedef Cpp::AllocType AllocType;
 
 // direct interpreter access -------------------------------------------------
 RPY_EXPORTED
@@ -297,6 +298,10 @@ RPY_EXPORTED
 std::string GetDoxygenComment(TCppScope_t scope, bool strip_markers = true);
 RPY_EXPORTED
 bool IsConstMethod(TCppMethod_t);
+RPY_EXPORTED
+AllocType IsAllocator(TCppMethod_t);
+RPY_EXPORTED
+AllocType GetAllocType(TCppMethod_t);
 // Templated method/function reflection information
 // ------------------------------------
 RPY_EXPORTED

src/interop/interop_wrapper.cxx

@@ -1200,6 +1200,16 @@ interop::TCppType_t interop::GetMethodReturnType(TCppMethod_t method) {
   return Cpp::GetFunctionReturnType(method);
 }
 
+interop::AllocType interop::IsAllocator(TCppMethod_t method) {
+  std::lock_guard<std::recursive_mutex> Lock(InterOpMutex);
+  return Cpp::IsAllocator(method);
+}
+
+interop::AllocType interop::GetAllocType(TCppMethod_t method) {
+  std::lock_guard<std::recursive_mutex> Lock(InterOpMutex);
+  return Cpp::GetAllocType(method);
+}
+
 std::string interop::GetMethodReturnTypeAsString(TCppMethod_t method) {
   std::lock_guard<std::recursive_mutex> Lock(InterOpMutex);
   return Cpp::GetTypeAsString(

test/Makefile

@@ -10,6 +10,7 @@ dictnames = advancedcpp \
             doc_helper \
             example01 \
             fragile \
+            memory_analysis \
             operators \
             overloads \
             pythonizables \

test/cpp/MemoryOwnership/MemOwnrship.apinotes

@@ -0,0 +1,9 @@
+Name: MemOwnrship
+Functions:
+    - Name: memOwnAllocGlobal
+      SwiftReturnOwnership: cppAllocNew
+Tags:
+    - Name: memOwn
+      Methods:
+        - Name: memOwnAllocator
+          SwiftReturnOwnership: cppAllocNew

test/cpp/MemoryOwnership/memory_analysis_redecl.h

@@ -0,0 +1,10 @@
+#ifndef MEMORY_ANALYSIS_REDECL_H
+#define MEMORY_ANALYSIS_REDECL_H
+#include "../memory_analysis.h"
+
+namespace memory {
+[[clang::annotate("cppAllocNew")]]
+memOwn* allocDefaultMemOwn();
+}
+
+#endif

test/cpp/MemoryOwnership/module.modulemap

@@ -0,0 +1 @@
+module MemOwnrship { header "../memory_analysis.h" }

test/cpp/memory_analysis.cxx

@@ -0,0 +1,24 @@
+#include "memory_analysis.h"
+
+namespace memory {
+
+__attribute__((malloc)) memAnalysisKlass* mallocAttr() {
+  return new memAnalysisKlass;
+}
+
+__attribute__((ownership_returns(malloc))) memAnalysisKlass*
+ownershipReturnsAttr() {
+  return new memAnalysisKlass;
+}
+
+// Expected to not return ownership when analysis is off, and there is just
+// attr-check
+memAnalysisKlass* noAttr() { return new memAnalysisKlass; }
+
+memOwn* memOwnAllocGlobal() { return (memOwn*)malloc(sizeof(memOwn)); }
+
+memOwn* allocDefaultMemOwn() { return new memOwn; }
+
+memOwn* noAttrAlloc() { return new memOwn; }
+
+} // namespace memory

test/cpp/memory_analysis.h

@@ -0,0 +1,38 @@
+#ifndef MEMORY_ANALYSIS_H
+#define MEMORY_ANALYSIS_H
+
+#include <new>
+#include <stdlib.h>
+namespace memory {
+
+class memAnalysisKlass {
+public:
+  int val;
+};
+__attribute__((malloc)) memAnalysisKlass* mallocAttr();
+__attribute__((ownership_returns(malloc))) memAnalysisKlass*
+ownershipReturnsAttr();
+memAnalysisKlass* noAttr();
+
+struct memOwn {
+  int val;
+  memOwn(int value) : val(value) {}
+  memOwn() { val = 0; }
+  // Attribute injected by APINotes
+  static memOwn* memOwnAllocator(int x) { return new memOwn(x); }
+};
+
+// Attribute injected by APINotes
+memOwn* memOwnAllocGlobal();
+
+// Attribute injected by redeclaration
+memOwn* allocDefaultMemOwn();
+
+// No ownership attribute anywhere
+memOwn* noAttrAlloc();
+
+inline memAnalysisKlass* allocAnalyzerOn() { return new memAnalysisKlass; }
+inline memAnalysisKlass* allocAnalyzerOff() { return new memAnalysisKlass; }
+} // namespace memory
+
+#endif // MEMORY_ANALYSIS_H