[AbstractCallSite] Look though constant cast expression when checking for callee use

[sdmitriev]

Summary

That makes AbstractCallSite::isCallee(const Use *) behavior consistent with AbstractCallSite constructor.

[jdoerfert]

LGTM, THX!

llvm/include/llvm/IR/AbstractCallSite.h

148

I checked and it still matches. At some point I locally removed the single use restriction actually. Nothing for this patch but just for your information ;)

llvm/include/llvm/IR/AbstractCallSite.h

public:
   /// Return true if @p U is the use that defines the callee of this ACS.
   bool isCallee(const Use *U) const {
     if (isDirectCall())
       return CB->isCallee(U);
 
     assert(!CI.ParameterEncoding.empty() &&
            "Callback without parameter encoding!");
 
+    // If the use is actually in a constant cast expression which itself
+    // has only one use, we look through the constant cast expression.
+    if (auto *CE = dyn_cast<ConstantExpr>(U->getUser()))
+      if (CE->getNumUses() == 1 && CE->isCast())
+        U = &*CE->use_begin();
+
     return (int)CB->getArgOperandNo(U) == CI.ParameterEncoding[0];
   }
 
   /// Return the number of parameters of the callee.
   unsigned getNumArgOperands() const {
     if (isDirectCall())
       return CB->getNumArgOperands();
     // Subtract 1 for the callee encoding.

llvm/unittests/IR/AbstractCallSiteTest.cpp

+//===----- AbstractCallSiteTest.cpp - AbstractCallSite Unittests ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/AbstractCallSite.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
+  SMDiagnostic Err;
+  std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+  if (!Mod)
+    Err.print("AbstractCallSiteTests", errs());
+  return Mod;
+}
+
+TEST(AbstractCallSite, CallbackCall) {
+  LLVMContext C;
+
+  const char *IR =
+      "define void @callback(i8* %X, i32* %A) {\n"
+      "  ret void\n"
+      "}\n"
+      "define void @foo(i32* %A) {\n"
+      "  call void (i32, void (i8*, ...)*, ...) @broker(i32 1, void (i8*, ...)* bitcast (void (i8*, i32*)* @callback to void (i8*, ...)*), i32* %A)\n"
+      "  ret void\n"
+      "}\n"
+      "declare !callback !0 void @broker(i32, void (i8*, ...)*, ...)\n"
+      "!0 = !{!1}\n"
+      "!1 = !{i64 1, i64 -1, i1 true}";
+
+  std::unique_ptr<Module> M = parseIR(C, IR);
+  ASSERT_TRUE(M);
+
+  Function *Callback = M->getFunction("callback");
+  ASSERT_NE(Callback, nullptr);
+
+  const Use *CallbackUse = Callback->getSingleUndroppableUse();
+  ASSERT_NE(CallbackUse, nullptr);
+
+  AbstractCallSite ACS(CallbackUse);
+  EXPECT_TRUE(ACS);
+  EXPECT_TRUE(ACS.isCallbackCall());
+  EXPECT_TRUE(ACS.isCallee(CallbackUse));
+  EXPECT_EQ(ACS.getCalledFunction(), Callback);
+}

llvm/unittests/IR/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
   AsmParser
   Core
   Support
   Passes
   TransformUtils
   )
 
 add_llvm_unittest(IRTests
+  AbstractCallSiteTest.cpp
   AsmWriterTest.cpp
   AttributesTest.cpp
   BasicBlockTest.cpp
   CFGBuilder.cpp
   ConstantRangeTest.cpp
   ConstantsTest.cpp
   DataLayoutTest.cpp
   DebugInfoTest.cpp