[WebAssembly] Only RAUW a constant once in FixFunctionBitcasts

[dschuff]

Summary

When we collect 2 uses of a function in FindUses and then RAUW when we visit
the first, we end up visiting the wrapper (because the second was RAUW'd).
We still want to use RAUW instead of just Use->set() because it has special handling for
Constants, so this patch just ensures that only one use of each constant is added to the work list.

[dschuff]

The alternative would of course be to just collect the uses in a set, deduplicated by callee, and then RAUW everything. But I'm actually unclear why the original has distinction between whether the user is a constant or not; it seems to work either way for me locally?

[dschuff]

OK, actually I forgot that our use list doesn't contain all the uses, just the ones with mismatched types. So RAUWing when U is a constant was intended to really just replaces the uses within the one constexpr? but because of uniquing it affects the other calls as well and we visit that use again.

[sunfish]

Constants are meant to be immutable at some level in LLVM, so my idea with the RAUW was to avoid mutating them (RAUW has special logic for constants so that it doesn't actually mutate them). However, you're right that this breaks when there are multiple uses of the same function.

[sunfish]

Oops, didn't mean to accept.

[dschuff]

Use RAUW for constants, but only add a single use to the worklist

[dschuff]

• clang-format

[dschuff]

OK makes sense. Patch updated to still use RAUW but only once per constant.

[sunfish]

That should do it.

lib/Target/WebAssembly/WebAssemblyFixFunctionBitcasts.cpp

 char FixFunctionBitcasts::ID = 0;
 ModulePass *llvm::createWebAssemblyFixFunctionBitcasts() {
   return new FixFunctionBitcasts();
 }
 
 // Recursively descend the def-use lists from V to find non-bitcast users of
 // bitcasts of V.
 static void FindUses(Value *V, Function &F,
-                     SmallVectorImpl<std::pair<Use *, Function *>> &Uses) {
+                     SmallVectorImpl<std::pair<Use *, Function *>> &Uses,
+                     SmallPtrSetImpl<Constant*>& ConstantBCs) {
   for (Use &U : V->uses()) {
     if (BitCastOperator *BC = dyn_cast<BitCastOperator>(U.getUser()))
-      FindUses(BC, F, Uses);
-    else if (U.get()->getType() != F.getType())
+      FindUses(BC, F, Uses, ConstantBCs);
+    else if (U.get()->getType() != F.getType()) {
+      if (isa<Constant>(U.get())) {
+        // Only add constant bitcasts to the list once; they get RAUW'd
+        auto c = ConstantBCs.insert(cast<Constant>(U.get()));
+        if (!c.second)
+          continue;
+      }
       Uses.push_back(std::make_pair(&U, &F));
     }
   }
+}
 
 // Create a wrapper function with type Ty that calls F (which may have a
 // different type). Attempt to support common bitcasted function idioms:
 //  - Call with more arguments than needed: arguments are dropped
 //  - Call with fewer arguments than needed: arguments are filled in with undef
 //  - Return value is not needed: drop it
 //  - Return value needed but not present: supply an undef
 //  
else {
     return nullptr;
   }
 
   return Wrapper;
 }
 
 bool FixFunctionBitcasts::runOnModule(Module &M) {
   SmallVector<std::pair<Use *, Function *>, 0> Uses;
+  SmallPtrSet<Constant*, 2> ConstantBCs;
 
   // Collect all the places that need wrappers.
   for (Function &F : M)
-    FindUses(&F, F, Uses);
+    FindUses(&F, F, Uses, ConstantBCs);
 
   DenseMap<std::pair<Function *, FunctionType *>, Function *> Wrappers;
 
   for (auto &UseFunc : Uses) {
     Use *U = UseFunc.first;
     Function *F = UseFunc.second;
     PointerType *PTy = cast<PointerType>(U->get()->getType());
     FunctionType *Ty = dyn_cast<FunctionType>(PTy->getElementType());
for (auto &UseFunc : Uses) {
     Function *Wrapper = Pair.first->second;
     if (!Wrapper)
       continue;
 
     if (isa<Constant>(U->get()))
       U->get()->replaceAllUsesWith(Wrapper);
     else
       U->set(Wrapper);
+
   }
 
   return true;
 }

test/CodeGen/WebAssembly/function-bitcasts.ll

 ; RUN: llc < %s -asm-verbose=false | FileCheck %s
 
 ; Test that function pointer casts are replaced with wrappers.
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
 ; CHECK-LABEL: test:
 ; CHECK-NEXT: call        .Lbitcast@FUNCTION{{$}}
+; CHECK-NEXT: call        .Lbitcast@FUNCTION{{$}}
 ; CHECK-NEXT: call        .Lbitcast.1@FUNCTION{{$}}
 ; CHECK-NEXT: i32.const   $push[[L0:[0-9]+]]=, 0
 ; CHECK-NEXT: call        .Lbitcast.2@FUNCTION, $pop[[L0]]{{$}}
+; CHECK-NEXT: i32.const   $push[[L1:[0-9]+]]=, 0
+; CHECK-NEXT: call        .Lbitcast.2@FUNCTION, $pop[[L1]]{{$}}
+; CHECK-NEXT: i32.const   $push[[L2:[0-9]+]]=, 0
+; CHECK-NEXT: call        .Lbitcast.2@FUNCTION, $pop[[L2]]{{$}}
+; CHECK-NEXT: call        foo0@FUNCTION
 ; CHECK-NEXT: i32.call    $drop=, .Lbitcast.3@FUNCTION{{$}}
 ; CHECK-NEXT: call        foo2@FUNCTION{{$}}
+; CHECK-NEXT: call        foo1@FUNCTION{{$}}
 ; CHECK-NEXT: call        foo3@FUNCTION{{$}}
 ; CHECK-NEXT: .endfunc
 
 ; CHECK-LABEL: .Lbitcast:
 ; CHECK-NEXT: .local      i32
 ; CHECK-NEXT: call        has_i32_arg@FUNCTION, $0{{$}}
 ; CHECK-NEXT: .endfunc
 

 declare void @foo0()
 declare void @foo1()
 declare void @foo2()
 declare void @foo3()
 
 define void @test() {
 entry:
   call void bitcast (void (i32)* @has_i32_arg to void ()*)()
+  call void bitcast (void (i32)* @has_i32_arg to void ()*)()
   call void bitcast (i32 ()* @has_i32_ret to void ()*)()
   call void bitcast (void ()* @foo0 to void (i32)*)(i32 0)
+  %p = bitcast void ()* @foo0 to void (i32)*
+  call void %p(i32 0)
+  %q = bitcast void ()* @foo0 to void (i32)*
+  call void %q(i32 0)
+  %r = bitcast void (i32)* %q to void ()*
+  call void %r()
   %t = call i32 bitcast (void ()* @foo1 to i32 ()*)()
   call void bitcast (void ()* @foo2 to void ()*)()
+  call void @foo1()
   call void @foo3()
+
   ret void
 }