PNaCl: Add ExpandVarArgs pass for expanding out variable-args function calls

lib/Transforms/NaCl/ExpandVarArgs.cpp

Index: lib/Transforms/NaCl/ExpandVarArgs.cpp
diff --git a/lib/Transforms/NaCl/ExpandVarArgs.cpp b/lib/Transforms/NaCl/ExpandVarArgs.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..0a353bbd33614136b847a8383aaa5ffa76e243b7
--- /dev/null
+++ b/lib/Transforms/NaCl/ExpandVarArgs.cpp
@@ -0,0 +1,318 @@
+//===- ExpandVarArgs.cpp - Expand out variable argument function calls-----===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass expands out all use of variable argument functions.
+//
+// This pass replaces a varargs function call with a function call in
+// which a pointer to the variable arguments is passed explicitly.
+// The callee explicitly allocates space for the variable arguments on
+// the stack using "alloca".
+//
+// Alignment:
+//
+// This pass does not add any alignment padding between the arguments
+// that are copied onto the stack.  We assume that the only argument
+// types that need to be handled are 32-bit and 64-bit -- i32, i64,
+// pointers and double:
+//
+//  * We won't see i1, i8, i16 and float as varargs arguments because
+//    the C standard requires the compiler to promote these to the
+//    types "int" and "double".
+//
+//  * We won't see va_arg instructions of struct type because Clang
+//    does not yet support them in PNaCl mode.  See
+//    https://code.google.com/p/nativeclient/issues/detail?id=2381
+//
+// If such arguments do appear in the input, this pass will generate
+// correct, working code, but this code might be inefficient due to
+// using unaligned memory accesses.
+//
+//===----------------------------------------------------------------------===//
+
+#include <vector>
+
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Transforms/NaCl.h"
+
+using namespace llvm;
+
+namespace {
+  class ExpandVarArgs : public ModulePass {
+  public:
+    static char ID; // Pass identification, replacement for typeid
+    ExpandVarArgs() : ModulePass(ID) {
+      initializeExpandVarArgsPass(*PassRegistry::getPassRegistry());
+    }
+
+    virtual bool runOnModule(Module &M);
+  };
+}
+
+char ExpandVarArgs::ID = 0;
+INITIALIZE_PASS(ExpandVarArgs, "expand-varargs",
+                "Expand out variable argument function definitions and calls",
+                false, false)
+
+static void ExpandVarArgFunc(Function *Func) {

[eliben, 2013/03/21 16:10:11]

Since you only call this from runOnModule, why not pass the Module* in, instead
of ->getParent().

[Mark Seaborn, 2013/03/21 18:07:55]

Using getParent() seems to be idiomatic in LLVM.  e.g. grep for
"getParent.*getParent".  It simplifies the call site, it reduces the chance of
having out-of-sync arguments, and it reduces the chance of ending up with an
unused argument after refactoring.

But I did notice that getContext() works at all levels of the hierarchy, not
just Modules, so that lets me remove many getParent()s.

As an aside, it's odd that there's no getModule() or getFunction() in addition
to getParent().

[eliben, 2013/03/21 18:36:40]

Alas, LLVM is full of cargo-cult programming so grepping is not an indication of
good practices. However, the reasons you provide make sense to me so I'm not
insisting.
There are also some cases (see lib/Analysis/AliasAnalysisCounter.cpp and also
Lint.cpp) where the Module is saved to a member variable in the beginning of
runOn* and then used. This is shorter and likely more efficient than calling
getParent() many times.

I'm deferring it to your judgement what to choose. The code is good either way.

+  Module *Module = Func->getParent();
+  Type *PtrType = Type::getInt8Ty(Module->getContext())->getPointerTo();

[eliben, 2013/03/21 16:10:11]

Type::getInt8PtrTy

[Mark Seaborn, 2013/03/21 18:07:55]

Done.

+
+  FunctionType *FTy = Func->getFunctionType();
+  std::vector<Type*> Params(FTy->param_begin(), FTy->param_end());

[eliben, 2013/03/21 16:10:11]

Did you consider SmallVector? It can be more efficient here because it 99% of
the cases there won't be more than, say, 8 parameters

[Mark Seaborn, 2013/03/21 18:07:55]

Done.

+  Params.push_back(PtrType);
+  FunctionType *NFTy = FunctionType::get(FTy->getReturnType(), Params, false);
+
+  // In order to change the function's arguments, we have to recreate
+  // the function.
+  Function *NewFunc = Function::Create(NFTy, Func->getLinkage());
+  NewFunc->copyAttributesFrom(Func);
+  Func->getParent()->getFunctionList().insert(Func, NewFunc);

[eliben, 2013/03/21 16:10:11]

Again, would be nice to just have a Module* here instead of repeatedly calling
Func->getParent

[Mark Seaborn, 2013/03/21 18:07:55]

I don't think that's idiomatic.  Note that I based this code on
DeleteDeadVarargs() in lib/Transforms/IPO/DeadArgumentElimination.cpp: 
http://llvm.org/docs/doxygen/html/DeadArgumentElimination_8cpp_source.html

+  NewFunc->takeName(Func);
+  NewFunc->getBasicBlockList().splice(NewFunc->begin(),
+                                      Func->getBasicBlockList());
+
+  // Move the arguments across to the new function.
+  for (Function::arg_iterator Arg = Func->arg_begin(), E = Func->arg_end(),
+         NewArg = NewFunc->arg_begin();
+       Arg != E; ++Arg, ++NewArg) {
+    Arg->replaceAllUsesWith(NewArg);
+    NewArg->takeName(Arg);
+  }
+
+  Func->replaceAllUsesWith(

[eliben, 2013/03/21 16:10:11]

Would it also be useful to preserve the debug information for the function?

[Mark Seaborn, 2013/03/21 18:07:55]

Yes; I'll have to investigate this.

[Mark Seaborn, 2013/03/21 19:50:54]

AFAICT, replaceAllUsesWith() preserves the debug metadata's reference to the
function OK, and I see debug info generated in the assembly output OK.

I've added setDebugLoc() calls to preserve debug line numbers for the
instructions we rewrite.  Please take a look to see if the CopyDebug() helper is
OK with you.

+      ConstantExpr::getBitCast(NewFunc, FTy->getPointerTo()));
+  Func->eraseFromParent();
+
+  Value *VarArgsArg = --NewFunc->arg_end();
+  VarArgsArg->setName("varargs");
+
+  // Expand out uses of llvm.va_start in this function.
+  for (Function::iterator BB = NewFunc->begin(), E = NewFunc->end();
+       BB != E;
+       ++BB) {
+    for (BasicBlock::iterator Iter = BB->begin(), E = BB->end();
+         Iter != E; ) {
+      Instruction *Inst = Iter++;
+      if (VAStartInst *VAS = dyn_cast<VAStartInst>(Inst)) {
+        Value *Cast = new BitCastInst(VAS->getArgList(),
+                                      PtrType->getPointerTo(),
+                                      "arglist", VAS);
+        new StoreInst(VarArgsArg, Cast, VAS);
+        VAS->eraseFromParent();
+      }
+    }
+  }
+}
+
+static void ExpandVAArgInst(VAArgInst *Inst) {
+  Module *Module = Inst->getParent()->getParent()->getParent();

[eliben, 2013/03/21 16:10:11]

Here it's probably even better reason to pass Module in, or Context if you only
need that.

Note that it's also fine to save the Module the pass is operating on in a
private member if you don't want to pass it around.

[Mark Seaborn, 2013/03/21 18:07:55]

Cleaned up to use Inst->getContext().

+  LLVMContext *Context = &Module->getContext();
+
+  // Read the argument.  We assume that no realignment of the pointer
+  // is required.
+  Value *ArgList = new BitCastInst(
+      Inst->getPointerOperand(),
+      Inst->getType()->getPointerTo()->getPointerTo(), "arglist", Inst);
+  Value *CurrentPtr = new LoadInst(ArgList, "arglist_current", Inst);
+  Value *Result = new LoadInst(CurrentPtr, "va_arg", Inst);
+  Result->takeName(Inst);
+
+  // Update the va_list to point to the next argument.
+  std::vector<Value*> Indexes;
+  Indexes.push_back(ConstantInt::get(*Context, APInt(32, 1)));
+  Value *Next = GetElementPtrInst::Create(CurrentPtr, Indexes,

[eliben, 2013/03/21 16:10:11]

Is it worthwhile to create such a GEP if you're going to expand it later anyway?
[btw, note this imposes a pass running order]

[Mark Seaborn, 2013/03/21 18:07:55]

It's certainly easier to use GEP and have a later pass expand it out for us. 
Note that we use a struct for packing the arguments when calling a varargs
function.

I think it's OK for the pass ordering to matter.  e.g. Whether we expand out
struct arguments first or varargs first will affect the calling convention for
how structs are passed through varargs.

+                                          "arglist_next", Inst);
+  new StoreInst(Next, ArgList, Inst);
+
+  Inst->replaceAllUsesWith(Result);
+  Inst->eraseFromParent();
+}
+
+static void ExpandVACopyInst(VACopyInst *Inst, Module *M) {
+  // va_list may have more space reserved, but we only need to
+  // copy a single pointer.
+  Type *I8 = Type::getInt8Ty(M->getContext());
+  Type *PtrTy = I8->getPointerTo()->getPointerTo();
+  Value *Src = new BitCastInst(Inst->getSrc(), PtrTy, "vacopy_src", Inst);
+  Value *Dest = new BitCastInst(Inst->getDest(), PtrTy, "vacopy_dest", Inst);
+  Value *CurrentPtr = new LoadInst(Src, "vacopy_currentptr", Inst);
+  new StoreInst(CurrentPtr, Dest, Inst);
+  Inst->eraseFromParent();
+}
+
+static void LifetimeDecl(Intrinsic::ID id, Value *Ptr, Value *Size,
+                         Instruction *InsertPt) {
+  Module *M = InsertPt->getParent()->getParent()->getParent();
+  Value *Func = Intrinsic::getDeclaration(M, id);
+  std::vector<Value*> Args;

[eliben, 2013/03/21 16:10:11]

Here and elsewhere, use SmallVector. It was created in LLVM to avoid heap
allocation for small std::vectors

[Mark Seaborn, 2013/03/21 18:07:55]

Done.

+  Args.push_back(Size);
+  Args.push_back(Ptr);
+  CallInst::Create(Func, Args, "", InsertPt);
+}
+
+// CopyCall() uses argument overloading so that it can be used by the
+// template ExpandVarArgCall().
+static Instruction *CopyCall(CallInst *Original, Value *Callee,
+                             ArrayRef<Value*> Args) {
+  return CallInst::Create(Callee, Args, "", Original);
+}
+
+static Instruction *CopyCall(InvokeInst *Original, Value *Callee,
+                             ArrayRef<Value*> Args) {
+  return InvokeInst::Create(Callee, Original->getNormalDest(),
+                            Original->getUnwindDest(), Args, "", Original);
+}
+
+template <class InstType>
+static bool ExpandVarArgCall(InstType *Call, DataLayout *DL) {

[eliben, 2013/03/21 16:10:11]

Please add a comment explaining what this does

[Mark Seaborn, 2013/03/21 18:07:55]

Done.

+  FunctionType *FuncType = cast<FunctionType>(
+      Call->getCalledValue()->getType()->getPointerElementType());
+  if (!FuncType->isFunctionVarArg())
+    return false;
+
+  LLVMContext *Context =
+      &Call->getParent()->getParent()->getParent()->getContext();
+
+  // Split argument list into fixed and variable arguments.
+  std::vector<Value*> FixedArgs;
+  std::vector<Value*> VarArgs;
+  std::vector<Type*> VarArgsTypes;
+  for (unsigned I = 0; I < FuncType->getNumParams(); ++I)
+    FixedArgs.push_back(Call->getArgOperand(I));
+  for (unsigned I = FuncType->getNumParams();
+       I < Call->getNumArgOperands(); ++I) {
+    VarArgs.push_back(Call->getArgOperand(I));
+    VarArgsTypes.push_back(Call->getArgOperand(I)->getType());
+  }
+
+  StructType *VarArgsTy;
+  Value *ArgToAdd;
+  Instruction *BufPtr = NULL;
+  Value *BufSize = NULL;
+  if (VarArgs.size() == 0) {
+    // If there are no variable arguments being passed, we still want
+    // to add an extra argument to the function call so that the
+    // number of arguments matches the callee's type.
+    VarArgsTy = StructType::get(*Context);
+    ArgToAdd = UndefValue::get(VarArgsTy->getPointerTo());
+  } else {
+    // Create struct type for packing variable arguments into.  We
+    // create this as packed for now and assume that no alignment
+    // padding is desired.
+    VarArgsTy = StructType::create(VarArgsTypes, "vararg_call", true);
+
+    // Allocate space for the variable argument buffer.  Do this at the
+    // start of the function so that we don't leak space if the function
+    // is called in a loop.
+    Function *Func = Call->getParent()->getParent();
+    Instruction *Buf = new AllocaInst(VarArgsTy, "vararg_buffer");
+    Func->getEntryBlock().getInstList().push_front(Buf);
+    ArgToAdd = Buf;
+
+    // Call llvm.lifetime.start/end intrinsics to indicate that Buf is
+    // only used for the duration of the function call, so that the
+    // stack space can be reused elsewhere.
+    Type *I8Ptr = Type::getInt8Ty(*Context)->getPointerTo();
+    BufPtr = new BitCastInst(Buf, I8Ptr, "vararg_lifetime_bitcast");
+    BufPtr->insertAfter(Buf);
+    BufSize = ConstantInt::get(*Context,
+                               APInt(64, DL->getTypeAllocSize(VarArgsTy)));

[eliben, 2013/03/21 16:10:11]

I'm not sure you need APInts  here and elsewhere. Won't a uint64_t do? 

[Mark Seaborn, 2013/03/21 18:07:55]

There are two ways of creating an i64 ConstantInt:

ConstantInt::get(*Context, APInt(64, x))
ConstantInt::get(Type::getInt64Ty(*Context), x))

I don't know of a simpler way.

See http://llvm.org/docs/doxygen/html/classllvm_1_1ConstantInt.html

I originally did the second, but I think the APInt version is more concise.  I
can change it back if you'd like?

[eliben, 2013/03/21 18:36:40]

I initially thought that precomputing all useful types can be more efficient,
but then I looked in the code of ConstantInt::get and it calls itself with a
uint wrapped in APInt anyway (The member variable holding the value in
ConstantInt is a APInt), so it doesn't matter I guess

+    LifetimeDecl(Intrinsic::lifetime_start, BufPtr, BufSize, Call);
+
+    // Copy variable arguments into buffer.
+    int Index = 0;
+    for (std::vector<Value*>::iterator Iter = VarArgs.begin();
+         Iter != VarArgs.end();
+         ++Iter, ++Index) {
+      std::vector<Value*> Indexes;
+      Indexes.push_back(ConstantInt::get(*Context, APInt(32, 0)));
+      Indexes.push_back(ConstantInt::get(*Context, APInt(32, Index)));
+      Value *Ptr = GetElementPtrInst::Create(Buf, Indexes, "vararg_ptr", Call);
+      new StoreInst(*Iter, Ptr, Call);
+    }
+  }
+
+  // Cast function to new type to add our extra pointer argument.
+  std::vector<Type*> ArgTypes(FuncType->param_begin(), FuncType->param_end());
+  ArgTypes.push_back(VarArgsTy->getPointerTo());
+  FunctionType *NFTy = FunctionType::get(FuncType->getReturnType(),
+                                         ArgTypes, false);
+  Value *CastFunc = new BitCastInst(Call->getCalledValue(),
+                                    NFTy->getPointerTo(), "vararg_func", Call);
+
+  // Create the converted function call.
+  FixedArgs.push_back(ArgToAdd);
+  Value *NewCall = CopyCall(Call, CastFunc, FixedArgs);
+  NewCall->takeName(Call);
+
+  if (BufPtr) {
+    if (isa<CallInst>(Call)) {
+      LifetimeDecl(Intrinsic::lifetime_end, BufPtr, BufSize, Call);
+    } else if (InvokeInst *Invoke = dyn_cast<InvokeInst>(Call)) {
+      LifetimeDecl(Intrinsic::lifetime_end, BufPtr, BufSize,
+                   Invoke->getNormalDest()->getFirstInsertionPt());
+      LifetimeDecl(Intrinsic::lifetime_end, BufPtr, BufSize,
+                   Invoke->getUnwindDest()->getFirstInsertionPt());
+    }
+  }
+
+  Call->replaceAllUsesWith(NewCall);
+  Call->eraseFromParent();
+
+  return true;
+}
+
+bool ExpandVarArgs::runOnModule(Module &M) {
+  bool Changed = false;
+  DataLayout DL(&M);
+
+  for (Module::iterator Iter = M.begin(), E = M.end(); Iter != E; ) {
+    Function *Func = Iter++;
+
+    for (Function::iterator BB = Func->begin(), E = Func->end();
+         BB != E;
+         ++BB) {
+      for (BasicBlock::iterator Iter = BB->begin(), E = BB->end();
+           Iter != E; ) {
+        Instruction *Inst = Iter++;
+        if (VAArgInst *VI = dyn_cast<VAArgInst>(Inst)) {
+          Changed = true;
+          ExpandVAArgInst(VI);
+        } else if (isa<VAEndInst>(Inst)) {
+          // va_end() is a no-op in this implementation.
+          Changed = true;
+          Inst->eraseFromParent();
+        } else if (VACopyInst *VAC = dyn_cast<VACopyInst>(Inst)) {
+          Changed = true;
+          ExpandVACopyInst(VAC, &M);
+        } else if (CallInst *Call = dyn_cast<CallInst>(Inst)) {
+          Changed |= ExpandVarArgCall(Call, &DL);
+        } else if (InvokeInst *Call = dyn_cast<InvokeInst>(Inst)) {
+          Changed |= ExpandVarArgCall(Call, &DL);
+        }
+      }
+    }
+
+    if (Func->isVarArg()) {
+      Changed = true;
+      ExpandVarArgFunc(Func);
+    }
+  }
+
+  return Changed;
+}
+
+ModulePass *llvm::createExpandVarArgsPass() {
+  return new ExpandVarArgs();
+}