[WebAssembly] Fix broken assumption that all bitcasts are to functions types

[sbc100]

Summary

Handle bitcasts to non-function (integer) pointer types and also handle cast
instructions as well as cast const exprs.

Fixes PR39591

[sbc100]

• improve tests

[sbc100]

• improve tests

[dschuff]

lib/Target/WebAssembly/WebAssemblyAddMissingPrototypes.cpp

123 ↗

(On Diff #173760)

nit: this (and below) could be auto because of the dyn_cast

129 ↗

(On Diff #173760)

isn't this modifying F.uses while iterating over it?

[sbc100]

• feedbac + clang-format

lib/Target/WebAssembly/WebAssemblyAddMissingPrototypes.cpp

129 ↗

(On Diff #173760)

I'm not sure. Is the eraseFromParent call that makes you think that? Any idea how would I validate that?

[sbc100]

• erase after uses iteration

[dschuff]

lib/Target/WebAssembly/WebAssemblyAddMissingPrototypes.cpp

129 ↗

(On Diff #173760)

Yes, Inst is a user of F, and I'm pretty sure eraseFromParent removes the instruction from the BB list and also drops its use references.
edit: yes, it's the operator delete of User that does that.

llvm/trunk/lib/Target/WebAssembly/WebAssemblyAddMissingPrototypes.cpp

if (F.getFunctionType()->getNumParams() != 0)
           F.getName());
 
     // Create a function prototype based on the first call site (first bitcast)
     // that we find.
     FunctionType *NewType = nullptr;
     Function *NewF = nullptr;
     for (Use &U : F.uses()) {
       LLVM_DEBUG(dbgs() << "prototype-less use: " << F.getName() << "\n");
-      if (BitCastOperator *BC = dyn_cast<BitCastOperator>(U.getUser())) {
-        FunctionType *DestType =
-            cast<FunctionType>(BC->getDestTy()->getPointerElementType());
-
+      if (auto *BC = dyn_cast<BitCastOperator>(U.getUser())) {
+        if (auto *DestType = dyn_cast<FunctionType>(
+                BC->getDestTy()->getPointerElementType())) {
+          if (!NewType) {
             // Create a new function with the correct type
             NewType = DestType;
             NewF = Function::Create(NewType, F.getLinkage(), F.getName());
             NewF->setAttributes(F.getAttributes());
             NewF->removeFnAttr("no-prototype");
-        break;
+          } else {
+            if (NewType != DestType) {
+              report_fatal_error("Prototypeless function used with "
+                                 "conflicting signatures: " +
+                                 F.getName());
+            }
+          }
+        }
       }
     }
 
     if (!NewType) {
       LLVM_DEBUG(
           dbgs() << "could not derive a function prototype from usage: " +
                         F.getName() + "\n");
       continue;
     }
 
-    for (Use &U : F.uses()) {
-      if (BitCastOperator *BC = dyn_cast<BitCastOperator>(U.getUser())) {
-        FunctionType *DestType =
-            cast<FunctionType>(BC->getDestTy()->getPointerElementType());
-        if (NewType != DestType) {
-          report_fatal_error(
-              "Prototypeless function used with conflicting signatures: " +
-              F.getName());
-        }
-        BC->replaceAllUsesWith(NewF);
-        Replacements.emplace_back(&F, NewF);
+    SmallVector<Instruction *, 4> DeadInsts;
+
+    for (Use &US : F.uses()) {
+      User *U = US.getUser();
+      if (auto *BC = dyn_cast<BitCastOperator>(U)) {
+        if (auto *Inst = dyn_cast<BitCastInst>(U)) {
+          // Replace with a new bitcast
+          IRBuilder<> Builder(Inst);
+          Value *NewCast = Builder.CreatePointerCast(NewF, BC->getDestTy());
+          Inst->replaceAllUsesWith(NewCast);
+          DeadInsts.push_back(Inst);
+        } else if (auto *Const = dyn_cast<ConstantExpr>(U)) {
+          Constant *NewConst =
+              ConstantExpr::getPointerCast(NewF, BC->getDestTy());
+          Const->replaceAllUsesWith(NewConst);
         } else {
-        dbgs() << *U.getUser()->getType() << "\n";
+          dbgs() << *U->getType() << "\n";
 #ifndef NDEBUG
-        U.getUser()->dump();
+          U->dump();
 #endif
-        report_fatal_error(
-            "unexpected use of prototypeless function: " + F.getName() + "\n");
+          report_fatal_error("unexpected use of prototypeless function: " +
+                             F.getName() + "\n");
         }
       }
     }
 
+    for (auto I : DeadInsts)
+      I->eraseFromParent();
+    Replacements.emplace_back(&F, NewF);
+  }
+
+
   // Finally replace the old function declarations with the new ones
   for (auto &Pair : Replacements) {
     Function *Old = Pair.first;
     Function *New = Pair.second;
     Old->eraseFromParent();
     M.getFunctionList().push_back(New);
   }
 
   return !Replacements.empty();
 }

llvm/trunk/test/CodeGen/WebAssembly/add-prototypes.ll

 ; RUN: opt -S -wasm-add-missing-prototypes %s | FileCheck %s
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
+; CHECK: @foo_addr = global i64 (i32)* @foo, align 8
 @foo_addr = global i64 (i32)* bitcast (i64 (...)* @foo to i64 (i32)*), align 8
 
-define void @bar(i32 %a) {
-entry:
+; CHECK: @foo_addr_i8 = global i8* bitcast (i64 (i32)* @foo to i8*), align 8
+@foo_addr_i8 = global i8* bitcast (i64 (...)* @foo to i8*), align 8
+
+; CHECK-LABEL: @call_foo
+; CHECK: %call = call i64 @foo(i32 42)
+define void @call_foo(i32 %a) {
   %call = call i64 bitcast (i64 (...)* @foo to i64 (i32)*)(i32 42)
   ret void
 }
 
-declare i64 @foo(...) #1
+; CHECK-LABEL: @call_foo_ptr
+; CHECK: %call = call i64 @foo(i32 43)
+define i64 @call_foo_ptr(i32 %a) {
+  %1 = bitcast i64 (...)* @foo to i64 (i32)*
+  %call = call i64 (i32) %1(i32 43)
+  ret i64 %call
+}
 
-attributes #1 = { "no-prototype" }
+; CHECK-LABEL: @to_intptr_inst
+; CHECK: ret i8* bitcast (i64 (i32)* @foo to i8*)
+define i8* @to_intptr_inst() {
+  %1 = bitcast i64 (...)* @foo to i8*
+  ret i8* %1
+}
+
+; CHECK-LABEL: @to_intptr_constexpr
+; CHECK: ret i8* bitcast (i64 (i32)* @foo to i8*)
+define i8* @to_intptr_constexpr() {
+  ret i8* bitcast (i64 (...)* @foo to i8*)
+}
 
-; CHECK: %call = call i64 @foo(i32 42)
 ; CHECK: declare i64 @foo(i32)
+declare i64 @foo(...) #1
+
 ; CHECK-NOT: attributes {{.*}} = { {{.*}}"no-prototype"{{.*}} }
+attributes #1 = { "no-prototype" }