IsAllocator redesigned and intercommunicating with GetAllocType

include/CppInterOp/CppInterOpTypes.h

@@ -90,6 +90,8 @@ typedef struct CppConstFuncRef {
 
 #else // __cplusplus
 
+#include <optional>
+
 namespace Cpp {
 
 struct DeclRef {
@@ -411,11 +413,10 @@ enum class AllocType : unsigned char {
   New,
   NewArr,
   Malloc,
-  Unknown,
-  CustomAlloc,
-  Null,
   OperatorNew,
-  OperatorNewArr
+  OperatorNewArr,
+  Null,
+  Unknown
 };
 
 enum class DeallocType : unsigned char {

lib/CppInterOp/CppInterOp.cpp

@@ -1767,31 +1767,56 @@ TypeRef GetFunctionReturnType(ConstFuncRef func) {
   return INTEROP_RETURN(nullptr);
 }
 
-bool IsAllocator(ConstFuncRef Fn) {
+AllocType IsAllocator(ConstFuncRef Fn) {
   INTEROP_TRACE(Fn);
   if (!Fn)
-    return INTEROP_RETURN(false);
-  const auto* D = unwrap<clang::Decl>(Fn);
+    return INTEROP_RETURN(AllocType::Unknown);
+  const auto* D = UnwrapUsingShadowToFunction(unwrap<clang::Decl>(Fn));
+  if (const auto* FTD = dyn_cast<FunctionTemplateDecl>(D))
+    D = FTD->getTemplatedDecl();
   if (const auto* FD = dyn_cast<FunctionDecl>(D)) {
     if (FD->getBuiltinID() == Builtin::ID::BImalloc)
-      return INTEROP_RETURN(true);
+      return INTEROP_RETURN(AllocType::Malloc);
     if (const auto* FDA = FD->getAttr<RestrictAttr>()) {
       if (FDA->getSemanticSpelling() != RestrictAttr::Declspec_restrict)
-        return INTEROP_RETURN(true);
+        return INTEROP_RETURN(AllocType::Malloc);
     }
 
     if (const auto* FDA = FD->getAttr<OwnershipAttr>()) {
       if (FDA->getOwnKind() == OwnershipAttr::Returns)
-        return INTEROP_RETURN(true);
+        return INTEROP_RETURN(AllocType::Malloc);
     }
 
     if (FD->hasAttr<CFReturnsRetainedAttr>() ||
         FD->hasAttr<NSReturnsRetainedAttr>() ||
         FD->hasAttr<OSReturnsRetainedAttr>())
-      return INTEROP_RETURN(true);
+      return INTEROP_RETURN(AllocType::Malloc);
+
+    for (const auto* attr : FD->attrs()) {
+      llvm::StringRef attrName;
+      if (const auto* swiftAttr = dyn_cast<clang::SwiftAttrAttr>(attr))
+        attrName = swiftAttr->getAttribute();
+      else if (const auto* annotateAttr = dyn_cast<clang::AnnotateAttr>(attr))
+        attrName = annotateAttr->getAnnotation();
+      else
+        continue;
+      attrName.consume_front("returns_");
+      if (attrName == "cppAllocNone")
+        return INTEROP_RETURN(AllocType::None);
+      if (attrName == "cppAllocNew")
+        return INTEROP_RETURN(AllocType::New);
+      if (attrName == "cppAllocNewArr")
+        return INTEROP_RETURN(AllocType::NewArr);
+      if (attrName == "cppAllocMalloc")
+        return INTEROP_RETURN(AllocType::Malloc);
+      if (attrName == "cppAllocOperatorNew")
+        return INTEROP_RETURN(AllocType::OperatorNew);
+      if (attrName == "cppAllocOperatorNewArr")
+        return INTEROP_RETURN(AllocType::OperatorNewArr);
+    }
   }
 
-  return INTEROP_RETURN(false);
+  return INTEROP_RETURN(AllocType::Unknown);
 }
 
 bool IsDeallocator(ConstFuncRef Fn) {
@@ -2007,7 +2032,10 @@ struct AllocationTraverser : RecursiveASTVisitor<AllocationTraverser> {
       }
       auto it = visitedFuncs.find(FD);
       if (it == visitedFuncs.end()) {
-        visitedFuncs[FD] = std::nullopt;
+        auto storedResult = IsAllocator(wrap<ConstFuncRef>(FD));
+        visitedFuncs[FD] = storedResult;
+        if (storedResult != AllocType::Unknown)
+          return storedResult;
         return AnalyzeAllocType(FD, visitedFuncs);
       }
       return it->second;

lib/CppInterOp/CppInterOp.td

@@ -1304,7 +1304,7 @@ def GetFunctionReturnType : CppInterOpAPI {
 
 def IsAllocator : CppInterOpAPI {
   let Doc = "Checks if the provided function is an allocator function";
-  let ReturnType = "bool";
+  let ReturnType = "AllocType";
   let Args = [
     Arg<"ConstFuncRef", "func">
   ];

unittests/CppInterOp/APINotes/TestAttributeMerge.h

@@ -0,0 +1,14 @@
+#ifndef UNITTESTS_CPPINTEROP_TESTATTRIBUTEMERGE_H
+#define UNITTESTS_CPPINTEROP_TESTATTRIBUTEMERGE_H
+
+[[clang::annotate("cppAllocMalloc")]] void* mergeFunc();
+
+#pragma clang attribute push([[clang::annotate("cppAllocNew")]],               \
+                             apply_to = function)
+int* overloadFunc();
+int* overloadFunc(int n);
+#pragma clang attribute pop
+
+[[clang::annotate("cppAllocMalloc")]] void* func71_helper();
+
+#endif

unittests/CppInterOp/APINotes/TestHeader.apinotes

@@ -4,3 +4,17 @@ Functions:
     RetainCountConvention: CFReturnsRetained
   - Name: testNotAlloc
     RetainCountConvention: CFReturnsNotRetained
+  - Name: testMalloc
+    SwiftReturnOwnership: cppAllocMalloc
+  - Name: testNew
+    SwiftReturnOwnership: cppAllocNew
+  - Name: testNewArr
+    SwiftReturnOwnership: cppAllocNewArr
+  - Name: testOperatorNew
+    SwiftReturnOwnership: cppAllocOperatorNew
+  - Name: testOperatorNewArr
+    SwiftReturnOwnership: cppAllocOperatorNewArr
+  - Name: testNone
+    SwiftReturnOwnership: cppAllocNone
+  - Name: testWeirdAttr
+    SwiftReturnOwnership: someWeirdAttr

unittests/CppInterOp/APINotes/TestHeader.h

@@ -4,4 +4,11 @@
 void* testAlloc(int value);
 void testNotAlloc(void* ptr);
 
+void* testMalloc();
+void* testNew();
+void* testNewArr();
+void* testOperatorNew();
+void* testOperatorNewArr();
+void* testNone();
+void* testWeirdAttr();
 #endif

unittests/CppInterOp/FunctionReflectionTest.cpp

@@ -844,7 +844,10 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsAllocator) {
   std::vector<Decl*> Decls;
   std::string code = R"(
     class Klass{
-      int val;
+      int val = 0;
+      int* __attribute__((annotate("cppAllocNone"))) getValAdress(){
+        return &val;
+      }
     };
     __attribute__((ownership_returns(malloc)))
     Klass* Allocator(){
@@ -856,24 +859,74 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsAllocator) {
       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
+
+    void __attribute__((ownership_takes(malloc, 1))) Deallocator(void* p);
+    void* __attribute__((cf_returns_retained)) CFAllocFunc();
+
+    int* __attribute__((annotate("cppAllocNone"))) NoneFunc();
+    int* __attribute__((annotate("cppAllocNew"))) NewFunc();
+    int* __attribute__((annotate("cppAllocNewArr"))) NewArrFunc();
+    int* __attribute__((annotate("cppAllocMalloc"))) MallocFunc();
+    int* __attribute__((annotate("cppAllocOperatorNew"))) OpNewFunc();
+    int* __attribute__((annotate("cppAllocOperatorNewArr"))) OpNewArrFunc();
+    int* __attribute__((annotate("unrelatedAttr"))) UnrelatedFunc();
+    __declspec(restrict) int* DeclspecRestrictFunc();
+    template <typename T>
+    __attribute__((annotate("cppAllocNew"))) T* TemplatedFunc(){
+      return new T;
+    }
+    template <> int* TemplatedFunc();
+    template char* TemplatedFunc<char>();
   )";
-  TestFixture::CreateInterpreter({"-include", "stdlib.h"});
-  Interp->process(code);
-  auto mallocDecl = Cpp::GetNamed("malloc");
-  EXPECT_TRUE(Cpp::IsAllocator(Cpp::ConstFuncRef{mallocDecl.data}));
-  Cpp::DeleteInterpreter();
+  GetAllTopLevelDecls(code, Decls, true,
+                      {"-std=c++17", "-include", "stdlib.h", "-fdeclspec"});
+#define TESTIA(N, EXP)                                                         \
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef { Cpp::GetNamed(#N).data }), EXP)
+
+  TESTIA(malloc, Cpp::AllocType::Malloc);
+  TESTIA(Allocator, Cpp::AllocType::Malloc);
+  TESTIA(Allocator2, Cpp::AllocType::Malloc);
+  TESTIA(foo, Cpp::AllocType::Unknown);
+  TESTIA(Deallocator, Cpp::AllocType::Unknown);
+  TESTIA(CFAllocFunc, Cpp::AllocType::Malloc);
+  TESTIA(NoneFunc, Cpp::AllocType::None);
+  TESTIA(NewFunc, Cpp::AllocType::New);
+  TESTIA(NewArrFunc, Cpp::AllocType::NewArr);
+  TESTIA(MallocFunc, Cpp::AllocType::Malloc);
+  TESTIA(OpNewFunc, Cpp::AllocType::OperatorNew);
+  TESTIA(OpNewArrFunc, Cpp::AllocType::OperatorNewArr);
+  TESTIA(UnrelatedFunc, Cpp::AllocType::Unknown);
+  TESTIA(DeclspecRestrictFunc, Cpp::AllocType::Unknown);
+
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{Decls[14]}),
+            Cpp::AllocType::New);
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{Decls[15]}),
+            Cpp::AllocType::New);
+  ASTContext& C = Interp->getCI()->getASTContext();
+  std::vector<Cpp::TemplateArgInfo> charArg = {C.CharTy.getAsOpaquePtr()};
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{
+                Cpp::InstantiateTemplate(Decls[14], charArg).data}),
+            Cpp::AllocType::New);
+
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{
+                Cpp::GetNamed("getValAdress", Cpp::GetNamed("Klass")).data}),
+            Cpp::AllocType::None);
+
+  //! Fn coverage
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{nullptr}),
+            Cpp::AllocType::Unknown);
+  // casting coverage
+  EXPECT_EQ(Cpp::IsAllocator(Cpp::ConstFuncRef{Cpp::GetNamed("Klass").data}),
+            Cpp::AllocType::Unknown);
 
