Add IsAllocator and IsDeallocator API functions

lib/CppInterOp/CppInterOp.cpp

@@ -53,6 +53,7 @@
 #include "clang/AST/Stmt.h"
 #include "clang/AST/Type.h"
 #include "clang/AST/VTableBuilder.h"
+#include "clang/Basic/Builtins.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticSema.h"
@@ -1580,6 +1581,50 @@
   return INTEROP_RETURN(nullptr);
 }
 
+bool IsAllocator(ConstFuncRef Fn) {
+  INTEROP_TRACE(Fn);
+  if (!Fn)
+    return INTEROP_RETURN(false);
+  const auto* D = unwrap<clang::Decl>(Fn);
+  if (const auto* FD = dyn_cast<FunctionDecl>(D)) {
+    if (FD->getBuiltinID() == Builtin::ID::BImalloc)
+      return INTEROP_RETURN(true);
+    if (const auto* FDA = FD->getAttr<RestrictAttr>()) {
+      if (FDA->getSemanticSpelling() != RestrictAttr::Declspec_restrict)
+        return INTEROP_RETURN(true);
+    }
+
+    if (const auto* FDA = FD->getAttr<OwnershipAttr>()) {
+      if (FDA->getOwnKind() == OwnershipAttr::Returns)
+        return INTEROP_RETURN(true);
+    }
+
+    if (FD->hasAttr<CFReturnsRetainedAttr>() ||
+        FD->hasAttr<NSReturnsRetainedAttr>() ||
+        FD->hasAttr<OSReturnsRetainedAttr>())
+      return INTEROP_RETURN(true);
+  }
+
+  return INTEROP_RETURN(false);
+}
+
+bool IsDeallocator(ConstFuncRef Fn) {
+  INTEROP_TRACE(Fn);
+  if (!Fn)
+    INTEROP_RETURN(false);
+  const auto* D = unwrap<clang::Decl>(Fn);
+  if (const auto* FD = dyn_cast<FunctionDecl>(D)) {
+    if (FD->getBuiltinID() == Builtin::ID::BIfree)
+      return INTEROP_RETURN(true);
+    if (const auto* FDA = FD->getAttr<OwnershipAttr>()) {
+      if (FDA->getOwnKind() == OwnershipAttr::Takes)
+        return INTEROP_RETURN(true);
+    }
+  }
+
+  return INTEROP_RETURN(false);
+}
+
 bool IsFunctionProtoType(ConstTypeRef TyRef) {
   INTEROP_TRACE(TyRef);
   QualType QT = QualType::getFromOpaquePtr(TyRef.data);

lib/CppInterOp/CppInterOp.td

@@ -1262,6 +1262,22 @@ def GetFunctionReturnType : CppInterOpAPI {
   let Args = [Arg<"ConstFuncRef", "func">];
 }
 
+def IsAllocator : CppInterOpAPI {
+  let Doc = "Checks if the provided function is an allocator function";
+  let ReturnType = "bool";
+  let Args = [
+    Arg<"ConstFuncRef", "func">
+  ];
+}
+
+def IsDeallocator : CppInterOpAPI {
+  let Doc = "Checks if the provided function is a deallocator function";
+  let ReturnType = "bool";
+  let Args = [
+    Arg<"ConstFuncRef", "func">
+  ];
+}
+
 def GetFunctionSignature : CppInterOpAPI {
   let Doc = [{\\returns a stringified version of a given function signature in the
 form: void N::f(int i, double d, long l = 0, char ch = 'a').}];

unittests/CppInterOp/APINotes/TestHeader.apinotes

@@ -0,0 +1,6 @@
+Name: TestHeader
+Functions:
+  - Name: testAlloc
+    RetainCountConvention: CFReturnsRetained
+  - Name: testNotAlloc
+    RetainCountConvention: CFReturnsNotRetained

unittests/CppInterOp/APINotes/TestHeader.h

@@ -0,0 +1,7 @@
+#ifndef UNITTESTS_CPPINTEROP_APINOTES_TESTHEADER_H
+#define UNITTESTS_CPPINTEROP_APINOTES_TESTHEADER_H
+
+void* testAlloc(int value);
+void testNotAlloc(void* ptr);
+
+#endif

unittests/CppInterOp/APINotes/module.modulemap

@@ -0,0 +1 @@
+module TestHeader {header "TestHeader.h"}

unittests/CppInterOp/CMakeLists.txt

@@ -85,6 +85,9 @@ if(TARGET CppInterOpTableGen)
   add_dependencies(CAPITestC CppInterOpTableGen)
 endif()
 target_link_libraries(CppInterOpTests PRIVATE CAPITestC)
+target_compile_definitions(CppInterOpTests PRIVATE
+  "CPPINTEROP_DIR=\"${CMAKE_CURRENT_SOURCE_DIR}/../../\""
+)
 
 set_source_files_properties(InterpreterTest.cpp PROPERTIES COMPILE_DEFINITIONS
   "LLVM_BINARY_DIR=\"${LLVM_BINARY_DIR}\""

unittests/CppInterOp/FunctionReflectionTest.cpp

@@ -575,6 +575,86 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_GetFunctionReturnType) {
       "double");
 }
 
+TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsAllocator) {
+  std::vector<Decl*> Decls;
+  std::string code = R"(
+    class Klass{
+      int val;
+    };
+    __attribute__((ownership_returns(malloc)))
+    Klass* Allocator(){
+      Klass* obj = new Klass;
+      return obj;
+    }
+    Klass* __attribute__((malloc)) Allocator2(){
+      Klass* obj = new Klass;
+      return obj;
+    }
+    void foo();
+    )";
+  GetAllTopLevelDecls(code, Decls, true);
+  EXPECT_TRUE(Cpp::IsAllocator(Decls[1]));
+  EXPECT_TRUE(Cpp::IsAllocator(Decls[2]));
+  EXPECT_FALSE(Cpp::IsAllocator(Decls[3]));
+  // Builtin check
+  code = R"(
+  //There is nothing lstdlib.h is included at args
+  )";
+  TestFixture::CreateInterpreter({"-include", "stdlib.h"});
+  Interp->process(code);
+  auto mallocDecl = Cpp::GetNamed("malloc");
+  EXPECT_TRUE(Cpp::IsAllocator(Cpp::ConstFuncRef{mallocDecl.data}));
+  Cpp::DeleteInterpreter();
+
+  // APINotes check
+#if !defined(CPPINTEROP_USE_CLING) && !defined(__EMSCRIPTEN__)
+  std::string include_flag =
+      "-I" + std::string(CPPINTEROP_DIR) + "unittests/CppInterOp/APINotes";
+  std::vector<const char*> interpreter_args = {
+      "-fmodules", "-fimplicit-module-maps", "-fapinotes-modules",
+      include_flag.c_str()};
+  TestFixture::CreateInterpreter(interpreter_args);
+  code = R"(
+  #include "TestHeader.h"
+  )";
+  Interp->process(code);
+  auto testAllocDecl = Cpp::GetNamed("testAlloc");
+  EXPECT_TRUE(Cpp::IsAllocator(Cpp::ConstFuncRef{testAllocDecl.data}));
+
+  auto testNotAllocDecl = Cpp::GetNamed("testNotAlloc");
+  EXPECT_FALSE(Cpp::IsAllocator(Cpp::ConstFuncRef{testNotAllocDecl.data}));
+  Cpp::DeleteInterpreter();
+#endif
+}
+
+TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsDeallocator) {
+  std::vector<Decl*> Decls;
+  std::string code = R"(
+    class Klass{
+      int val;
+    };
+    __attribute__((ownership_takes(malloc, 1)))
+    void Deallocator(Klass* arg){
+      delete arg;
+    }
+    void foo();
+    )";
+  GetAllTopLevelDecls(code, Decls, true);
+  EXPECT_TRUE(Cpp::IsDeallocator(Decls[1]));
+  EXPECT_FALSE(Cpp::IsDeallocator(Decls[2]));
+
+  code = R"(
+  #include <stdlib.h>
+    void test(){
+      //Do Nothing
+    }
+  )";
+  TestFixture::CreateInterpreter();
+  Interp->process(code);
+  auto freeDecl = Cpp::GetNamed("free");
+  EXPECT_TRUE(Cpp::IsDeallocator(Cpp::ConstFuncRef{freeDecl.data}));
+}
+
 TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_GetFunctionNumArgs) {
   std::vector<Decl*> Decls, TemplateSubDecls;
   std::string code = R"(