IRMover: Steal arguments when moving functions, NFC

Instead of copying arguments from the source function to the
destination, steal them.  This has a few advantages.

  - The ValueMap doesn't need to be seeded with (or cleared of)
    Arguments.

  - Often the destination function won't have created any arguments yet,
    so this avoids malloc traffic.

  - Argument names don't need to be copied.

Because argument lists are lazy, this required a new
Function::stealArgumentListFrom helper.

llvm-svn: 265519

llvm/include/llvm/IR/Function.h

@@ -88,9 +88,12 @@ private:
   /// built on demand, so that the list isn't allocated until the first client
   /// needs it.  The hasLazyArguments predicate returns true if the arg list
   /// hasn't been set up yet.
+public:
   bool hasLazyArguments() const {
     return getSubclassDataFromValue() & (1<<0);
   }
+
+private:
   void CheckLazyArguments() const {
     if (hasLazyArguments())
       BuildLazyArguments();
@@ -442,6 +445,12 @@ public:
   ///
   void eraseFromParent() override;
 
+  /// Steal arguments from another function.
+  ///
+  /// Drop this function's arguments and splice in the ones from \c Src.
+  /// Requires that this has no function body.
+  void stealArgumentListFrom(Function &Src);
+
   /// Get the underlying elements of the Function... the basic block list is
   /// empty for external functions.
   ///

llvm/lib/IR/Function.cpp

@@ -302,6 +302,28 @@ void Function::BuildLazyArguments() const {
   const_cast<Function*>(this)->setValueSubclassData(SDC &= ~(1<<0));
 }
 
+void Function::stealArgumentListFrom(Function &Src) {
+  assert(isDeclaration() && "Expected no references to current arguments");
+
+  // Drop the current arguments, if any, and set the lazy argument bit.
+  if (!hasLazyArguments()) {
+    assert(llvm::all_of(ArgumentList,
+                        [](const Argument &A) { return A.use_empty(); }) &&
+           "Expected arguments to be unused in declaration");
+    ArgumentList.clear();
+    setValueSubclassData(getSubclassDataFromValue() | (1 << 0));
+  }
+
+  // Nothing to steal if Src has lazy arguments.
+  if (Src.hasLazyArguments())
+    return;
+
+  // Steal arguments from Src, and fix the lazy argument bits.
+  ArgumentList.splice(ArgumentList.end(), Src.ArgumentList);
+  setValueSubclassData(getSubclassDataFromValue() & ~(1 << 0));
+  Src.setValueSubclassData(Src.getSubclassDataFromValue() | (1 << 0));
+}
+
 size_t Function::arg_size() const {
   return getFunctionType()->getNumParams();
 }

llvm/lib/Linker/IRMover.cpp

@@ -980,16 +980,6 @@ bool IRLinker::linkFunctionBody(Function &Dst, Function &Src) {
                                   ValueMapperFlags, &TypeMap,
                                   &GValMaterializer));
 
-  // Go through and convert function arguments over, remembering the mapping.
-  Function::arg_iterator DI = Dst.arg_begin();
-  for (Argument &Arg : Src.args()) {
-    DI->setName(Arg.getName()); // Copy the name over.
-
-    // Add a mapping to our mapping.
-    ValueMap[&Arg] = &*DI;
-    ++DI;
-  }
-
   // Copy over the metadata attachments.
   SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
   Src.getAllMetadata(MDs);
@@ -997,22 +987,19 @@ bool IRLinker::linkFunctionBody(Function &Dst, Function &Src) {
     Dst.setMetadata(I.first, MapMetadata(I.second, ValueMap, ValueMapperFlags,
                                          &TypeMap, &GValMaterializer));
 
-  // Splice the body of the source function into the dest function.
+  // Steal arguments and splice the body of Src into Dst.
+  Dst.stealArgumentListFrom(Src);
   Dst.getBasicBlockList().splice(Dst.end(), Src.getBasicBlockList());
 
-  // At this point, all of the instructions and values of the function are now
-  // copied over.  The only problem is that they are still referencing values in
-  // the Source function as operands.  Loop through all of the operands of the
-  // functions and patch them up to point to the local versions.
+  // At this point, everything has been moved over, but the types and non-local
+  // operands will be wrong.  Loop through everything and patch it up.
+  for (Argument &A : Dst.args())
+    A.mutateType(TypeMap.get(A.getType()));
   for (BasicBlock &BB : Dst)
     for (Instruction &I : BB)
       RemapInstruction(&I, ValueMap, RF_IgnoreMissingEntries | ValueMapperFlags,
                        &TypeMap, &GValMaterializer);
 
-  // There is no need to map the arguments anymore.
-  for (Argument &Arg : Src.args())
-    ValueMap.erase(&Arg);
-
   return false;
 }
 

