Enhance TLI detection of __size_returning_new lib funcs.

Snehasish Kumar · Snehasish Kumar · commit c87dba49dce5 · 2024-08-09T02:22:10.000Z
Previously the return types of __size_returning_new variants were not
validated based on their members. This patch checks the members
manually, also generalizes the size_t checks to be based on the module
instead of being hardcoded.
diff --git a/llvm/include/llvm/Analysis/TargetLibraryInfo.def b/llvm/include/llvm/Analysis/TargetLibraryInfo.def
@@ -365,22 +365,22 @@ TLI_DEFINE_SIG_INTERNAL(Ptr, Long, Long, Ptr, Bool)
 /// __sized_ptr_t __size_returning_new(size_t size)
 TLI_DEFINE_ENUM_INTERNAL(size_returning_new)
 TLI_DEFINE_STRING_INTERNAL("__size_returning_new")
-TLI_DEFINE_SIG_INTERNAL(Struct, Long)
+TLI_DEFINE_SIG_INTERNAL(/* Checked manually. */)
 
 /// __sized_ptr_t __size_returning_new_hot_cold(size_t, __hot_cold_t)
 TLI_DEFINE_ENUM_INTERNAL(size_returning_new_hot_cold)
 TLI_DEFINE_STRING_INTERNAL("__size_returning_new_hot_cold")
-TLI_DEFINE_SIG_INTERNAL(Struct, Long, Bool)
+TLI_DEFINE_SIG_INTERNAL(/* Checked manually. */)
 
 /// __sized_ptr_t __size_returning_new_aligned(size_t, std::align_val_t)
 TLI_DEFINE_ENUM_INTERNAL(size_returning_new_aligned)
 TLI_DEFINE_STRING_INTERNAL("__size_returning_new_aligned")
-TLI_DEFINE_SIG_INTERNAL(Struct, Long, Long)
+TLI_DEFINE_SIG_INTERNAL(/* Checked manually. */)
 
 /// __sized_ptr_t __size_returning_new_aligned(size_t, std::align_val_t, __hot_cold_t)
 TLI_DEFINE_ENUM_INTERNAL(size_returning_new_aligned_hot_cold)
 TLI_DEFINE_STRING_INTERNAL("__size_returning_new_aligned_hot_cold")
-TLI_DEFINE_SIG_INTERNAL(Struct, Long, Long, Bool)
+TLI_DEFINE_SIG_INTERNAL(/* Checked manually. */)
 
 /// double __acos_finite(double x);
 TLI_DEFINE_ENUM_INTERNAL(acos_finite)
diff --git a/llvm/lib/Analysis/TargetLibraryInfo.cpp b/llvm/lib/Analysis/TargetLibraryInfo.cpp
@@ -1049,6 +1049,49 @@ static bool matchType(FuncArgTypeID ArgTy, const Type *Ty, unsigned IntBits,
   llvm_unreachable("Invalid type");
 }
 
+static bool isValidProtoForSizeReturningNew(const FunctionType& FTy, LibFunc F, const Module& M, int SizeTSizeBits) {
+  switch (F) {
+    case LibFunc_size_returning_new: {
+      if(FTy.getNumParams() != 1 || !FTy.getParamType(0)->isIntegerTy(SizeTSizeBits)) {
+        return false;
+      }
+    }
+    break;
+    case LibFunc_size_returning_new_hot_cold: {
+      if(FTy.getNumParams() != 2 ||
+         !FTy.getParamType(0)->isIntegerTy(SizeTSizeBits) ||
+         !FTy.getParamType(1)->isIntegerTy(8)) {
+        return false;
+      }
+    }
+    break;
+    case LibFunc_size_returning_new_aligned: {
+      if(FTy.getNumParams() !=2 ||
+        !FTy.getParamType(0)->isIntegerTy(SizeTSizeBits) ||
+         !FTy.getParamType(1)->isIntegerTy(SizeTSizeBits)) {
+        return false;
+      }
+    }
+    break;
+    case LibFunc_size_returning_new_aligned_hot_cold:
+      if( FTy.getNumParams() != 3 ||
+          !FTy.getParamType(0)->isIntegerTy(SizeTSizeBits) ||
+         !FTy.getParamType(1)->isIntegerTy(SizeTSizeBits) ||
+         !FTy.getParamType(2)->isIntegerTy(8)) {
+        return false;
+      }
+    break;
+    default:
+      return false;
+  }
+
+  auto& Context = M.getContext();
+  llvm::Type *I8Ty = Type::getInt8Ty(Context);
+  llvm::PointerType *I8PtrTy = PointerType::get(I8Ty, 0);
+  llvm::StructType *SizedPtrTy = llvm::StructType::get(Context, {I8PtrTy, Type::getIntNTy(Context, SizeTSizeBits)});
+  return FTy.getReturnType() == SizedPtrTy;
+}
+
 bool TargetLibraryInfoImpl::isValidProtoForLibFunc(const FunctionType &FTy,
                                                    LibFunc F,
                                                    const Module &M) const {
@@ -1099,7 +1142,13 @@ bool TargetLibraryInfoImpl::isValidProtoForLibFunc(const FunctionType &FTy,
 
     return false;
   }
-
+    // Special handling of __size_returning_new functions that return a struct
+    // of type {void*, size_t}.
+  case LibFunc_size_returning_new:
+  case LibFunc_size_returning_new_hot_cold:
+  case LibFunc_size_returning_new_aligned:
+  case LibFunc_size_returning_new_aligned_hot_cold:
+    return isValidProtoForSizeReturningNew(FTy, F, M, getSizeTSize(M));
   default:
     break;
   }
