Adding a pass - RewritePNaClLibraryCalls, that replaces known library calls with stable bitcode int…

lib/Transforms/NaCl/RewritePNaClLibraryCalls.cpp

+//===- RewritePNaClLibraryCalls.cpp - PNaCl library calls to intrinsics ---===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass replaces calls to known library functions with calls to intrinsics
+// that are part of the PNaCl stable bitcode ABI.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/NaCl.h"
+
+using namespace llvm;
+
+namespace {
+  class RewritePNaClLibraryCalls : public ModulePass {
+  public:
+    static char ID;
+    RewritePNaClLibraryCalls() :
+      ModulePass(ID), TheModule(0),

[Mark Seaborn, 2013/05/13 20:09:00]

Indent by 2 more spaces so that this is more indented than the method body

[eliben, 2013/05/13 22:07:18]

Done.

+      LongjmpFunc(0), LongjmpWrapperFunc(0) {

[Mark Seaborn, 2013/05/13 20:09:00]

0 -> NULL?

[eliben, 2013/05/13 22:07:18]

Done.

+      // This is a module pass because it may have to introduce
+      // intrinsic declarations into the module, and rewrite global
+      // initializers.
+      initializeRewritePNaClLibraryCallsPass(*PassRegistry::getPassRegistry());
+    }
+
+    virtual bool runOnModule(Module &M);
+  private:
+    /// Rewrites the given \p Call to setjmp.
+    void rewriteSetjmpCall(CallInst *Call);
+
+    /// Rewrites the given \p Call to longjmp.
+    void rewriteLongjmpCall(CallInst *Call);
+
+    /// Rewrites the given \p Store instruction that stores a pointer to

[Mark Seaborn, 2013/05/13 20:09:00]

Why are you checking specifically for Store instructions?  What about other
uses, e.g. ptrtoint?

+    /// longjmp.
+    void rewriteLongjmpStore(StoreInst *Store);
+
+    void rewriteAddrOfLongjmpInGlobal(GlobalVariable *Global);

[Mark Seaborn, 2013/05/13 20:09:00]

This also doesn't make sense.  You should handle all uses generically, with
replacesAllUsesWith().

+
+    /// Creates the longjmp wrapper on-demand.
+    Function *getOrCreateLongjmpWrapper();
+
+    /// The module this pass runs on.
+    Module *TheModule;
+
+    /// Global longjmp function the original code calls. NULL if it's not in
+    /// the module.
+    Function *LongjmpFunc;
+
+    /// The wrapper for longjmp is created on-demand and cached here.
+    Function *LongjmpWrapperFunc;
+  };
+}
+
+static const char *LONGJMP_NAME = "longjmp";
+static const char *SETJMP_NAME = "setjmp";
+static const char *WRAP_FUNC_PREFIX = "__nacl_wrap_";
+
+char RewritePNaClLibraryCalls::ID = 0;
+INITIALIZE_PASS(RewritePNaClLibraryCalls, "rewrite-pnacl-library-calls",
+                "Rewrite PNaCl library calls to stable intrinsics",
+                false, false)
+
+bool RewritePNaClLibraryCalls::runOnModule(Module &M) {
+  TheModule = &M;
+  bool Changed = false;
+
+  // Iterate over all uses of the setjmp and longjmp functions, if they exist
+  // in the module with external linkage.
+  // For setjmp, calls are replaced by calls to intrinsics.
+  // For longjmp calls are also replaced; however, longjmp can have its
+  // address taken, so in these places an address of a wrapper function is
+  // stored instead (or used as an initializer in a global variable).
+  // The wrapper function calls the intrinsic with an appropriate cast.
+  Function *SetjmpFunc = TheModule->getFunction(SETJMP_NAME);
+  if (SetjmpFunc && SetjmpFunc->hasExternalLinkage()) {
+    for (Value::use_iterator UI = SetjmpFunc->use_begin(),
+                             UE = SetjmpFunc->use_end(); UI != UE;) {
+      Value *Use = *UI++;
+      if (CallInst *Call = dyn_cast<CallInst>(Use)) {
+        rewriteSetjmpCall(Call);
+        Changed = true;
+      } else if (isa<StoreInst>(Use) || isa<GlobalVariable>(Use)) {
+        report_fatal_error("Taking the address of setjmp is invalid");
+      } else {
+        report_fatal_error("Invalid use of setjmp");
+      }
+    }
+  }
+
+  LongjmpFunc = TheModule->getFunction(LONGJMP_NAME);
+  if (LongjmpFunc && LongjmpFunc->hasExternalLinkage()) {
+    for (Value::use_iterator UI = LongjmpFunc->use_begin(),
+                             UE = LongjmpFunc->use_end(); UI != UE;) {
+      Value *Use = *UI++;
+      if (CallInst *Call = dyn_cast<CallInst>(Use)) {
+        rewriteLongjmpCall(Call);
+        Changed = true;
+      } else if (StoreInst *Store = dyn_cast<StoreInst>(Use)) {
+        rewriteLongjmpStore(Store);
+        Changed = true;
+      } else if (GlobalVariable *Global = dyn_cast<GlobalVariable>(Use)) {
+        rewriteAddrOfLongjmpInGlobal(Global);
+        Changed = true;
+      } else {
+        report_fatal_error("Invalid use of longjmp");
+      }
+    }
+  }
+
+  return Changed;
+}
+
+void RewritePNaClLibraryCalls::rewriteSetjmpCall(CallInst *Call) {
+  FunctionType *FTy = Call->getCalledFunction()->getFunctionType();
+
+  // Sanity check on the setjmp call.
+  if (FTy->getNumParams() != 1 ||
+      !FTy->getReturnType()->isIntegerTy() ||
+      !FTy->getParamType(0)->isPointerTy()) {
+    report_fatal_error(Twine("Wrong signature of function ") +
+                       Call->getCalledFunction()->getName());
+  }
+  const DebugLoc &DLoc = Call->getDebugLoc();
+
+  // Find the intrinsic function.
+  Function *NaClSetjmpFunc = Intrinsic::getDeclaration(TheModule,
+                                                       Intrinsic::nacl_setjmp);
+  // Cast the jmp_buf argument to the type NaClSetjmpCall expects.
+  Type *PtrTy = NaClSetjmpFunc->getFunctionType()->getParamType(0);
+  BitCastInst *JmpBufCast = new BitCastInst(Call->getArgOperand(0), PtrTy,
+                                            "jmp_buf_i8", Call);
+  JmpBufCast->setDebugLoc(DLoc);
+
+  // Emit the updated call.
+  SmallVector<Value *, 2> Args;

[Mark Seaborn, 2013/05/13 20:20:42]

Nit: 1 will do here instead of 2

[eliben, 2013/05/13 22:07:18]

Done.

+  Args.push_back(JmpBufCast);
+  CallInst *NaClSetjmpCall = CallInst::Create(NaClSetjmpFunc, Args, "", Call);
+  NaClSetjmpCall->setDebugLoc(DLoc);
+  NaClSetjmpCall->takeName(Call);
+
+  // Replace the original call.
+  Call->replaceAllUsesWith(NaClSetjmpCall);
+  Call->eraseFromParent();
+}
+
+void RewritePNaClLibraryCalls::rewriteLongjmpCall(CallInst *Call) {
+  FunctionType *FTy = Call->getCalledFunction()->getFunctionType();
+
+  // Sanity check on the longjmp call.
+  if (FTy->getNumParams() != 2 ||
+      !FTy->getReturnType()->isVoidTy() ||
+      !FTy->getParamType(0)->isPointerTy()) {
+    report_fatal_error(Twine("Wrong signature of function ") +
+                       Call->getCalledFunction()->getName());
+  }
+  const DebugLoc &DLoc = Call->getDebugLoc();
+
+  // Find the intrinsic function.
+  Function *NaClLongjmpFunc = Intrinsic::getDeclaration(
+      TheModule, Intrinsic::nacl_longjmp);
+  // Cast the jmp_buf argument to the type NaClLongjmpCall expects.
+  Type *PtrTy = NaClLongjmpFunc->getFunctionType()->getParamType(0);
+  BitCastInst *JmpBufCast = new BitCastInst(Call->getArgOperand(0), PtrTy,
+                                            "jmp_buf_i8", Call);
+  JmpBufCast->setDebugLoc(DLoc);
+
+  // Emit the call.
+  SmallVector<Value *, 2> Args;
+  Args.push_back(JmpBufCast);
+  Args.push_back(Call->getArgOperand(1));
+  CallInst *NaClLongjmpCall = CallInst::Create(NaClLongjmpFunc, Args, "", Call);
+  NaClLongjmpCall->setDebugLoc(DLoc);
+  // No takeName here since longjmp is a void call that does not get assigned to
+  // a value.
+
+  // Remove the original call. There's no need for RAUW because longjmp
+  // returns void.
+  Call->eraseFromParent();
+}
+
+void RewritePNaClLibraryCalls::rewriteLongjmpStore(StoreInst *Store) {
+  // Sanity check: this method was called for a store that uses a function
+  // pointer to setjmp.
+  Function *CalledFunc = dyn_cast_or_null<Function>(Store->getValueOperand());
+  if (CalledFunc != LongjmpFunc) {
+    report_fatal_error(Twine("Badly formed store of ") +
+                       LongjmpFunc->getName());
+  }
+  
+  Function *LongjmpWrapperFunc = getOrCreateLongjmpWrapper();
+
+  // Replace Store by a new store that stores a pointer to the wrapper.

[Mark Seaborn, 2013/05/13 20:20:42]

You should bitcast your wrapper function rather than recreating store
instructions and globals with a different type.  You can use a bitcast
ConstantExpr (this pass can run before ExpandConstantExpr).

Actually, you don't even need a bitcast, since your wrapper function has the
same type as the original longjmp().  See comment below...

+  StoreInst *WrapperStore = new StoreInst(LongjmpWrapperFunc,
+                                          Store->getPointerOperand(),
+                                          Store);
+  WrapperStore->setDebugLoc(Store->getDebugLoc());
+  WrapperStore->takeName(Store);
+  Store->replaceAllUsesWith(WrapperStore);
+  Store->eraseFromParent();
+}
+
+void RewritePNaClLibraryCalls::rewriteAddrOfLongjmpInGlobal(
+    GlobalVariable *Global) {
+  if (!Global->hasUniqueInitializer()) {
+    report_fatal_error(Twine("Badly formed assignment to global ") +
+                             Global->getName());
+  }
+
+  Function *LongjmpWrapperFunc = getOrCreateLongjmpWrapper();
+  
+  GlobalVariable *NewGlobal = new GlobalVariable(
+      /* M                */ *TheModule,
+      /* Ty               */ LongjmpWrapperFunc->getType(),
+      /* isConstant       */ Global->isConstant(),
+      /* Linkage          */ GlobalValue::PrivateLinkage,
+      /* Initializer      */ LongjmpWrapperFunc,
+      /* Name             */ "",
+      /* InsertBefore     */ Global,
+      /* ThreadLocalMode  */ Global->getThreadLocalMode());
+  NewGlobal->setAlignment(Global->getAlignment());
+  // Using takeName instead of assigning Name in the constructor of
+  // GlobalVariable because otherwise a numeric suffix is added to the name,
+  // while we want it to stay exactly the same.
+  NewGlobal->takeName(Global);
+  Global->eraseFromParent();
+}
+
+Function *RewritePNaClLibraryCalls::getOrCreateLongjmpWrapper() {
+  if (LongjmpWrapperFunc)
+    return LongjmpWrapperFunc;
+
+  // A wrapper function is created that has the same type as longjmp; it

[Mark Seaborn, 2013/05/13 20:20:42]

You actually don't need to create a new function.

A function declaration is just a function with no basic blocks.  You can just
insert a basic block into the existing longjmp() function.

Then your wrapper will still be named longjmp() (without a takeName() call). 
You might need to change the linkage to be internal.

[eliben, 2013/05/13 22:07:18]

Done.

+  // redirects the call to the intrinsic after casting pointer argument.
+  FunctionType *WrapFuncTy = LongjmpFunc->getFunctionType();
+  SmallString<32> Name(WRAP_FUNC_PREFIX); Name.append(LONGJMP_NAME);

[Mark Seaborn, 2013/05/13 20:20:42]

Split onto separate lines.

[eliben, 2013/05/13 22:07:18]

Function removed.

+  Function *LongjmpWrapperFunc = dyn_cast<Function>(
+      TheModule->getOrInsertFunction(Name, WrapFuncTy));
+  LongjmpWrapperFunc->setLinkage(Function::PrivateLinkage);
+
+  // This function has just been created. Fill its contents.
+  LLVMContext &Context = TheModule->getContext();
+  BasicBlock *BB = BasicBlock::Create(Context, "entry", LongjmpWrapperFunc);
+
+  Function::arg_iterator LongjmpArgs = LongjmpWrapperFunc->arg_begin();
+  Value *EnvArg = LongjmpArgs++;
+  EnvArg->setName("env");
+  Value *ValArg = LongjmpArgs++;
+  ValArg->setName("val");
+
+  // Insert a return instruction

[Mark Seaborn, 2013/05/13 20:20:42]

Nit: could insert 'unreachable' instead

[eliben, 2013/05/13 22:07:18]

Done.

+  Instruction *Ret = ReturnInst::Create(Context, 0, BB);
+  
+  // Find the intrinsic function.
+  Function *NaClLongjmpFunc = Intrinsic::getDeclaration(
+      TheModule, Intrinsic::nacl_longjmp);
+  // Cast the jmp_buf argument to the type NaClLongjmpCall expects.
+  Type *PtrTy = NaClLongjmpFunc->getFunctionType()->getParamType(0);
+  BitCastInst *JmpBufCast = new BitCastInst(EnvArg,  PtrTy, "jmp_buf_i8", Ret);
+
+  // Emit the call.
+  SmallVector<Value *, 2> Args;
+  Args.push_back(JmpBufCast);
+  Args.push_back(ValArg);
+  CallInst::Create(NaClLongjmpFunc, Args, "", Ret);
+ 
+  return LongjmpWrapperFunc;
+}
+
+ModulePass *llvm::createRewritePNaClLibraryCallsPass() {
+  return new RewritePNaClLibraryCalls();
+}
+