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),
+      SetjmpFunc(0), LongjmpFunc(0), LongjmpWrapperFunc(0) {
+      // This is a module pass because it may have to introduce
+      // intrindic declarations into the module.

[jvoung (off chromium), 2013/05/13 18:12:30]

intrindic -> intrinsic

Another reason (if you want to add to the comment) is that it may need to
rewrite global intializers that take the address of longjmp.

[eliben, 2013/05/13 18:33:36]

Done.

+      initializeExpandTlsPass(*PassRegistry::getPassRegistry());

[jvoung (off chromium), 2013/05/13 18:12:30]

initializeExpandTlsPass?

[eliben, 2013/05/13 18:33:36]

Oops :-/ 

Thanks, fixed.

+    }
+
+    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
+    /// longjmp.
+    void rewriteLongjmpStore(StoreInst *Store);
+
+    void rewriteAddrOfLongjmpInGlobal(GlobalVariable *Global);
+
+    /// Creates the longjmp wrapper on-demand.
+    Function *getOrCreateLongjmpWrapper();
+
+    /// The module this pass runs on.
+    Module *TheModule;
+
+    /// Global setjmp/longjmp functions the original code calls. NULL if they
+    /// are not in the module.
+    Function *SetjmpFunc;

[jvoung (off chromium), 2013/05/13 18:12:30]

Does SetjmpFunc really need to be a field?  It looks like it's only used in the
runOnModule function.

LongjmpFunc could be passed along from runOnModule to the rewrite functions as
an argument, so it doesn't really need to be a field either.

Then you wouldn't need to initialize these to 0.

The LongjmpWrapperFunc makes sense, since that is state from the getOrCreate...
method.

[eliben, 2013/05/13 18:33:36]

Yeah, SetjmpFunc is definitely a leftover from refactoring. I'll remove it.
LongjmpFunc though is used in getOrCreateLongjmpWrapper too so I don't know if
it's worth piping through a bunch of places. Do you think that would be better?

+    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

[jvoung (off chromium), 2013/05/13 18:12:30]

no need for two spaces between ";" and "however"?

[eliben, 2013/05/13 18:33:36]

Done.

+  // 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.
+  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);

[jvoung (off chromium), 2013/05/13 18:12:30]

Now that you aren't checking "Call->getCalledFunction()->getName();", should it
check?

Is it possible to have a use of setjmp as a function parameter?  E.g., if you
have "printf("my setjmp is at addr: %p\n", &setjmp);"

[eliben, 2013/05/13 18:33:36]

Not sure what you mean. Only going over the uses of setjmp here.
No, it's illegal to take the address of setjmp in any way

[jvoung (off chromium), 2013/05/13 19:52:07]

I was wondering if you had IL like this:

%call = call i32 (i8*, ...)* @printf(i8* getelementptr inbounds ([4 x i8]*
@.str, i32 0, i32 0), i32 (i8*)* @setjmp)

Would that be considered a "user" of setjmp, within a CallInst, even though
setjmp is an operand, and not the called function.  I guess not, because it's
only an operand, and not a value defined by a call to setjmp?


It seems like you are also relying on the ABI checker rejecting constant exprs
(and the const expr expansion pass running), otherwise the address could have
been taken in arbitrary places (like the above IL).

+        Changed = true;
+      } else if (isa<StoreInst>(Use)) {
+        report_fatal_error("Invalid store of setjmp's address");
+      }

[jvoung (off chromium), 2013/05/13 18:12:30]

Do we expect to see other uses of setjmp (else branch -> error)?

Same with longjmp (else branch -> error)?

[eliben, 2013/05/13 18:33:36]

Done.

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