update handling of rust (de)alloc fn to use demangled names

Author: ZuseZ4

enzyme/Enzyme/GradientUtils.cpp

@@ -58,6 +58,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/TimeProfiler.h"
 
+#include "llvm/Demangle/Demangle.h"
+
 #if LLVM_VERSION_MAJOR >= 14
 #define addAttribute addAttributeAtIndex
 #define hasAttribute hasAttributeAtIndex
@@ -9472,7 +9474,8 @@ llvm::CallInst *freeKnownAllocation(llvm::IRBuilder<> &builder,
                                     GradientUtils *gutils) {
   assert(isAllocationFunction(allocationfn, TLI));
 
-  if (allocationfn == "__rust_alloc" || allocationfn == "__rust_alloc_zeroed") {
+  std::string demangledName = llvm::demangle(allocationfn);
+  if (demangledName == "__rustc::__rust_alloc" || demangledName == "__rustc::__rust_alloc_zeroed" ) {
     Type *VoidTy = Type::getVoidTy(tofree->getContext());
     Type *IntPtrTy = orig->getType();
     Type *RustSz = orig->getArgOperand(0)->getType();

enzyme/Enzyme/LibraryFuncs.h

@@ -26,6 +26,7 @@
 #include <llvm/ADT/StringMap.h>
 #include <llvm/Analysis/AliasAnalysis.h>
 #include <llvm/Analysis/TargetLibraryInfo.h>
+#include "llvm/Demangle/Demangle.h"
 #include <llvm/IR/IRBuilder.h>
 #include <llvm/IR/InlineAsm.h>
 #include <llvm/IR/Instructions.h>
@@ -55,7 +56,8 @@ static inline bool isAllocationFunction(const llvm::StringRef name,
     return true;
   if (name == "__size_returning_new_experiment")
     return true;
-  if (name == "__rust_alloc" || name == "__rust_alloc_zeroed")
+  std::string demangledName = llvm::demangle(name.str());
+  if (demangledName == "__rustc::__rust_alloc" || demangledName == "__rustc::__rust_alloc_zeroed")
     return true;
   if (name == "julia.gc_alloc_obj" || name == "jl_gc_alloc_typed" ||
       name == "ijl_gc_alloc_typed")
@@ -131,7 +133,8 @@ static inline bool isDeallocationFunction(const llvm::StringRef name,
       return true;
     if (name == "_mlir_memref_to_llvm_free")
       return true;
-    if (name == "__rust_dealloc")
+    std::string demangledName = llvm::demangle(name.str());
+    if (demangledName == "__rustc::__rust_dealloc")
       return true;
     if (name == "swift_release")
       return true;
@@ -207,7 +210,8 @@ static inline void zeroKnownAllocation(llvm::IRBuilder<> &bb,
   assert(isAllocationFunction(funcName, TLI));
 
   // Don't re-zero an already-zero buffer
-  if (funcName == "calloc" || funcName == "__rust_alloc_zeroed")
+    std::string demangledName = llvm::demangle(funcName.str());
+  if (funcName == "calloc" || demangledName == "__rustc::__rust_alloc_zeroed")
     return;
 
   Value *allocSize = argValues[0];