llvm/unittests/IR/CMakeLists.txt

@@ -12,6 +12,7 @@ set(IRSources
   ConstantsTest.cpp
   DebugInfoTest.cpp
   DominatorTreeTest.cpp
+  FunctionTest.cpp
   IRBuilderTest.cpp
   InstructionsTest.cpp
   IntrinsicsTest.cpp

llvm/unittests/IR/FunctionTest.cpp

@@ -0,0 +1,106 @@
+//===- FunctionTest.cpp - Function unit tests -----------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/IR/Function.h"
+#include "gtest/gtest.h"
+using namespace llvm;
+
+namespace {
+
+TEST(FunctionTest, hasLazyArguments) {
+  LLVMContext C;
+
+  Type *ArgTypes[] = {Type::getInt8Ty(C), Type::getInt32Ty(C)};
+  FunctionType *FTy = FunctionType::get(Type::getVoidTy(C), ArgTypes, false);
+
+  // Functions start out with lazy arguments.
+  std::unique_ptr<Function> F(
+      Function::Create(FTy, GlobalValue::ExternalLinkage, "F"));
+  EXPECT_TRUE(F->hasLazyArguments());
+
+  // Checking for empty or size shouldn't force arguments to be instantiated.
+  EXPECT_FALSE(F->arg_empty());
+  EXPECT_TRUE(F->hasLazyArguments());
+  EXPECT_EQ(2u, F->arg_size());
+  EXPECT_TRUE(F->hasLazyArguments());
+
+  // The argument list should be populated at first access.
+  (void)F->arg_begin();
+  EXPECT_FALSE(F->hasLazyArguments());
+}
+
+TEST(FunctionTest, stealArgumentListFrom) {
+  LLVMContext C;
+
+  Type *ArgTypes[] = {Type::getInt8Ty(C), Type::getInt32Ty(C)};
+  FunctionType *FTy = FunctionType::get(Type::getVoidTy(C), ArgTypes, false);
+  std::unique_ptr<Function> F1(
+      Function::Create(FTy, GlobalValue::ExternalLinkage, "F1"));
+  std::unique_ptr<Function> F2(
+      Function::Create(FTy, GlobalValue::ExternalLinkage, "F1"));
+  EXPECT_TRUE(F1->hasLazyArguments());
+  EXPECT_TRUE(F2->hasLazyArguments());
+
+  // Steal arguments before they've been accessed.  Nothing should change; both
+  // functions should still have lazy arguments.
+  //
+  //   steal(empty); drop (empty)
+  F1->stealArgumentListFrom(*F2);
+  EXPECT_TRUE(F1->hasLazyArguments());
+  EXPECT_TRUE(F2->hasLazyArguments());
+
+  // Save arguments from F1 for later assertions.  F1 won't have lazy arguments
+  // anymore.
+  SmallVector<Argument *, 4> Args;
+  for (Argument &A : F1->args())
+    Args.push_back(&A);
+  EXPECT_EQ(2u, Args.size());
+  EXPECT_FALSE(F1->hasLazyArguments());
+
+  // Steal arguments from F1 to F2.  F1's arguments should be lazy again.
+  //
+  //   steal(real); drop (empty)
+  F2->stealArgumentListFrom(*F1);
+  EXPECT_TRUE(F1->hasLazyArguments());
+  EXPECT_FALSE(F2->hasLazyArguments());
+  unsigned I = 0;
+  for (Argument &A : F2->args())
+    EXPECT_EQ(Args[I++], &A);
+  EXPECT_EQ(2u, I);
+
+  // Check that arguments in F1 don't have pointer equality with the saved ones.
+  // This also instantiates F1's arguments.
+  I = 0;
+  for (Argument &A : F1->args())
+    EXPECT_NE(Args[I++], &A);
+  EXPECT_EQ(2u, I);
+  EXPECT_FALSE(F1->hasLazyArguments());
+  EXPECT_FALSE(F2->hasLazyArguments());
+
+  // Steal back from F2.  F2's arguments should be lazy again.
+  //
+  //   steal(real); drop (real)
+  F1->stealArgumentListFrom(*F2);
+  EXPECT_FALSE(F1->hasLazyArguments());
+  EXPECT_TRUE(F2->hasLazyArguments());
+  I = 0;
+  for (Argument &A : F1->args())
+    EXPECT_EQ(Args[I++], &A);
+  EXPECT_EQ(2u, I);
+
+  // Steal from F2 a second time.  Now both functions should have lazy
+  // arguments.
+  //
+  //   steal(empty); drop (real)
+  F1->stealArgumentListFrom(*F2);
+  EXPECT_TRUE(F1->hasLazyArguments());
+  EXPECT_TRUE(F2->hasLazyArguments());
+}
+
+} // end namespace