diff --git a/llvm/unittests/Analysis/TargetLibraryInfoTest.cpp b/llvm/unittests/Analysis/TargetLibraryInfoTest.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/SourceMgr.h"
@@ -81,6 +82,30 @@ TEST_F(TargetLibraryInfoTest, InvalidProto) {
   }
 }
 
+TEST_F(TargetLibraryInfoTest, SizeReturningNewProto) {
+  parseAssembly("target datalayout = \"p:64:64:64\"\n"
+                "declare {i8*, i64} @__size_returning_new(i64)\n"
+                "declare {i8*, i64} @__size_returning_new_hot_cold(i64, i8)\n"
+                "declare {i8*, i64} @__size_returning_new_aligned(i64, i64)\n"
+                "declare {i8*, i64} "
+                "@__size_returning_new_aligned_hot_cold(i64, i64, i8)\n");
+
+  llvm::Type *I8Ty = Type::getInt8Ty(Context);
+  llvm::PointerType *I8PtrTy = PointerType::get(I8Ty, 0);
+  llvm::StructType *SizedPtrT =
+      llvm::StructType::get(Context, {I8PtrTy, Type::getInt64Ty(Context)});
+
+  for (const LibFunc LF :
+       {LibFunc_size_returning_new, LibFunc_size_returning_new_aligned,
+        LibFunc_size_returning_new_hot_cold,
+        LibFunc_size_returning_new_aligned_hot_cold}) {
+    TLII.setAvailable(LF);
+    Function *F = M->getFunction(TLI.getName(LF));
+    ASSERT_NE(F, nullptr);
+    EXPECT_EQ(F->getReturnType(), SizedPtrT);
+  }
+}
+
 // Check that we do accept know-correct prototypes.
 TEST_F(TargetLibraryInfoTest, ValidProto) {
   parseAssembly(
@@ -472,10 +497,11 @@ TEST_F(TargetLibraryInfoTest, ValidProto) {
       "declare i8* @_ZnwmSt11align_val_tRKSt9nothrow_t(i64, i64, %struct*)\n"
       "declare i8* @_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t(i64, i64, "
       "%struct*, i8)\n"
-      "declare %struct @__size_returning_new(i64)\n"
-      "declare %struct @__size_returning_new_hot_cold(i64, i8)\n"
-      "declare %struct @__size_returning_new_aligned(i64, i64)\n"
-      "declare %struct @__size_returning_new_aligned_hot_cold(i64, i64, i8)\n"
+      "declare {i8*, i64} @__size_returning_new(i64)\n"
+      "declare {i8*, i64} @__size_returning_new_hot_cold(i64, i8)\n"
+      "declare {i8*, i64} @__size_returning_new_aligned(i64, i64)\n"
+      "declare {i8*, i64} @__size_returning_new_aligned_hot_cold(i64, i64, "
+      "i8)\n"
 
       "declare void @\"??3@YAXPEAX@Z\"(i8*)\n"
       "declare void @\"??3@YAXPEAXAEBUnothrow_t@std@@@Z\"(i8*, %struct*)\n"