+  Cpp::DeleteInterpreter();
+#ifdef EMSCRIPTEN
+  GTEST_SKIP() << "Test fails for Emscipten builds";
+#endif
+  std::string include_flag;
   // APINotes check
-#if !defined(CPPINTEROP_USE_CLING) && !defined(__EMSCRIPTEN__)
-  std::string include_flag =
+#ifndef CPPINTEROP_USE_CLING
+  include_flag =
       "-I" + std::string(CPPINTEROP_DIR) + "unittests/CppInterOp/APINotes";
   std::vector<const char*> interpreter_args = {
       "-fmodules", "-fimplicit-module-maps", "-fapinotes-modules",
@@ -883,13 +936,48 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsAllocator) {
   #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}));
+  TESTIA(testAlloc, Cpp::AllocType::Malloc);
+  TESTIA(testNotAlloc, Cpp::AllocType::Unknown);
+  TESTIA(testMalloc, Cpp::AllocType::Malloc);
+  TESTIA(testNew, Cpp::AllocType::New);
+  TESTIA(testNewArr, Cpp::AllocType::NewArr);
+  TESTIA(testOperatorNew, Cpp::AllocType::OperatorNew);
+  TESTIA(testOperatorNewArr, Cpp::AllocType::OperatorNewArr);
+  TESTIA(testNone, Cpp::AllocType::None);
+  TESTIA(testWeirdAttr, Cpp::AllocType::Unknown);
+
   Cpp::DeleteInterpreter();
 #endif
+#undef TESTIA
+  include_flag =
+      "-I" + std::string(CPPINTEROP_DIR) + "unittests/CppInterOp/APINotes";
+  Decls.clear();
+  code = R"(
+    void* mergeFunc() {
+      return malloc(sizeof(int));
+    }
+  )";
+  GetAllTopLevelDecls(code, Decls, true,
+                      {"-std=c++17", include_flag.c_str(), "-include",
+                       "stdlib.h", "-include", "TestAttributeMerge.h"});
+  EXPECT_EQ(Cpp::IsAllocator(Decls[0]), Cpp::AllocType::Malloc);
+
+  Decls.clear();
+  code = R"(
+    int* overloadFunc(){
+      return new int;
+    }
+
+    int* overloadFunc(int n){
+      return new int(n);
+    }
+  )";
+  GetAllTopLevelDecls(
+      code, Decls, true,
+      {"-std=c++17", include_flag.c_str(), "-include", "TestAttributeMerge.h"});
+  EXPECT_EQ(Cpp::IsAllocator(Decls[0]), Cpp::AllocType::New);
+  EXPECT_EQ(Cpp::IsAllocator(Decls[1]), Cpp::AllocType::New);
 }
 
 TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_IsDeallocator) {
@@ -1749,14 +1837,27 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_GetAllocType) {
       a = new int;
       return a;
     }
+
+    //Not analyzed, attribute is merged in header
+    void* func71_helper();
+
+    void* func71(){
+      return func71_helper();
+    }
     )";
+  std::string include_flag =
+      "-I" + std::string(CPPINTEROP_DIR) + "unittests/CppInterOp/APINotes";
+#ifndef EMSCRIPTEN
+  TestFixture::CreateInterpreter(
+      {"-std=c++17", include_flag.c_str(), "-include", "TestAttributeMerge.h"});
+#else
   TestFixture::CreateInterpreter({"-std=c++17"});
-  Interp->declare(code);
+#endif
 
+  Interp->declare(code);
 #define TESTAC(N, EXP)                                                         \
-  EXPECT_EQ(                                                                   \
-      Cpp::GetAllocType(Cpp::ConstFuncRef { Cpp::GetNamed("func" #N).data }),  \
-      Cpp::AllocType::EXP)
+  EXPECT_EQ(Cpp::GetAllocType(Cpp::FuncRef { Cpp::GetNamed("func" #N).data }), \
+            Cpp::AllocType::EXP)
 
   TESTAC(0, New);
   TESTAC(1, New);
@@ -1819,10 +1920,14 @@ TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_GetAllocType) {
   TESTAC(58, OperatorNew);
   TESTAC(59, Malloc);
   TESTAC(70, Unknown);
+#ifndef EMSCRIPTEN
+  TESTAC(71, Malloc);
+#endif
 #undef TESTAC
 
   Cpp::DeleteInterpreter();
 }
+
 TYPED_TEST(CPPINTEROP_TEST_MODE, FunctionReflection_GetFunctionSignature) {
   std::vector<Decl*> Decls;
   std::string code = R"(

unittests/CppInterOp/TracingTests.cpp

@@ -1142,7 +1142,7 @@ TEST_F(TracingTest, WriteToFileSuppressesSelfTrace) {
 #ifndef EMSCRIPTEN
 TEST_F(TracingTest, ReproducerCompilesViaInterpreter) {
   // Create an interpreter so we can exercise real API calls.
-  Cpp::CreateInterpreter({});
+  Cpp::CreateInterpreter({"-std=c++17"});
 
   // Verify tracing is active after interpreter creation.
   ASSERT_NE(TheTraceInfo, nullptr)
@@ -1192,7 +1192,7 @@ TEST_F(TracingTest, ReproducerCompilesViaInterpreter) {
   EXPECT_THAT(content, HasSubstr("Cpp::"));
 
   // Create a fresh interpreter and #include the reproducer file as-is.
-  Cpp::CreateInterpreter({});
+  Cpp::CreateInterpreter({"-std=c++17"});
   Cpp::AddIncludePath(CPPINTEROP_DIR "/include");
   // Generated .inc files live under the build tree.
   Cpp::AddIncludePath(CPPINTEROP_BINARY_DIR "/include");