<commit_before>#include "CfgBuilderUtils.hh"

#include "llvm/ADT/APInt.h"
#include "llvm/IR/CallSite.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"

namespace clam {

using namespace ikos;
using namespace llvm;

// Any integer that cannot be represented by 64 bits is considered a bignum.
bool isSignedBigNum(const APInt &v) {
  unsigned b = v.getBitWidth();
  if (b <= 64) {
    return false;
  } else {
    // if bitwidth > 64 then we check the actual value
    APInt max(b, APInt::getSignedMaxValue(64).getSExtValue(), true);
    APInt min(b, APInt::getSignedMinValue(64).getSExtValue(), true);
    return (v.sgt(max) || v.slt(min));
  }
}

bool isBool(const Type *t) { return (t->isIntegerTy(1)); }

bool isBool(const Value &v) { return isBool(v.getType()); }

bool isInteger(const Type *t) { return (t->isIntegerTy() && !isBool(t)); }

bool isInteger(const Value &v) { return isInteger(v.getType()); }

bool isReference(const Type *t, const CrabBuilderParams &params) {
  return (t->isPointerTy() && params.trackMemory());
}

bool isReference(const Value &v, const CrabBuilderParams &params) {
  return isReference(v.getType(), params);
}

z_number toZNumber(const APInt &v, const CrabBuilderParams &params,
                   bool &is_bignum) {
  is_bignum = false;
  if (!params.enable_bignums) {
    is_bignum = isSignedBigNum(v);
  }
#if 0
  // Convert to strings is not ideal but it shouldn't be a big
  // bottleneck.
  std::string val = v.toString(10,true /*is signed*/);
  return z_number(val);
#else
  // Based on:
  // https://llvm.org/svn/llvm-project/polly/trunk/lib/Support/GICHelper.cpp
  APInt abs;
  abs = v.isNegative() ? v.abs() : v;
  const uint64_t *rawdata = abs.getRawData();
  unsigned numWords = abs.getNumWords();

  ikos::z_number res;
  mpz_import(res.get_mpz_t(), numWords, -1, sizeof(uint64_t), 0, 0, rawdata);
  return v.isNegative() ? -res : res;
#endif
}

z_number getIntConstant(const ConstantInt *CI, const CrabBuilderParams &params,
                        bool &is_bignum) {
  is_bignum = false;
  if (CI->getType()->isIntegerTy(1)) {
    return z_number((int64_t)CI->getZExtValue());
  } else {
    return toZNumber(CI->getValue(), params, is_bignum);
  }
}

bool isTrackedType(const Type &ty, const CrabBuilderParams &params) {
  return isReference(&ty, params) || ty.isIntegerTy();
}

bool isTracked(const Value &v, const CrabBuilderParams &params) {
  // -- ignore any shadow variable created by seahorn
  // if (v.getName().startswith("shadow.mem"))
  // return false;

  return isTrackedType(*v.getType(), params);
}

bool ShouldCallSiteReturn(CallInst &I, const CrabBuilderParams &params) {
  CallSite CS(&I);
  if (Function *Callee = CS.getCalledFunction()) {
    Type *RT = Callee->getReturnType();
    return (!(RT->isVoidTy()) && isTrackedType(*RT, params));
  }
  return false;
}

bool DoesCallSiteReturn(CallInst &I, const CrabBuilderParams &params) {
  return (!I.getType()->isVoidTy() && isTracked(I, params));
}

bool hasDebugLoc(const Instruction *inst) {
  if (!inst)
    return false;
  const DebugLoc &dloc = inst->getDebugLoc();
  return dloc;
}

crab::cfg::debug_info getDebugLoc(const Instruction *inst) {
  if (!hasDebugLoc(inst))
    return crab::cfg::debug_info();
  const DebugLoc &dloc = inst->getDebugLoc();
  unsigned Line = dloc.getLine();
  unsigned Col = dloc.getCol();
  std::string File = (*dloc).getFilename();
  if (File == "")
    File = "unknown file";

  // HACK: append assertion ID to the filename for bookeeping
  if (MDNode *MDN = inst->getMetadata("clam-assertion")) {
    StringRef Check = cast<MDString>(MDN->getOperand(0))->getString();
    if (Check != "") {
      File += "#" + Check.str();
    }
  }
  return crab::cfg::debug_info(File, Line, Col);
}

uint64_t storageSize(const Type *t, const DataLayout &dl) {
  return dl.getTypeStoreSize(const_cast<Type *>(t));
}

void normalizeCmpInst(CmpInst &I) {
  switch (I.getPredicate()) {
  case ICmpInst::ICMP_UGT:
  case ICmpInst::ICMP_SGT:
    I.swapOperands();
    break;
  case ICmpInst::ICMP_UGE:
  case ICmpInst::ICMP_SGE:
    I.swapOperands();
    break;
  default:;
  }
}

bool isIntToBool(const CastInst &I) {
  return (isa<TruncInst>(I) && I.getDestTy()->isIntegerTy(1));
}

bool isBoolToInt(const CastInst &I) {
  return ((isa<ZExtInst>(I) || isa<SExtInst>(I)) &&
          I.getSrcTy()->isIntegerTy(1));
}

bool isBoolArray(const Type &T) {
  return (T.isArrayTy() && T.getArrayElementType()->isIntegerTy(1));
}

bool isIntArray(const Type &T) {
  return (T.isArrayTy() && T.getArrayElementType()->isIntegerTy() &&
          !(T.getArrayElementType()->isIntegerTy(1)));
}

// bool isPointerArray(const Type &T) {
//   return (T.isArrayTy() && T.getArrayElementType()->isPointerTy());
// }

bool isAssertFn(const Function &F) {
  return (F.getName().equals("verifier.assert") ||
          F.getName().equals("crab.assert") ||
          F.getName().equals("__CRAB_assert"));
}

bool isSeaHornFail(const Function &F) {
  return (F.getName().equals("seahorn.fail"));
}

bool isErrorFn(const Function &F) {
  return (F.getName().equals("seahorn.error") ||
          F.getName().equals("verifier.error") ||
          F.getName().equals("__VERIFIER_error") ||
          F.getName().equals("__SEAHORN_error"));
}

bool isAssumeFn(const Function &F) {
  return (F.getName().equals("verifier.assume") ||
          F.getName().equals("__VERIFIER_assume") ||
          F.getName().equals("__CRAB_assume"));
}

bool isNotAssumeFn(const Function &F) {
  return (F.getName().equals("verifier.assume.not") ||
          F.getName().equals("__VERIFIER_assume_not") ||
          F.getName().equals("__CRAB_assume_not"));
}

bool isVerifierCall(const Function &F) {
  return (isAssertFn(F) || isErrorFn(F) || isAssumeFn(F) || isNotAssumeFn(F) ||
          isSeaHornFail(F));
}

bool isCrabIntrinsic(const Function &F) {
  return (F.getName().startswith("__CRAB_intrinsic_"));
}

std::string getCrabIntrinsicName(const Function &F) {
  assert(isCrabIntrinsic(F));
  StringRef res = F.getName().split("__CRAB_intrinsic_").second;
  return res.str();
}

bool isZeroInitializer(const Function &F) {
  return F.getName().startswith("verifier.zero_initializer");
}

bool isZeroInitializer(const CallInst &CI) {
  ImmutableCallSite CS(&CI);
  const Value *calleeV = CS.getCalledValue();
  if (const Function *callee =
          dyn_cast<Function>(calleeV->stripPointerCasts())) {
    return isZeroInitializer(*callee);
  }
  return false;
}

bool isIntInitializer(const Function &F) {
  return F.getName().startswith("verifier.int_initializer");
}

bool isIntInitializer(const CallInst &CI) {
  ImmutableCallSite CS(&CI);
  const Value *calleeV = CS.getCalledValue();
  if (const Function *callee =
          dyn_cast<Function>(calleeV->stripPointerCasts())) {
    return isIntInitializer(*callee);
  }
  return false;
}

// Return true if all uses are BranchInst's
bool AllUsesAreBrInst(Value &V) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses())
    if (!isa<BranchInst>(U.getUser()))
      return false;
  return true;
}

// Return true if all uses are BranchInst's or Select's
bool AllUsesAreBrOrIntSelectCondInst(Value &V, const CrabBuilderParams &params) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses()) {
    if ((!isa<BranchInst>(U.getUser())) && (!isa<SelectInst>(U.getUser())))
      return false;
    if (SelectInst *SI = dyn_cast<SelectInst>(U.getUser())) {
      if (isBool(*SI) || SI->getCondition() != &V ||
	  isReference(*SI, params)) {
        // if the operands are bool or V is not the condition
        return false;
      }
    }
  }
  return true;
}

// Return true if all uses are the callee at callsites
bool AllUsesAreIndirectCalls(Value &V) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses()) {
    if (CallInst *CI = dyn_cast<CallInst>(U.getUser())) {
      CallSite CS(CI);
      const Value *callee = CS.getCalledValue();
      if (callee == &V)
        continue;
    }
    return false;
  }
  return true;
}

// Return true if all uses are verifier calls (assume/assert)
bool AllUsesAreVerifierCalls(Value &V,
			     bool goThroughIntegerCasts, bool nonBoolCond,
			     SmallVector<CallInst*, 4> &verifierCalls) {
  for (auto &U : V.uses()) {
    Value *User = U.getUser();
    if (goThroughIntegerCasts) {
      if (isa<ZExtInst>(User) || isa<SExtInst>(User)) {
	return AllUsesAreVerifierCalls(*User,
				       goThroughIntegerCasts, nonBoolCond,
				       verifierCalls);
      }
    }

    if (CallInst *CI = dyn_cast<CallInst>(User)) {
      CallSite CS(CI);
      const Value *calleeV = CS.getCalledValue();
      const Function *callee = dyn_cast<Function>(calleeV->stripPointerCasts());
      if (callee && (isAssertFn(*callee) || isAssumeFn(*callee) ||
                     isNotAssumeFn(*callee))) {
	if (nonBoolCond) {
	  FunctionType *FTy = callee->getFunctionType();
	  if (!FTy->isVarArg() &&
	      FTy->getReturnType()->isVoidTy() &&
	      FTy->getNumParams() == 1 && !isBool(FTy->getParamType(0))) {
	    verifierCalls.push_back(CI);
	    continue;
	  }
	} else {
	  verifierCalls.push_back(CI);
	  continue;
	}
      }
    }
    verifierCalls.clear();
    return false;
  }
  return true;
}

bool AllUsesAreVerifierCalls(Value &V) {
  SmallVector<CallInst*, 4> verifierCalls /*unused*/;
  return AllUsesAreVerifierCalls(V, false, false, verifierCalls);
}

// Return true if all uses are GEPs
bool AllUsesAreGEP(Value &V) {
  for (auto &U : V.uses())
    if (!isa<GetElementPtrInst>(U.getUser()))
      return false;
  return true;
}

bool AllUsesAreIgnoredInst(llvm::Value &V) {
  for (auto &U : V.uses()) {
    if (CallInst *CI = dyn_cast<CallInst>(U.getUser())) {
      if (Function * CalledF = dyn_cast<Function>(CI->getCalledOperand())) {
	if (CalledF->getName().startswith("llvm.dbg.value") ||
	    CalledF->getName().startswith("llvm.lifetime")) {
	  continue;
	}
      }
    }
    return false;
  }
  return true;
}

} // namespace clam
<commit_msg>fix: strip pointer casts when looking for callee<commit_after>#include "CfgBuilderUtils.hh"

#include "llvm/ADT/APInt.h"
#include "llvm/IR/CallSite.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/DebugLoc.h"
#include "llvm/IR/Type.h"
#include "llvm/IR/Value.h"

namespace clam {

using namespace ikos;
using namespace llvm;

// Any integer that cannot be represented by 64 bits is considered a bignum.
bool isSignedBigNum(const APInt &v) {
  unsigned b = v.getBitWidth();
  if (b <= 64) {
    return false;
  } else {
    // if bitwidth > 64 then we check the actual value
    APInt max(b, APInt::getSignedMaxValue(64).getSExtValue(), true);
    APInt min(b, APInt::getSignedMinValue(64).getSExtValue(), true);
    return (v.sgt(max) || v.slt(min));
  }
}

bool isBool(const Type *t) { return (t->isIntegerTy(1)); }

bool isBool(const Value &v) { return isBool(v.getType()); }

bool isInteger(const Type *t) { return (t->isIntegerTy() && !isBool(t)); }

bool isInteger(const Value &v) { return isInteger(v.getType()); }

bool isReference(const Type *t, const CrabBuilderParams &params) {
  return (t->isPointerTy() && params.trackMemory());
}

bool isReference(const Value &v, const CrabBuilderParams &params) {
  return isReference(v.getType(), params);
}

z_number toZNumber(const APInt &v, const CrabBuilderParams &params,
                   bool &is_bignum) {
  is_bignum = false;
  if (!params.enable_bignums) {
    is_bignum = isSignedBigNum(v);
  }
#if 0
  // Convert to strings is not ideal but it shouldn't be a big
  // bottleneck.
  std::string val = v.toString(10,true /*is signed*/);
  return z_number(val);
#else
  // Based on:
  // https://llvm.org/svn/llvm-project/polly/trunk/lib/Support/GICHelper.cpp
  APInt abs;
  abs = v.isNegative() ? v.abs() : v;
  const uint64_t *rawdata = abs.getRawData();
  unsigned numWords = abs.getNumWords();

  ikos::z_number res;
  mpz_import(res.get_mpz_t(), numWords, -1, sizeof(uint64_t), 0, 0, rawdata);
  return v.isNegative() ? -res : res;
#endif
}

z_number getIntConstant(const ConstantInt *CI, const CrabBuilderParams &params,
                        bool &is_bignum) {
  is_bignum = false;
  if (CI->getType()->isIntegerTy(1)) {
    return z_number((int64_t)CI->getZExtValue());
  } else {
    return toZNumber(CI->getValue(), params, is_bignum);
  }
}

bool isTrackedType(const Type &ty, const CrabBuilderParams &params) {
  return isReference(&ty, params) || ty.isIntegerTy();
}

bool isTracked(const Value &v, const CrabBuilderParams &params) {
  // -- ignore any shadow variable created by seahorn
  // if (v.getName().startswith("shadow.mem"))
  // return false;

  return isTrackedType(*v.getType(), params);
}

bool ShouldCallSiteReturn(CallInst &I, const CrabBuilderParams &params) {
  CallSite CS(&I);
  if (Function *Callee = dyn_cast<Function>(CS.getCalledValue()->stripPointerCasts())) {
    Type *RT = Callee->getReturnType();
    return (!(RT->isVoidTy()) && isTrackedType(*RT, params));
  }
  return false;
}

bool DoesCallSiteReturn(CallInst &I, const CrabBuilderParams &params) {
  return (!I.getType()->isVoidTy() && isTracked(I, params));
}

bool hasDebugLoc(const Instruction *inst) {
  if (!inst)
    return false;
  const DebugLoc &dloc = inst->getDebugLoc();
  return dloc;
}

crab::cfg::debug_info getDebugLoc(const Instruction *inst) {
  if (!hasDebugLoc(inst))
    return crab::cfg::debug_info();
  const DebugLoc &dloc = inst->getDebugLoc();
  unsigned Line = dloc.getLine();
  unsigned Col = dloc.getCol();
  std::string File = (*dloc).getFilename();
  if (File == "")
    File = "unknown file";

  // HACK: append assertion ID to the filename for bookeeping
  if (MDNode *MDN = inst->getMetadata("clam-assertion")) {
    StringRef Check = cast<MDString>(MDN->getOperand(0))->getString();
    if (Check != "") {
      File += "#" + Check.str();
    }
  }
  return crab::cfg::debug_info(File, Line, Col);
}

uint64_t storageSize(const Type *t, const DataLayout &dl) {
  return dl.getTypeStoreSize(const_cast<Type *>(t));
}

void normalizeCmpInst(CmpInst &I) {
  switch (I.getPredicate()) {
  case ICmpInst::ICMP_UGT:
  case ICmpInst::ICMP_SGT:
    I.swapOperands();
    break;
  case ICmpInst::ICMP_UGE:
  case ICmpInst::ICMP_SGE:
    I.swapOperands();
    break;
  default:;
  }
}

bool isIntToBool(const CastInst &I) {
  return (isa<TruncInst>(I) && I.getDestTy()->isIntegerTy(1));
}

bool isBoolToInt(const CastInst &I) {
  return ((isa<ZExtInst>(I) || isa<SExtInst>(I)) &&
          I.getSrcTy()->isIntegerTy(1));
}

bool isBoolArray(const Type &T) {
  return (T.isArrayTy() && T.getArrayElementType()->isIntegerTy(1));
}

bool isIntArray(const Type &T) {
  return (T.isArrayTy() && T.getArrayElementType()->isIntegerTy() &&
          !(T.getArrayElementType()->isIntegerTy(1)));
}

// bool isPointerArray(const Type &T) {
//   return (T.isArrayTy() && T.getArrayElementType()->isPointerTy());
// }

bool isAssertFn(const Function &F) {
  return (F.getName().equals("verifier.assert") ||
          F.getName().equals("crab.assert") ||
          F.getName().equals("__CRAB_assert"));
}

bool isSeaHornFail(const Function &F) {
  return (F.getName().equals("seahorn.fail"));
}

bool isErrorFn(const Function &F) {
  return (F.getName().equals("seahorn.error") ||
          F.getName().equals("verifier.error") ||
          F.getName().equals("__VERIFIER_error") ||
          F.getName().equals("__SEAHORN_error"));
}

bool isAssumeFn(const Function &F) {
  return (F.getName().equals("verifier.assume") ||
          F.getName().equals("__VERIFIER_assume") ||
          F.getName().equals("__CRAB_assume"));
}

bool isNotAssumeFn(const Function &F) {
  return (F.getName().equals("verifier.assume.not") ||
          F.getName().equals("__VERIFIER_assume_not") ||
          F.getName().equals("__CRAB_assume_not"));
}

bool isVerifierCall(const Function &F) {
  return (isAssertFn(F) || isErrorFn(F) || isAssumeFn(F) || isNotAssumeFn(F) ||
          isSeaHornFail(F));
}

bool isCrabIntrinsic(const Function &F) {
  return (F.getName().startswith("__CRAB_intrinsic_"));
}

std::string getCrabIntrinsicName(const Function &F) {
  assert(isCrabIntrinsic(F));
  StringRef res = F.getName().split("__CRAB_intrinsic_").second;
  return res.str();
}

bool isZeroInitializer(const Function &F) {
  return F.getName().startswith("verifier.zero_initializer");
}

bool isZeroInitializer(const CallInst &CI) {
  ImmutableCallSite CS(&CI);
  const Value *calleeV = CS.getCalledValue();
  if (const Function *callee =
          dyn_cast<Function>(calleeV->stripPointerCasts())) {
    return isZeroInitializer(*callee);
  }
  return false;
}

bool isIntInitializer(const Function &F) {
  return F.getName().startswith("verifier.int_initializer");
}

bool isIntInitializer(const CallInst &CI) {
  ImmutableCallSite CS(&CI);
  const Value *calleeV = CS.getCalledValue();
  if (const Function *callee =
          dyn_cast<Function>(calleeV->stripPointerCasts())) {
    return isIntInitializer(*callee);
  }
  return false;
}

// Return true if all uses are BranchInst's
bool AllUsesAreBrInst(Value &V) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses())
    if (!isa<BranchInst>(U.getUser()))
      return false;
  return true;
}

// Return true if all uses are BranchInst's or Select's
bool AllUsesAreBrOrIntSelectCondInst(Value &V, const CrabBuilderParams &params) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses()) {
    if ((!isa<BranchInst>(U.getUser())) && (!isa<SelectInst>(U.getUser())))
      return false;
    if (SelectInst *SI = dyn_cast<SelectInst>(U.getUser())) {
      if (isBool(*SI) || SI->getCondition() != &V ||
	  isReference(*SI, params)) {
        // if the operands are bool or V is not the condition
        return false;
      }
    }
  }
  return true;
}

// Return true if all uses are the callee at callsites
bool AllUsesAreIndirectCalls(Value &V) {
  // XXX: do not strip pointers here
  for (auto &U : V.uses()) {
    if (CallInst *CI = dyn_cast<CallInst>(U.getUser())) {
      CallSite CS(CI);
      const Value *callee = CS.getCalledValue();
      if (callee == &V)
        continue;
    }
    return false;
  }
  return true;
}

// Return true if all uses are verifier calls (assume/assert)
bool AllUsesAreVerifierCalls(Value &V,
			     bool goThroughIntegerCasts, bool nonBoolCond,
			     SmallVector<CallInst*, 4> &verifierCalls) {
  for (auto &U : V.uses()) {
    Value *User = U.getUser();
    if (goThroughIntegerCasts) {
      if (isa<ZExtInst>(User) || isa<SExtInst>(User)) {
	return AllUsesAreVerifierCalls(*User,
				       goThroughIntegerCasts, nonBoolCond,
				       verifierCalls);
      }
    }

    if (CallInst *CI = dyn_cast<CallInst>(User)) {
      CallSite CS(CI);
      const Value *calleeV = CS.getCalledValue();
      const Function *callee = dyn_cast<Function>(calleeV->stripPointerCasts());
      if (callee && (isAssertFn(*callee) || isAssumeFn(*callee) ||
                     isNotAssumeFn(*callee))) {
	if (nonBoolCond) {
	  FunctionType *FTy = callee->getFunctionType();
	  if (!FTy->isVarArg() &&
	      FTy->getReturnType()->isVoidTy() &&
	      FTy->getNumParams() == 1 && !isBool(FTy->getParamType(0))) {
	    verifierCalls.push_back(CI);
	    continue;
	  }
	} else {
	  verifierCalls.push_back(CI);
	  continue;
	}
      }
    }
    verifierCalls.clear();
    return false;
  }
  return true;
}

bool AllUsesAreVerifierCalls(Value &V) {
  SmallVector<CallInst*, 4> verifierCalls /*unused*/;
  return AllUsesAreVerifierCalls(V, false, false, verifierCalls);
}

// Return true if all uses are GEPs
bool AllUsesAreGEP(Value &V) {
  for (auto &U : V.uses())
    if (!isa<GetElementPtrInst>(U.getUser()))
      return false;
  return true;
}

bool AllUsesAreIgnoredInst(llvm::Value &V) {
  for (auto &U : V.uses()) {
    if (CallInst *CI = dyn_cast<CallInst>(U.getUser())) {
      if (Function * CalledF = dyn_cast<Function>(CI->getCalledOperand())) {
	if (CalledF->getName().startswith("llvm.dbg.value") ||
	    CalledF->getName().startswith("llvm.lifetime")) {
	  continue;
	}
      }
    }
    return false;
  }
  return true;
}

} // namespace clam
<|endoftext|>


<commit_before>//===-- MachineSink.cpp - Sinking for machine instructions ----------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This pass moves instructions into successor blocks, when possible, so that
// they aren't executed on paths where their results aren't needed.
//
// This pass is not intended to be a replacement or a complete alternative
// for an LLVM-IR-level sinking pass. It is only designed to sink simple
// constructs that are not exposed before lowering and instruction selection.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "machine-sink"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/MachineDominators.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Target/TargetRegisterInfo.h"
#include "llvm/Target/TargetInstrInfo.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

STATISTIC(NumSunk, "Number of machine instructions sunk");

namespace {
  class MachineSinking : public MachineFunctionPass {
    const TargetInstrInfo *TII;
    const TargetRegisterInfo *TRI;
    MachineRegisterInfo  *RegInfo; // Machine register information
    MachineDominatorTree *DT;   // Machine dominator tree
    AliasAnalysis *AA;
    BitVector AllocatableSet;   // Which physregs are allocatable?

  public:
    static char ID; // Pass identification
    MachineSinking() : MachineFunctionPass(&ID) {}

    virtual bool runOnMachineFunction(MachineFunction &MF);

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.setPreservesCFG();
      MachineFunctionPass::getAnalysisUsage(AU);
      AU.addRequired<AliasAnalysis>();
      AU.addRequired<MachineDominatorTree>();
      AU.addPreserved<MachineDominatorTree>();
    }
  private:
    bool ProcessBlock(MachineBasicBlock &MBB);
    bool SinkInstruction(MachineInstr *MI, bool &SawStore);
    bool AllUsesDominatedByBlock(unsigned Reg, MachineBasicBlock *MBB) const;
  };
} // end anonymous namespace

char MachineSinking::ID = 0;
static RegisterPass<MachineSinking>
X("machine-sink", "Machine code sinking");

FunctionPass *llvm::createMachineSinkingPass() { return new MachineSinking(); }

/// AllUsesDominatedByBlock - Return true if all uses of the specified register
/// occur in blocks dominated by the specified block.
bool MachineSinking::AllUsesDominatedByBlock(unsigned Reg,
                                             MachineBasicBlock *MBB) const {
  assert(TargetRegisterInfo::isVirtualRegister(Reg) &&
         "Only makes sense for vregs");
  // Ignoring debug uses is necessary so debug info doesn't affect the code.
  // This may leave a referencing dbg_value in the original block, before
  // the definition of the vreg.  Dwarf generator handles this although the
  // user might not get the right info at runtime.
  for (MachineRegisterInfo::use_nodbg_iterator I =
       RegInfo->use_nodbg_begin(Reg),
       E = RegInfo->use_nodbg_end(); I != E; ++I) {
    // Determine the block of the use.
    MachineInstr *UseInst = &*I;
    MachineBasicBlock *UseBlock = UseInst->getParent();
    if (UseInst->isPHI()) {
      // PHI nodes use the operand in the predecessor block, not the block with
      // the PHI.
      UseBlock = UseInst->getOperand(I.getOperandNo()+1).getMBB();
    }
    // Check that it dominates.
    if (!DT->dominates(MBB, UseBlock))
      return false;
  }
  return true;
}

bool MachineSinking::runOnMachineFunction(MachineFunction &MF) {
  DEBUG(dbgs() << "******** Machine Sinking ********\n");

  const TargetMachine &TM = MF.getTarget();
  TII = TM.getInstrInfo();
  TRI = TM.getRegisterInfo();
  RegInfo = &MF.getRegInfo();
  DT = &getAnalysis<MachineDominatorTree>();
  AA = &getAnalysis<AliasAnalysis>();
  AllocatableSet = TRI->getAllocatableSet(MF);

  bool EverMadeChange = false;

  while (1) {
    bool MadeChange = false;

    // Process all basic blocks.
    for (MachineFunction::iterator I = MF.begin(), E = MF.end();
         I != E; ++I)
      MadeChange |= ProcessBlock(*I);

    // If this iteration over the code changed anything, keep iterating.
    if (!MadeChange) break;
    EverMadeChange = true;
  }
  return EverMadeChange;
}

bool MachineSinking::ProcessBlock(MachineBasicBlock &MBB) {
  // Can't sink anything out of a block that has less than two successors.
  if (MBB.succ_size() <= 1 || MBB.empty()) return false;

  // Don't bother sinking code out of unreachable blocks. In addition to being
  // unprofitable, it can also lead to infinite looping, because in an unreachable
  // loop there may be nowhere to stop.
  if (!DT->isReachableFromEntry(&MBB)) return false;

  bool MadeChange = false;

  // Walk the basic block bottom-up.  Remember if we saw a store.
  MachineBasicBlock::iterator I = MBB.end();
  --I;
  bool ProcessedBegin, SawStore = false;
  do {
    MachineInstr *MI = I;  // The instruction to sink.

    // Predecrement I (if it's not begin) so that it isn't invalidated by
    // sinking.
    ProcessedBegin = I == MBB.begin();
    if (!ProcessedBegin)
      --I;

    if (MI->isDebugValue())
      continue;

    if (SinkInstruction(MI, SawStore))
      ++NumSunk, MadeChange = true;

    // If we just processed the first instruction in the block, we're done.
  } while (!ProcessedBegin);

  return MadeChange;
}

/// SinkInstruction - Determine whether it is safe to sink the specified machine
/// instruction out of its current block into a successor.
bool MachineSinking::SinkInstruction(MachineInstr *MI, bool &SawStore) {
  // Check if it's safe to move the instruction.
  if (!MI->isSafeToMove(TII, AA, SawStore))
    return false;

  // FIXME: This should include support for sinking instructions within the
  // block they are currently in to shorten the live ranges.  We often get
  // instructions sunk into the top of a large block, but it would be better to
  // also sink them down before their first use in the block.  This xform has to
  // be careful not to *increase* register pressure though, e.g. sinking
  // "x = y + z" down if it kills y and z would increase the live ranges of y
  // and z and only shrink the live range of x.

  // Loop over all the operands of the specified instruction.  If there is
  // anything we can't handle, bail out.
  MachineBasicBlock *ParentBlock = MI->getParent();

  // SuccToSinkTo - This is the successor to sink this instruction to, once we
  // decide.
  MachineBasicBlock *SuccToSinkTo = 0;

  for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
    const MachineOperand &MO = MI->getOperand(i);
    if (!MO.isReg()) continue;  // Ignore non-register operands.

    unsigned Reg = MO.getReg();
    if (Reg == 0) continue;

    if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
      if (MO.isUse()) {
        // If the physreg has no defs anywhere, it's just an ambient register
        // and we can freely move its uses. Alternatively, if it's allocatable,
        // it could get allocated to something with a def during allocation.
        if (!RegInfo->def_empty(Reg))
          return false;
        if (AllocatableSet.test(Reg))
          return false;
        // Check for a def among the register's aliases too.
        for (const unsigned *Alias = TRI->getAliasSet(Reg); *Alias; ++Alias) {
          unsigned AliasReg = *Alias;
          if (!RegInfo->def_empty(AliasReg))
            return false;
          if (AllocatableSet.test(AliasReg))
            return false;
        }
      } else if (!MO.isDead()) {
        // A def that isn't dead. We can't move it.
        return false;
      }
    } else {
      // Virtual register uses are always safe to sink.
      if (MO.isUse()) continue;

      // If it's not safe to move defs of the register class, then abort.
      if (!TII->isSafeToMoveRegClassDefs(RegInfo->getRegClass(Reg)))
        return false;

      // FIXME: This picks a successor to sink into based on having one
      // successor that dominates all the uses.  However, there are cases where
      // sinking can happen but where the sink point isn't a successor.  For
      // example:
      //   x = computation
      //   if () {} else {}
      //   use x
      // the instruction could be sunk over the whole diamond for the
      // if/then/else (or loop, etc), allowing it to be sunk into other blocks
      // after that.

      // Virtual register defs can only be sunk if all their uses are in blocks
      // dominated by one of the successors.
      if (SuccToSinkTo) {
        // If a previous operand picked a block to sink to, then this operand
        // must be sinkable to the same block.
        if (!AllUsesDominatedByBlock(Reg, SuccToSinkTo))
          return false;
        continue;
      }

      // Otherwise, we should look at all the successors and decide which one
      // we should sink to.
      for (MachineBasicBlock::succ_iterator SI = ParentBlock->succ_begin(),
           E = ParentBlock->succ_end(); SI != E; ++SI) {
        if (AllUsesDominatedByBlock(Reg, *SI)) {
          SuccToSinkTo = *SI;
          break;
        }
      }

      // If we couldn't find a block to sink to, ignore this instruction.
      if (SuccToSinkTo == 0)
        return false;
    }
  }

  // If there are no outputs, it must have side-effects.
  if (SuccToSinkTo == 0)
    return false;

  // It's not safe to sink instructions to EH landing pad. Control flow into
  // landing pad is implicitly defined.
  if (SuccToSinkTo->isLandingPad())
    return false;

  // It is not possible to sink an instruction into its own block.  This can
  // happen with loops.
  if (MI->getParent() == SuccToSinkTo)
    return false;

  DEBUG(dbgs() << "Sink instr " << *MI);
  DEBUG(dbgs() << "to block " << *SuccToSinkTo);

  // If the block has multiple predecessors, this would introduce computation on
  // a path that it doesn't already exist.  We could split the critical edge,
  // but for now we just punt.
  // FIXME: Split critical edges if not backedges.
  if (SuccToSinkTo->pred_size() > 1) {
    // We cannot sink a load across a critical edge - there may be stores in
    // other code paths.
    bool store = true;
    if (!MI->isSafeToMove(TII, AA, store)) {
      DEBUG(dbgs() << " *** PUNTING: Wont sink load along critical edge.\n");
      return false;
    }

    // We don't want to sink across a critical edge if we don't dominate the
    // successor. We could be introducing calculations to new code paths.
    if (!DT->dominates(ParentBlock, SuccToSinkTo)) {
      DEBUG(dbgs() << " *** PUNTING: Critical edge found\n");
      return false;
    }

    // Otherwise we are OK with sinking along a critical edge.
    DEBUG(dbgs() << "Sinking along critical edge.\n");
  }

  // Determine where to insert into.  Skip phi nodes.
  MachineBasicBlock::iterator InsertPos = SuccToSinkTo->begin();
  while (InsertPos != SuccToSinkTo->end() && InsertPos->isPHI())
    ++InsertPos;

  // Move the instruction.
  SuccToSinkTo->splice(InsertPos, ParentBlock, MI,
                       ++MachineBasicBlock::iterator(MI));
  return true;
}
<commit_msg>Avoid sinking machine instructions into a loop.<commit_after>//===-- MachineSink.cpp - Sinking for machine instructions ----------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This pass moves instructions into successor blocks, when possible, so that
// they aren't executed on paths where their results aren't needed.
//
// This pass is not intended to be a replacement or a complete alternative
// for an LLVM-IR-level sinking pass. It is only designed to sink simple
// constructs that are not exposed before lowering and instruction selection.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "machine-sink"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/MachineDominators.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Target/TargetRegisterInfo.h"
#include "llvm/Target/TargetInstrInfo.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

STATISTIC(NumSunk, "Number of machine instructions sunk");

namespace {
  class MachineSinking : public MachineFunctionPass {
    const TargetInstrInfo *TII;
    const TargetRegisterInfo *TRI;
    MachineRegisterInfo  *RegInfo; // Machine register information
    MachineDominatorTree *DT;   // Machine dominator tree
    MachineLoopInfo *LI;
    AliasAnalysis *AA;
    BitVector AllocatableSet;   // Which physregs are allocatable?

  public:
    static char ID; // Pass identification
    MachineSinking() : MachineFunctionPass(&ID) {}

    virtual bool runOnMachineFunction(MachineFunction &MF);

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.setPreservesCFG();
      MachineFunctionPass::getAnalysisUsage(AU);
      AU.addRequired<AliasAnalysis>();
      AU.addRequired<MachineDominatorTree>();
      AU.addRequired<MachineLoopInfo>();
      AU.addPreserved<MachineDominatorTree>();
      AU.addPreserved<MachineLoopInfo>();
    }
  private:
    bool ProcessBlock(MachineBasicBlock &MBB);
    bool SinkInstruction(MachineInstr *MI, bool &SawStore);
    bool AllUsesDominatedByBlock(unsigned Reg, MachineBasicBlock *MBB) const;
  };
} // end anonymous namespace

char MachineSinking::ID = 0;
static RegisterPass<MachineSinking>
X("machine-sink", "Machine code sinking");

FunctionPass *llvm::createMachineSinkingPass() { return new MachineSinking(); }

/// AllUsesDominatedByBlock - Return true if all uses of the specified register
/// occur in blocks dominated by the specified block.
bool MachineSinking::AllUsesDominatedByBlock(unsigned Reg,
                                             MachineBasicBlock *MBB) const {
  assert(TargetRegisterInfo::isVirtualRegister(Reg) &&
         "Only makes sense for vregs");
  // Ignoring debug uses is necessary so debug info doesn't affect the code.
  // This may leave a referencing dbg_value in the original block, before
  // the definition of the vreg.  Dwarf generator handles this although the
  // user might not get the right info at runtime.
  for (MachineRegisterInfo::use_nodbg_iterator I =
       RegInfo->use_nodbg_begin(Reg),
       E = RegInfo->use_nodbg_end(); I != E; ++I) {
    // Determine the block of the use.
    MachineInstr *UseInst = &*I;
    MachineBasicBlock *UseBlock = UseInst->getParent();
    if (UseInst->isPHI()) {
      // PHI nodes use the operand in the predecessor block, not the block with
      // the PHI.
      UseBlock = UseInst->getOperand(I.getOperandNo()+1).getMBB();
    }
    // Check that it dominates.
    if (!DT->dominates(MBB, UseBlock))
      return false;
  }
  return true;
}

bool MachineSinking::runOnMachineFunction(MachineFunction &MF) {
  DEBUG(dbgs() << "******** Machine Sinking ********\n");

  const TargetMachine &TM = MF.getTarget();
  TII = TM.getInstrInfo();
  TRI = TM.getRegisterInfo();
  RegInfo = &MF.getRegInfo();
  DT = &getAnalysis<MachineDominatorTree>();
  LI = &getAnalysis<MachineLoopInfo>();
  AA = &getAnalysis<AliasAnalysis>();
  AllocatableSet = TRI->getAllocatableSet(MF);

  bool EverMadeChange = false;

  while (1) {
    bool MadeChange = false;

    // Process all basic blocks.
    for (MachineFunction::iterator I = MF.begin(), E = MF.end();
         I != E; ++I)
      MadeChange |= ProcessBlock(*I);

    // If this iteration over the code changed anything, keep iterating.
    if (!MadeChange) break;
    EverMadeChange = true;
  }
  return EverMadeChange;
}

bool MachineSinking::ProcessBlock(MachineBasicBlock &MBB) {
  // Can't sink anything out of a block that has less than two successors.
  if (MBB.succ_size() <= 1 || MBB.empty()) return false;

  // Don't bother sinking code out of unreachable blocks. In addition to being
  // unprofitable, it can also lead to infinite looping, because in an unreachable
  // loop there may be nowhere to stop.
  if (!DT->isReachableFromEntry(&MBB)) return false;

  bool MadeChange = false;

  // Walk the basic block bottom-up.  Remember if we saw a store.
  MachineBasicBlock::iterator I = MBB.end();
  --I;
  bool ProcessedBegin, SawStore = false;
  do {
    MachineInstr *MI = I;  // The instruction to sink.

    // Predecrement I (if it's not begin) so that it isn't invalidated by
    // sinking.
    ProcessedBegin = I == MBB.begin();
    if (!ProcessedBegin)
      --I;

    if (MI->isDebugValue())
      continue;

    if (SinkInstruction(MI, SawStore))
      ++NumSunk, MadeChange = true;

    // If we just processed the first instruction in the block, we're done.
  } while (!ProcessedBegin);

  return MadeChange;
}

/// SinkInstruction - Determine whether it is safe to sink the specified machine
/// instruction out of its current block into a successor.
bool MachineSinking::SinkInstruction(MachineInstr *MI, bool &SawStore) {
  // Check if it's safe to move the instruction.
  if (!MI->isSafeToMove(TII, AA, SawStore))
    return false;

  // FIXME: This should include support for sinking instructions within the
  // block they are currently in to shorten the live ranges.  We often get
  // instructions sunk into the top of a large block, but it would be better to
  // also sink them down before their first use in the block.  This xform has to
  // be careful not to *increase* register pressure though, e.g. sinking
  // "x = y + z" down if it kills y and z would increase the live ranges of y
  // and z and only shrink the live range of x.

  // Loop over all the operands of the specified instruction.  If there is
  // anything we can't handle, bail out.
  MachineBasicBlock *ParentBlock = MI->getParent();

  // SuccToSinkTo - This is the successor to sink this instruction to, once we
  // decide.
  MachineBasicBlock *SuccToSinkTo = 0;

  for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
    const MachineOperand &MO = MI->getOperand(i);
    if (!MO.isReg()) continue;  // Ignore non-register operands.

    unsigned Reg = MO.getReg();
    if (Reg == 0) continue;

    if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
      if (MO.isUse()) {
        // If the physreg has no defs anywhere, it's just an ambient register
        // and we can freely move its uses. Alternatively, if it's allocatable,
        // it could get allocated to something with a def during allocation.
        if (!RegInfo->def_empty(Reg))
          return false;
        if (AllocatableSet.test(Reg))
          return false;
        // Check for a def among the register's aliases too.
        for (const unsigned *Alias = TRI->getAliasSet(Reg); *Alias; ++Alias) {
          unsigned AliasReg = *Alias;
          if (!RegInfo->def_empty(AliasReg))
            return false;
          if (AllocatableSet.test(AliasReg))
            return false;
        }
      } else if (!MO.isDead()) {
        // A def that isn't dead. We can't move it.
        return false;
      }
    } else {
      // Virtual register uses are always safe to sink.
      if (MO.isUse()) continue;

      // If it's not safe to move defs of the register class, then abort.
      if (!TII->isSafeToMoveRegClassDefs(RegInfo->getRegClass(Reg)))
        return false;

      // FIXME: This picks a successor to sink into based on having one
      // successor that dominates all the uses.  However, there are cases where
      // sinking can happen but where the sink point isn't a successor.  For
      // example:
      //   x = computation
      //   if () {} else {}
      //   use x
      // the instruction could be sunk over the whole diamond for the
      // if/then/else (or loop, etc), allowing it to be sunk into other blocks
      // after that.

      // Virtual register defs can only be sunk if all their uses are in blocks
      // dominated by one of the successors.
      if (SuccToSinkTo) {
        // If a previous operand picked a block to sink to, then this operand
        // must be sinkable to the same block.
        if (!AllUsesDominatedByBlock(Reg, SuccToSinkTo))
          return false;
        continue;
      }

      // Otherwise, we should look at all the successors and decide which one
      // we should sink to.
      for (MachineBasicBlock::succ_iterator SI = ParentBlock->succ_begin(),
           E = ParentBlock->succ_end(); SI != E; ++SI) {
        if (AllUsesDominatedByBlock(Reg, *SI)) {
          SuccToSinkTo = *SI;
          break;
        }
      }

      // If we couldn't find a block to sink to, ignore this instruction.
      if (SuccToSinkTo == 0)
        return false;
    }
  }

  // If there are no outputs, it must have side-effects.
  if (SuccToSinkTo == 0)
    return false;

  // It's not safe to sink instructions to EH landing pad. Control flow into
  // landing pad is implicitly defined.
  if (SuccToSinkTo->isLandingPad())
    return false;

  // It is not possible to sink an instruction into its own block.  This can
  // happen with loops.
  if (MI->getParent() == SuccToSinkTo)
    return false;

  DEBUG(dbgs() << "Sink instr " << *MI);
  DEBUG(dbgs() << "to block " << *SuccToSinkTo);

  // If the block has multiple predecessors, this would introduce computation on
  // a path that it doesn't already exist.  We could split the critical edge,
  // but for now we just punt.
  // FIXME: Split critical edges if not backedges.
  if (SuccToSinkTo->pred_size() > 1) {
    // We cannot sink a load across a critical edge - there may be stores in
    // other code paths.
    bool store = true;
    if (!MI->isSafeToMove(TII, AA, store)) {
      DEBUG(dbgs() << " *** PUNTING: Wont sink load along critical edge.\n");
      return false;
    }

    // We don't want to sink across a critical edge if we don't dominate the
    // successor. We could be introducing calculations to new code paths.
    if (!DT->dominates(ParentBlock, SuccToSinkTo)) {
      DEBUG(dbgs() << " *** PUNTING: Critical edge found\n");
      return false;
    }

    // Don't sink instructions into a loop.
    if (LI->isLoopHeader(SuccToSinkTo)) {
      DEBUG(dbgs() << " *** PUNTING: Loop header found\n");
      return false;
    }

    // Otherwise we are OK with sinking along a critical edge.
    DEBUG(dbgs() << "Sinking along critical edge.\n");
  }

  // Determine where to insert into.  Skip phi nodes.
  MachineBasicBlock::iterator InsertPos = SuccToSinkTo->begin();
  while (InsertPos != SuccToSinkTo->end() && InsertPos->isPHI())
    ++InsertPos;

  // Move the instruction.
  SuccToSinkTo->splice(InsertPos, ParentBlock, MI,
                       ++MachineBasicBlock::iterator(MI));
  return true;
}
<|endoftext|>


<commit_before>//  This example is a generalization of the on/off problem.

/*
FourBin Instructional Tutorial:
 authors:
 Kyle Cranmer <cranmer@cern.ch>
 Tanja Rommerskirchen <tanja.rommerskirchen@cern.ch>

 date: June 1, 2010

 This example is a generalization of the on/off problem.
It's a common setup for SUSY searches.  Imagine that one has two
variables "x" and "y" (eg. missing ET and SumET), see figure.
The signal region has high values of both of these variables (top right).
One can see low values of "x" or "y" acting as side-bands.  If we
just used "y" as a sideband, we would have the on/off problem.
 - In the signal region we observe non events and expect s+b events.
 - In the region with low values of "y" (bottom right)
   we observe noff events and expect tau*b events.
Note the significance of tau.  In the background only case:
   tau ~ <expectation off> / <expectation on>
If tau is known, this model is sufficient, but often tau is not known exactly.
So one can use low values of "x" as an additional constraint for tau.
Note that this technique critically depends on the notion that the
joint distribution for "x" and "y" can be factorized.
Generally, these regions have many events, so it the ratio can be
measured very precisely there.  So we extend the model to describe the
left two boxes... denoted with "bar".
  - In the upper left we observe nonbar events and expect bbar events
  - In the bottom left we observe noffbar events and expect tau bbar events
Note again we have:
   tau ~ <expecation off bar> / <expectation on bar>
One can further expand the model to account for the systematic associated
to assuming the distribution of "x" and "y" factorizes (eg. that
tau is the same for off/on and offbar/onbar). This can be done in several
ways, but here we introduce an additional parameter rho, which so that
one set of models will use tau and the other tau*rho. The choice is arbitary,
but it has consequences on the numerical stability of the algorithms.
The "bar" measurements typically have more events (& smaller relative errors).
If we choose <expectation noffbar> = tau * rho * <expectation noonbar>, the
product tau*rho will be known very precisely (~1/sqrt(bbar)) and the contour
in those parameters will be narrow and have a non-trivial tau~1/rho shape.
However, if we choose to put rho on the non/noff measurements (where the
product will have an error ~1/sqrt(b)), the contours will be more ameanable
to numerical techniques.  Thus, here we choose to define
   tau := <expecation off bar> / (<expectation on bar>)
   rho := <expecation off> / (<expectation on> * tau)

^ y
|
|---------------------------+
|               |           |
|     nonbar    |    non    |
|      bbar     |    s+b    |
|               |           |
|---------------+-----------|
|               |           |
|    noffbar    |    noff   |
|    tau bbar   | tau b rho |
|               |           |
+-----------------------------> x

Left in this way, the problem is under-constrained.  However, one may
have some auxiliary measurement (usually based on Monte Carlo) to
constrain rho.  Let us call this auxiliary measurement that gives
the nominal value of rho "rhonom".  Thus, there is a 'constraint' term in
the full model: P(rhonom | rho).  In this case, we consider a Gaussian
constraint with standard deviation sigma.

In the example, the initial values of the parameters are
  - s    = 40
  - b    = 100
  - tau  = 5
  - bbar = 1000
  - rho  = 1
  (sigma for rho = 20%)
and in the toy dataset:
   - non = 139
   - noff = 528
   - nonbar = 993
   - noffbar = 4906
   - rhonom = 1.27824

Note, the covariance matrix of the parameters has large off-diagonal terms.
Clearly s,b are anti-correlated.  Similary, since noffbar >> nonbar, one would
expect bbar,tau to be anti-correlated.

This can be seen below.
            GLOBAL      b    bbar   rho      s     tau
        b  0.96820   1.000  0.191 -0.942 -0.762 -0.209
     bbar  0.91191   0.191  1.000  0.000 -0.146 -0.912
      rho  0.96348  -0.942  0.000  1.000  0.718 -0.000
        s  0.76250  -0.762 -0.146  0.718  1.000  0.160
      tau  0.92084  -0.209 -0.912 -0.000  0.160  1.000

Similarly, since tau*rho appears as a product, we expect rho,tau
to be anti-correlated. When the error on rho is significantly
larger than 1/sqrt(bbar), tau is essentially known and the
correlation is minimal (tau mainly cares about bbar, and rho about b,s).
In the alternate parametrizaton (bbar* tau * rho) the correlation coefficient
for rho,tau is large (and negative).

The code below uses best-practices for RooFit & RooStats as of
June 2010.  It proceeds as follows:
 - create a workspace to hold the model
 - use workspace factory to quickly create the terms of the model
 - use workspace factory to define total model (a prod pdf)
 - create a RooStats ModelConfig to specify observables, parameters of interest
 - add to the ModelConfig a prior on the parameters for Bayesian techniques
   - note, the pdf it is factorized for parameters of interest & nuisance params
 - visualize the model
 - write the workspace to a file
 - use several of RooStats IntervalCalculators & compare results
*/

#include "TStopwatch.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "RooPlot.h"
#include "RooAbsPdf.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooGlobalFunc.h"
#include "RooFitResult.h"
#include "RooRandom.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/BayesianCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/SimpleInterval.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/PointSetInterval.h"

using namespace RooFit;
using namespace RooStats;

void FourBinInstructional(bool doBayesian=false, bool doFeldmanCousins=false, bool doMCMC=false){

  // let's time this challenging example
  TStopwatch t;
  t.Start();

  // set RooFit random seed for reproducible results
  RooRandom::randomGenerator()->SetSeed(4357);

  // make model
  RooWorkspace* wspace = new RooWorkspace("wspace");
  wspace->factory("Poisson::on(non[0,1000], sum::splusb(s[40,0,100],b[100,0,300]))");
  wspace->factory("Poisson::off(noff[0,5000], prod::taub(b,tau[5,3,7],rho[1,0,2]))");
  wspace->factory("Poisson::onbar(nonbar[0,10000], bbar[1000,500,2000])");
  wspace->factory("Poisson::offbar(noffbar[0,1000000], prod::lambdaoffbar(bbar, tau))");
  wspace->factory("Gaussian::mcCons(rhonom[1.,0,2], rho, sigma[.2])");
  wspace->factory("PROD::model(on,off,onbar,offbar,mcCons)");
  wspace->defineSet("obs","non,noff,nonbar,noffbar,rhonom");

  wspace->factory("Uniform::prior_poi({s})");
  wspace->factory("Uniform::prior_nuis({b,bbar,tau, rho})");
  wspace->factory("PROD::prior(prior_poi,prior_nuis)");

  ///////////////////////////////////////////
  // Control some interesting variations
  // define parameers of interest
  // for 1-d plots
  wspace->defineSet("poi","s");
  wspace->defineSet("nuis","b,tau,rho,bbar");
  // for 2-d plots to inspect correlations:
  //  wspace->defineSet("poi","s,rho");

  // test simpler cases where parameters are known.
  //  wspace->var("tau")->setConstant();
  //  wspace->var("rho")->setConstant();
  //  wspace->var("b")->setConstant();
  //  wspace->var("bbar")->setConstant();

  // inspect workspace
  //  wspace->Print();

  ////////////////////////////////////////////////////////////
  // Generate toy data
  // generate toy data assuming current value of the parameters
  // import into workspace.
  // add Verbose() to see how it's being generated
  RooDataSet* data =   wspace->pdf("model")->generate(*wspace->set("obs"),1);
  //  data->Print("v");
  wspace->import(*data);

  /////////////////////////////////////////////////////
  // Now the statistical tests
  // model config
  ModelConfig* modelConfig = new ModelConfig("FourBins");
  modelConfig->SetWorkspace(*wspace);
  modelConfig->SetPdf(*wspace->pdf("model"));
  modelConfig->SetPriorPdf(*wspace->pdf("prior"));
  modelConfig->SetParametersOfInterest(*wspace->set("poi"));
  modelConfig->SetNuisanceParameters(*wspace->set("nuis"));
  wspace->import(*modelConfig);
  wspace->writeToFile("FourBin.root");

  //////////////////////////////////////////////////
  // If you want to see the covariance matrix uncomment
  //  wspace->pdf("model")->fitTo(*data);

  // use ProfileLikelihood
  ProfileLikelihoodCalculator plc(*data, *modelConfig);
  plc.SetConfidenceLevel(0.95);
  LikelihoodInterval* plInt = plc.GetInterval();
  RooFit::MsgLevel msglevel = RooMsgService::instance().globalKillBelow();
  RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);
  plInt->LowerLimit( *wspace->var("s") ); // get ugly print out of the way. Fix.
  RooMsgService::instance().setGlobalKillBelow(msglevel);

  // use FeldmaCousins (takes ~20 min)
  FeldmanCousins fc(*data, *modelConfig);
  fc.SetConfidenceLevel(0.95);
  //number counting: dataset always has 1 entry with N events observed
  fc.FluctuateNumDataEntries(false);
  fc.UseAdaptiveSampling(true);
  fc.SetNBins(40);
  PointSetInterval* fcInt = NULL;
  if(doFeldmanCousins){ // takes 7 minutes
    fcInt = (PointSetInterval*) fc.GetInterval(); // fix cast
  }

  // use BayesianCalculator (only 1-d parameter of interest, slow for this problem)
  BayesianCalculator bc(*data, *modelConfig);
  bc.SetConfidenceLevel(0.95);
  SimpleInterval* bInt = NULL;
  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    bInt = bc.GetInterval();
  } else{
    cout << "Bayesian Calc. only supports on parameter of interest" << endl;
  }

  // use MCMCCalculator  (takes about 1 min)
  // Want an efficient proposal function, so derive it from covariance
  // matrix of fit
  RooFitResult* fit = wspace->pdf("model")->fitTo(*data,Save());
  ProposalHelper ph;
  ph.SetVariables((RooArgSet&)fit->floatParsFinal());
  ph.SetCovMatrix(fit->covarianceMatrix());
  ph.SetUpdateProposalParameters(kTRUE); // auto-create mean vars and add mappings
  ph.SetCacheSize(100);
  ProposalFunction* pf = ph.GetProposalFunction();

  MCMCCalculator mc(*data, *modelConfig);
  mc.SetConfidenceLevel(0.95);
  mc.SetProposalFunction(*pf);
  mc.SetNumBurnInSteps(500); // first N steps to be ignored as burn-in
  mc.SetNumIters(50000);
  mc.SetLeftSideTailFraction(0.5); // make a central interval
  MCMCInterval* mcInt = NULL;
  if(doMCMC)
    mcInt = mc.GetInterval();

  //////////////////////////////////////
  // Make some  plots
  TCanvas* c1 = (TCanvas*) gROOT->Get("c1");
  if(!c1)
    c1 = new TCanvas("c1");

  if(doBayesian && doMCMC){
    c1->Divide(3);
    c1->cd(1);
  }
  else if (doBayesian || doMCMC){
    c1->Divide(2);
    c1->cd(1);
  }

  LikelihoodIntervalPlot* lrplot = new LikelihoodIntervalPlot(plInt);
  lrplot->Draw();

  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    c1->cd(2);
    // the plot takes a long time and print lots of error
    // using a scan it is better
    bc.SetScanOfPosterior(20);
    RooPlot* bplot = bc.GetPosteriorPlot();
    bplot->Draw();
  }

  if(doMCMC){
    if(doBayesian && wspace->set("poi")->getSize() == 1)
      c1->cd(3);
    else
      c1->cd(2);
    MCMCIntervalPlot mcPlot(*mcInt);
    mcPlot.Draw();
  }

  ////////////////////////////////////
  // querry intervals
  cout << "Profile Likelihood interval on s = [" <<
    plInt->LowerLimit( *wspace->var("s") ) << ", " <<
    plInt->UpperLimit( *wspace->var("s") ) << "]" << endl;
  //Profile Likelihood interval on s = [12.1902, 88.6871]

  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    cout << "Bayesian interval on s = [" <<
      bInt->LowerLimit( ) << ", " <<
      bInt->UpperLimit( ) << "]" << endl;
  }

  if(doFeldmanCousins){
    cout << "Feldman Cousins interval on s = [" <<
      fcInt->LowerLimit( *wspace->var("s") ) << ", " <<
      fcInt->UpperLimit( *wspace->var("s") ) << "]" << endl;
    //Feldman Cousins interval on s = [18.75 +/- 2.45, 83.75 +/- 2.45]
  }

  if(doMCMC){
    cout << "MCMC interval on s = [" <<
      mcInt->LowerLimit(*wspace->var("s") ) << ", " <<
      mcInt->UpperLimit(*wspace->var("s") ) << "]" << endl;
    //MCMC interval on s = [15.7628, 84.7266]

  }

  t.Print();

}
<commit_msg>Doxygen<commit_after>/// \file
/// \ingroup tutorial_roostats
///  This example is a generalization of the on/off problem.
///
///  This example is a generalization of the on/off problem.
/// It's a common setup for SUSY searches.  Imagine that one has two
/// variables "x" and "y" (eg. missing ET and SumET), see figure.
/// The signal region has high values of both of these variables (top right).
/// One can see low values of "x" or "y" acting as side-bands.  If we
/// just used "y" as a sideband, we would have the on/off problem.
///  - In the signal region we observe non events and expect s+b events.
///  - In the region with low values of "y" (bottom right)
///    we observe noff events and expect tau*b events.
/// Note the significance of tau.  In the background only case:
///
/// ~~~ {.cpp}
///    tau ~ <expectation off> / <expectation on>
/// ~~~
///
/// If tau is known, this model is sufficient, but often tau is not known exactly.
/// So one can use low values of "x" as an additional constraint for tau.
/// Note that this technique critically depends on the notion that the
/// joint distribution for "x" and "y" can be factorized.
/// Generally, these regions have many events, so it the ratio can be
/// measured very precisely there.  So we extend the model to describe the
/// left two boxes... denoted with "bar".
///   - In the upper left we observe nonbar events and expect bbar events
///   - In the bottom left we observe noffbar events and expect tau bbar events
/// Note again we have:
///
/// ~~~ {.cpp}
///    tau ~ <expecation off bar> / <expectation on bar>
/// ~~~
///
/// One can further expand the model to account for the systematic associated
/// to assuming the distribution of "x" and "y" factorizes (eg. that
/// tau is the same for off/on and offbar/onbar). This can be done in several
/// ways, but here we introduce an additional parameter rho, which so that
/// one set of models will use tau and the other tau*rho. The choice is arbitrary,
/// but it has consequences on the numerical stability of the algorithms.
/// The "bar" measurements typically have more events (& smaller relative errors).
/// If we choose
///
/// ~~~ {.cpp}
/// <expectation noffbar> = tau * rho * <expectation noonbar>
/// ~~~
///
/// the product tau*rho will be known very precisely (~1/sqrt(bbar)) and the contour
/// in those parameters will be narrow and have a non-trivial tau~1/rho shape.
/// However, if we choose to put rho on the non/noff measurements (where the
/// product will have an error `~1/sqrt(b))`, the contours will be more ameanable
/// to numerical techniques.  Thus, here we choose to define:
///
/// ~~~ {.cpp}
///    tau := <expecation off bar> / (<expectation on bar>)
///    rho := <expecation off> / (<expectation on> * tau)
///
/// ^ y
/// |
/// |---------------------------+
/// |               |           |
/// |     nonbar    |    non    |
/// |      bbar     |    s+b    |
/// |               |           |
/// |---------------+-----------|
/// |               |           |
/// |    noffbar    |    noff   |
/// |    tau bbar   | tau b rho |
/// |               |           |
/// +-----------------------------> x
/// ~~~
///
/// Left in this way, the problem is under-constrained.  However, one may
/// have some auxiliary measurement (usually based on Monte Carlo) to
/// constrain rho.  Let us call this auxiliary measurement that gives
/// the nominal value of rho "rhonom".  Thus, there is a 'constraint' term in
/// the full model: P(rhonom | rho).  In this case, we consider a Gaussian
/// constraint with standard deviation sigma.
///
/// In the example, the initial values of the parameters are:
///
/// ~~~ {.cpp}
///   - s    = 40
///   - b    = 100
///   - tau  = 5
///   - bbar = 1000
///   - rho  = 1
///   (sigma for rho = 20%)
/// ~~~
///
/// and in the toy dataset:
///
/// ~~~ {.cpp}
///    - non = 139
///    - noff = 528
///    - nonbar = 993
///    - noffbar = 4906
///    - rhonom = 1.27824
/// ~~~
///
/// Note, the covariance matrix of the parameters has large off-diagonal terms.
/// Clearly s,b are anti-correlated.  Similarly, since noffbar >> nonbar, one would
/// expect bbar,tau to be anti-correlated.
///
/// This can be seen below.
///
/// ~~~ {.cpp}
///             GLOBAL      b    bbar   rho      s     tau
///         b  0.96820   1.000  0.191 -0.942 -0.762 -0.209
///      bbar  0.91191   0.191  1.000  0.000 -0.146 -0.912
///       rho  0.96348  -0.942  0.000  1.000  0.718 -0.000
///         s  0.76250  -0.762 -0.146  0.718  1.000  0.160
///       tau  0.92084  -0.209 -0.912 -0.000  0.160  1.000
/// ~~~
///
/// Similarly, since tau*rho appears as a product, we expect rho,tau
/// to be anti-correlated. When the error on rho is significantly
/// larger than 1/sqrt(bbar), tau is essentially known and the
/// correlation is minimal (tau mainly cares about bbar, and rho about b,s).
/// In the alternate parametrization (bbar* tau * rho) the correlation coefficient
/// for rho,tau is large (and negative).
///
/// The code below uses best-practices for RooFit & RooStats as of June 2010.
///
/// It proceeds as follows:
///  - create a workspace to hold the model
///  - use workspace factory to quickly create the terms of the model
///  - use workspace factory to define total model (a prod pdf)
///  - create a RooStats ModelConfig to specify observables, parameters of interest
///  - add to the ModelConfig a prior on the parameters for Bayesian techniques
///    note, the pdf it is factorized for parameters of interest & nuisance params
///  - visualize the model
///  - write the workspace to a file
///  - use several of RooStats IntervalCalculators & compare results
///
/// \macro_image
/// \macro_output
/// \macro_code
///
/// \authors  authors: Kyle Cranmer, Tanja Rommerskirchen

#include "TStopwatch.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "RooPlot.h"
#include "RooAbsPdf.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooGlobalFunc.h"
#include "RooFitResult.h"
#include "RooRandom.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/BayesianCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/SimpleInterval.h"
#include "RooStats/FeldmanCousins.h"
#include "RooStats/PointSetInterval.h"

using namespace RooFit;
using namespace RooStats;

void FourBinInstructional(bool doBayesian=false, bool doFeldmanCousins=false, bool doMCMC=false){

  // let's time this challenging example
  TStopwatch t;
  t.Start();

  // set RooFit random seed for reproducible results
  RooRandom::randomGenerator()->SetSeed(4357);

  // make model
  RooWorkspace* wspace = new RooWorkspace("wspace");
  wspace->factory("Poisson::on(non[0,1000], sum::splusb(s[40,0,100],b[100,0,300]))");
  wspace->factory("Poisson::off(noff[0,5000], prod::taub(b,tau[5,3,7],rho[1,0,2]))");
  wspace->factory("Poisson::onbar(nonbar[0,10000], bbar[1000,500,2000])");
  wspace->factory("Poisson::offbar(noffbar[0,1000000], prod::lambdaoffbar(bbar, tau))");
  wspace->factory("Gaussian::mcCons(rhonom[1.,0,2], rho, sigma[.2])");
  wspace->factory("PROD::model(on,off,onbar,offbar,mcCons)");
  wspace->defineSet("obs","non,noff,nonbar,noffbar,rhonom");

  wspace->factory("Uniform::prior_poi({s})");
  wspace->factory("Uniform::prior_nuis({b,bbar,tau, rho})");
  wspace->factory("PROD::prior(prior_poi,prior_nuis)");

  ///////////////////////////////////////////
  // Control some interesting variations
  // define parameers of interest
  // for 1-d plots
  wspace->defineSet("poi","s");
  wspace->defineSet("nuis","b,tau,rho,bbar");
  // for 2-d plots to inspect correlations:
  //  wspace->defineSet("poi","s,rho");

  // test simpler cases where parameters are known.
  //  wspace->var("tau")->setConstant();
  //  wspace->var("rho")->setConstant();
  //  wspace->var("b")->setConstant();
  //  wspace->var("bbar")->setConstant();

  // inspect workspace
  //  wspace->Print();

  ////////////////////////////////////////////////////////////
  // Generate toy data
  // generate toy data assuming current value of the parameters
  // import into workspace.
  // add Verbose() to see how it's being generated
  RooDataSet* data =   wspace->pdf("model")->generate(*wspace->set("obs"),1);
  //  data->Print("v");
  wspace->import(*data);

  /////////////////////////////////////////////////////
  // Now the statistical tests
  // model config
  ModelConfig* modelConfig = new ModelConfig("FourBins");
  modelConfig->SetWorkspace(*wspace);
  modelConfig->SetPdf(*wspace->pdf("model"));
  modelConfig->SetPriorPdf(*wspace->pdf("prior"));
  modelConfig->SetParametersOfInterest(*wspace->set("poi"));
  modelConfig->SetNuisanceParameters(*wspace->set("nuis"));
  wspace->import(*modelConfig);
  wspace->writeToFile("FourBin.root");

  //////////////////////////////////////////////////
  // If you want to see the covariance matrix uncomment
  //  wspace->pdf("model")->fitTo(*data);

  // use ProfileLikelihood
  ProfileLikelihoodCalculator plc(*data, *modelConfig);
  plc.SetConfidenceLevel(0.95);
  LikelihoodInterval* plInt = plc.GetInterval();
  RooFit::MsgLevel msglevel = RooMsgService::instance().globalKillBelow();
  RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);
  plInt->LowerLimit( *wspace->var("s") ); // get ugly print out of the way. Fix.
  RooMsgService::instance().setGlobalKillBelow(msglevel);

  // use FeldmaCousins (takes ~20 min)
  FeldmanCousins fc(*data, *modelConfig);
  fc.SetConfidenceLevel(0.95);
  //number counting: dataset always has 1 entry with N events observed
  fc.FluctuateNumDataEntries(false);
  fc.UseAdaptiveSampling(true);
  fc.SetNBins(40);
  PointSetInterval* fcInt = NULL;
  if(doFeldmanCousins){ // takes 7 minutes
    fcInt = (PointSetInterval*) fc.GetInterval(); // fix cast
  }

  // use BayesianCalculator (only 1-d parameter of interest, slow for this problem)
  BayesianCalculator bc(*data, *modelConfig);
  bc.SetConfidenceLevel(0.95);
  SimpleInterval* bInt = NULL;
  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    bInt = bc.GetInterval();
  } else{
    cout << "Bayesian Calc. only supports on parameter of interest" << endl;
  }

  // use MCMCCalculator  (takes about 1 min)
  // Want an efficient proposal function, so derive it from covariance
  // matrix of fit
  RooFitResult* fit = wspace->pdf("model")->fitTo(*data,Save());
  ProposalHelper ph;
  ph.SetVariables((RooArgSet&)fit->floatParsFinal());
  ph.SetCovMatrix(fit->covarianceMatrix());
  ph.SetUpdateProposalParameters(kTRUE); // auto-create mean vars and add mappings
  ph.SetCacheSize(100);
  ProposalFunction* pf = ph.GetProposalFunction();

  MCMCCalculator mc(*data, *modelConfig);
  mc.SetConfidenceLevel(0.95);
  mc.SetProposalFunction(*pf);
  mc.SetNumBurnInSteps(500); // first N steps to be ignored as burn-in
  mc.SetNumIters(50000);
  mc.SetLeftSideTailFraction(0.5); // make a central interval
  MCMCInterval* mcInt = NULL;
  if(doMCMC)
    mcInt = mc.GetInterval();

  //////////////////////////////////////
  // Make some  plots
  TCanvas* c1 = (TCanvas*) gROOT->Get("c1");
  if(!c1)
    c1 = new TCanvas("c1");

  if(doBayesian && doMCMC){
    c1->Divide(3);
    c1->cd(1);
  }
  else if (doBayesian || doMCMC){
    c1->Divide(2);
    c1->cd(1);
  }

  LikelihoodIntervalPlot* lrplot = new LikelihoodIntervalPlot(plInt);
  lrplot->Draw();

  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    c1->cd(2);
    // the plot takes a long time and print lots of error
    // using a scan it is better
    bc.SetScanOfPosterior(20);
    RooPlot* bplot = bc.GetPosteriorPlot();
    bplot->Draw();
  }

  if(doMCMC){
    if(doBayesian && wspace->set("poi")->getSize() == 1)
      c1->cd(3);
    else
      c1->cd(2);
    MCMCIntervalPlot mcPlot(*mcInt);
    mcPlot.Draw();
  }

  ////////////////////////////////////
  // querry intervals
  cout << "Profile Likelihood interval on s = [" <<
    plInt->LowerLimit( *wspace->var("s") ) << ", " <<
    plInt->UpperLimit( *wspace->var("s") ) << "]" << endl;
  //Profile Likelihood interval on s = [12.1902, 88.6871]

  if(doBayesian && wspace->set("poi")->getSize() == 1)   {
    cout << "Bayesian interval on s = [" <<
      bInt->LowerLimit( ) << ", " <<
      bInt->UpperLimit( ) << "]" << endl;
  }

  if(doFeldmanCousins){
    cout << "Feldman Cousins interval on s = [" <<
      fcInt->LowerLimit( *wspace->var("s") ) << ", " <<
      fcInt->UpperLimit( *wspace->var("s") ) << "]" << endl;
    //Feldman Cousins interval on s = [18.75 +/- 2.45, 83.75 +/- 2.45]
  }

  if(doMCMC){
    cout << "MCMC interval on s = [" <<
      mcInt->LowerLimit(*wspace->var("s") ) << ", " <<
      mcInt->UpperLimit(*wspace->var("s") ) << "]" << endl;
    //MCMC interval on s = [15.7628, 84.7266]

  }

  t.Print();
}
<|endoftext|>


<commit_before>//===--- Immediate.cpp - the swift immediate mode -------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// This is the implementation of the swift interpreter, which takes a
// source file and JITs it.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "swift-immediate"
#include "swift/Immediate/Immediate.h"
#include "ImmediateImpl.h"

#include "swift/Subsystems.h"
#include "swift/AST/ASTContext.h"
#include "swift/AST/DiagnosticsFrontend.h"
#include "swift/AST/IRGenOptions.h"
#include "swift/AST/Module.h"
#include "swift/Basic/LLVM.h"
#include "swift/Basic/LLVMContext.h"
#include "swift/Frontend/Frontend.h"
#include "swift/SILOptimizer/PassManager/Passes.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/Config/config.h"
#include "llvm/ExecutionEngine/MCJIT.h"
#include "llvm/IR/DiagnosticPrinter.h"
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Linker/Linker.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/Support/Path.h"

#if defined(_MSC_VER)
#include "Windows.h"
#else
#include <dlfcn.h>
#endif

using namespace swift;
using namespace swift::immediate;

static void *loadRuntimeLib(StringRef runtimeLibPathWithName) {
#if defined(_MSC_VER)
  return LoadLibrary(runtimeLibPathWithName.str().c_str());
#else
  return dlopen(runtimeLibPathWithName.str().c_str(), RTLD_LAZY | RTLD_GLOBAL);
#endif
}

static void *loadRuntimeLib(StringRef sharedLibName, StringRef runtimeLibPath) {
  // FIXME: Need error-checking.
  llvm::SmallString<128> Path = runtimeLibPath;
  llvm::sys::path::append(Path, sharedLibName);
  return loadRuntimeLib(Path);
}

void *swift::immediate::loadSwiftRuntime(StringRef runtimeLibPath) {
  return loadRuntimeLib("libswiftCore" LTDL_SHLIB_EXT, runtimeLibPath);
}

static bool tryLoadLibrary(LinkLibrary linkLib,
                           SearchPathOptions searchPathOpts) {
  llvm::SmallString<128> path = linkLib.getName();

  // If we have an absolute or relative path, just try to load it now.
  if (llvm::sys::path::has_parent_path(path.str())) {
    return loadRuntimeLib(path);
  }

  bool success = false;
  switch (linkLib.getKind()) {
  case LibraryKind::Library: {
    llvm::SmallString<32> stem;
    if (llvm::sys::path::has_extension(path.str())) {
      stem = std::move(path);
    } else {
      // FIXME: Try the appropriate extension for the current platform?
      stem = "lib";
      stem += path;
      stem += LTDL_SHLIB_EXT;
    }

    // Try user-provided library search paths first.
    for (auto &libDir : searchPathOpts.LibrarySearchPaths) {
      path = libDir;
      llvm::sys::path::append(path, stem.str());
      success = loadRuntimeLib(path);
      if (success)
        break;
    }

    // Let loadRuntimeLib determine the best search paths.
    if (!success)
      success = loadRuntimeLib(stem);

    // If that fails, try our runtime library path.
    if (!success)
      success = loadRuntimeLib(stem, searchPathOpts.RuntimeLibraryPath);
    break;
  }
  case LibraryKind::Framework: {
    // If we have a framework, mangle the name to point to the framework
    // binary.
    llvm::SmallString<64> frameworkPart{std::move(path)};
    frameworkPart += ".framework";
    llvm::sys::path::append(frameworkPart, linkLib.getName());

    // Try user-provided framework search paths first; frameworks contain
    // binaries as well as modules.
    for (auto &frameworkDir : searchPathOpts.FrameworkSearchPaths) {
      path = frameworkDir;
      llvm::sys::path::append(path, frameworkPart.str());
      success = loadRuntimeLib(path);
      if (success)
        break;
    }

    // If that fails, let loadRuntimeLib search for system frameworks.
    if (!success)
      success = loadRuntimeLib(frameworkPart);
    break;
  }
  }

  return success;
}

bool swift::immediate::tryLoadLibraries(ArrayRef<LinkLibrary> LinkLibraries,
                                        SearchPathOptions SearchPathOpts,
                                        DiagnosticEngine &Diags) {
  SmallVector<bool, 4> LoadedLibraries;
  LoadedLibraries.append(LinkLibraries.size(), false);

  // Libraries are not sorted in the topological order of dependencies, and we
  // don't know the dependencies in advance.  Try to load all libraries until
  // we stop making progress.
  bool HadProgress;
  do {
    HadProgress = false;
    for (unsigned i = 0; i != LinkLibraries.size(); ++i) {
      if (!LoadedLibraries[i] &&
          tryLoadLibrary(LinkLibraries[i], SearchPathOpts)) {
        LoadedLibraries[i] = true;
        HadProgress = true;
      }
    }
  } while (HadProgress);

  return std::all_of(LoadedLibraries.begin(), LoadedLibraries.end(),
                     [](bool Value) { return Value; });
}

static void linkerDiagnosticHandlerNoCtx(const llvm::DiagnosticInfo &DI) {
  if (DI.getSeverity() != llvm::DS_Error)
    return;

  std::string MsgStorage;
  {
    llvm::raw_string_ostream Stream(MsgStorage);
    llvm::DiagnosticPrinterRawOStream DP(Stream);
    DI.print(DP);
  }
  llvm::errs() << "Error linking swift modules\n";
  llvm::errs() << MsgStorage << "\n";
}

static void linkerDiagnosticHandler(const llvm::DiagnosticInfo &DI,
                                    void *Context) {
  // This assert self documents our precondition that Context is always
  // nullptr. It seems that parts of LLVM are using the flexibility of having a
  // context. We don't really care about this.
  assert(Context == nullptr && "We assume Context is always a nullptr");

  return linkerDiagnosticHandlerNoCtx(DI);
}

bool swift::immediate::linkLLVMModules(llvm::Module *Module,
                                       std::unique_ptr<llvm::Module> SubModule
                            // TODO: reactivate the linker mode if it is
                            // supported in llvm again. Otherwise remove the
                            // commented code completely.
                            /*, llvm::Linker::LinkerMode LinkerMode */)
{
  llvm::LLVMContext &Ctx = SubModule->getContext();
  auto OldHandler = Ctx.getDiagnosticHandler();
  void *OldDiagnosticContext = Ctx.getDiagnosticContext();
  Ctx.setDiagnosticHandler(linkerDiagnosticHandler, nullptr);
  bool Failed = llvm::Linker::linkModules(*Module, std::move(SubModule));
  Ctx.setDiagnosticHandler(OldHandler, OldDiagnosticContext);

  return !Failed;
}

bool swift::immediate::IRGenImportedModules(
    CompilerInstance &CI,
    llvm::Module &Module,
    llvm::SmallPtrSet<swift::Module *, 8> &ImportedModules,
    SmallVectorImpl<llvm::Function*> &InitFns,
    IRGenOptions &IRGenOpts,
    const SILOptions &SILOpts) {
  swift::Module *M = CI.getMainModule();

  // Perform autolinking.
  SmallVector<LinkLibrary, 4> AllLinkLibraries(IRGenOpts.LinkLibraries);
  auto addLinkLibrary = [&](LinkLibrary linkLib) {
    AllLinkLibraries.push_back(linkLib);
  };

  M->forAllVisibleModules({}, /*includePrivateTopLevel=*/true,
                          [&](Module::ImportedModule import) {
    import.second->collectLinkLibraries(addLinkLibrary);
  });

  // Hack to handle thunks eagerly synthesized by the Clang importer.
  swift::Module *prev = nullptr;
  for (auto external : CI.getASTContext().ExternalDefinitions) {
    swift::Module *next = external->getModuleContext();
    if (next == prev)
      continue;
    next->collectLinkLibraries(addLinkLibrary);
    prev = next;
  }

  tryLoadLibraries(AllLinkLibraries, CI.getASTContext().SearchPathOpts,
                   CI.getDiags());

  ImportedModules.insert(M);
  if (!CI.hasSourceImport())
    return false;

  // IRGen the modules this module depends on. This is only really necessary
  // for imported source, but that's a very convenient thing to do in -i mode.
  // FIXME: Crawling all loaded modules is a hack.
  // FIXME: And re-doing SILGen, SIL-linking, SIL diagnostics, and IRGen is
  // expensive, because it's not properly being limited to new things right now.
  bool hadError = false;
  for (auto &entry : CI.getASTContext().LoadedModules) {
    swift::Module *import = entry.second;
    if (!ImportedModules.insert(import).second)
      continue;

    std::unique_ptr<SILModule> SILMod = performSILGeneration(import,
                                                             CI.getSILOptions());
    performSILLinking(SILMod.get());
    if (runSILDiagnosticPasses(*SILMod)) {
      hadError = true;
      break;
    }

    // FIXME: We shouldn't need to use the global context here, but
    // something is persisting across calls to performIRGeneration.
    auto SubModule = performIRGeneration(IRGenOpts, import,
                                         std::move(SILMod),
                                         import->getName().str(),
                                         getGlobalLLVMContext());

    if (CI.getASTContext().hadError()) {
      hadError = true;
      break;
    }

    if (!linkLLVMModules(&Module, std::move(SubModule)
                         // TODO: reactivate the linker mode if it is
                         // supported in llvm again. Otherwise remove the
                         // commented code completely.
                         /*, llvm::Linker::DestroySource */)) {
      hadError = true;
      break;
    }

    // FIXME: This is an ugly hack; need to figure out how this should
    // actually work.
    SmallVector<char, 20> NameBuf;
    StringRef InitFnName = (import->getName().str() + ".init").toStringRef(NameBuf);
    llvm::Function *InitFn = Module.getFunction(InitFnName);
    if (InitFn)
      InitFns.push_back(InitFn);
  }

  return hadError;
}

int swift::RunImmediately(CompilerInstance &CI, const ProcessCmdLine &CmdLine,
                          IRGenOptions &IRGenOpts, const SILOptions &SILOpts) {
  ASTContext &Context = CI.getASTContext();

  // IRGen the main module.
  auto *swiftModule = CI.getMainModule();
  // FIXME: We shouldn't need to use the global context here, but
  // something is persisting across calls to performIRGeneration.
  auto ModuleOwner = performIRGeneration(IRGenOpts, swiftModule,
                                         CI.takeSILModule(),
                                         swiftModule->getName().str(),
                                         getGlobalLLVMContext());
  auto *Module = ModuleOwner.get();

  if (Context.hadError())
    return -1;

  // Load libSwiftCore to setup process arguments.
  //
  // This must be done here, before any library loading has been done, to avoid
  // racing with the static initializers in user code.
  auto stdlib = loadSwiftRuntime(Context.SearchPathOpts.RuntimeLibraryPath);
  if (!stdlib) {
    CI.getDiags().diagnose(SourceLoc(),
                           diag::error_immediate_mode_missing_stdlib);
    return -1;
  }

  // Setup interpreted process arguments.
  using ArgOverride = void (*)(const char **, int);
  auto emplaceProcessArgs
          = (ArgOverride)dlsym(stdlib, "_swift_stdlib_overrideUnsafeArgvArgc");
  if (dlerror())
    return -1;

  SmallVector<const char *, 32> argBuf;
  for (size_t i = 0; i < CmdLine.size(); ++i) {
    argBuf.push_back(CmdLine[i].c_str());
  }
  argBuf.push_back(nullptr);

  (*emplaceProcessArgs)(argBuf.data(), CmdLine.size());

  SmallVector<llvm::Function*, 8> InitFns;
  llvm::SmallPtrSet<swift::Module *, 8> ImportedModules;
  if (IRGenImportedModules(CI, *Module, ImportedModules, InitFns,
                           IRGenOpts, SILOpts))
    return -1;

  llvm::PassManagerBuilder PMBuilder;
  PMBuilder.OptLevel = 2;
  PMBuilder.Inliner = llvm::createFunctionInliningPass(200);

  // Build the ExecutionEngine.
  llvm::EngineBuilder builder(std::move(ModuleOwner));
  std::string ErrorMsg;
  llvm::TargetOptions TargetOpt;
  std::string CPU;
  std::vector<std::string> Features;
  std::tie(TargetOpt, CPU, Features)
    = getIRTargetOptions(IRGenOpts, swiftModule->getASTContext());
  builder.setRelocationModel(llvm::Reloc::PIC_);
  builder.setTargetOptions(TargetOpt);
  builder.setMCPU(CPU);
  builder.setMAttrs(Features);
  builder.setErrorStr(&ErrorMsg);
  builder.setEngineKind(llvm::EngineKind::JIT);
  llvm::ExecutionEngine *EE = builder.create();
  if (!EE) {
    llvm::errs() << "Error loading JIT: " << ErrorMsg;
    return -1;
  }

  DEBUG(llvm::dbgs() << "Module to be executed:\n";
        Module->dump());

  EE->finalizeObject();

  // Run the generated program.
  for (auto InitFn : InitFns) {
    DEBUG(llvm::dbgs() << "Running initialization function "
            << InitFn->getName() << '\n');
    EE->runFunctionAsMain(InitFn, CmdLine, nullptr);
  }

  DEBUG(llvm::dbgs() << "Running static constructors\n");
  EE->runStaticConstructorsDestructors(false);
  DEBUG(llvm::dbgs() << "Running main\n");
  llvm::Function *EntryFn = Module->getFunction("main");
  return EE->runFunctionAsMain(EntryFn, CmdLine, nullptr);
}
<commit_msg>Port swift/Immediate to Windows (#5957)<commit_after>//===--- Immediate.cpp - the swift immediate mode -------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// This is the implementation of the swift interpreter, which takes a
// source file and JITs it.
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "swift-immediate"
#include "swift/Immediate/Immediate.h"
#include "ImmediateImpl.h"

#include "swift/Subsystems.h"
#include "swift/AST/ASTContext.h"
#include "swift/AST/DiagnosticsFrontend.h"
#include "swift/AST/IRGenOptions.h"
#include "swift/AST/Module.h"
#include "swift/Basic/LLVM.h"
#include "swift/Basic/LLVMContext.h"
#include "swift/Frontend/Frontend.h"
#include "swift/SILOptimizer/PassManager/Passes.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/Config/config.h"
#include "llvm/ExecutionEngine/MCJIT.h"
#include "llvm/IR/DiagnosticPrinter.h"
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Linker/Linker.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/Support/Path.h"

#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#else
#include <dlfcn.h>
#endif

using namespace swift;
using namespace swift::immediate;

static void *loadRuntimeLib(StringRef runtimeLibPathWithName) {
#if defined(_WIN32)
  return LoadLibraryA(runtimeLibPathWithName.str().c_str());
#else
  return dlopen(runtimeLibPathWithName.str().c_str(), RTLD_LAZY | RTLD_GLOBAL);
#endif
}

static void *loadRuntimeLib(StringRef sharedLibName, StringRef runtimeLibPath) {
  // FIXME: Need error-checking.
  llvm::SmallString<128> Path = runtimeLibPath;
  llvm::sys::path::append(Path, sharedLibName);
  return loadRuntimeLib(Path);
}

void *swift::immediate::loadSwiftRuntime(StringRef runtimeLibPath) {
  return loadRuntimeLib("libswiftCore" LTDL_SHLIB_EXT, runtimeLibPath);
}

static bool tryLoadLibrary(LinkLibrary linkLib,
                           SearchPathOptions searchPathOpts) {
  llvm::SmallString<128> path = linkLib.getName();

  // If we have an absolute or relative path, just try to load it now.
  if (llvm::sys::path::has_parent_path(path.str())) {
    return loadRuntimeLib(path);
  }

  bool success = false;
  switch (linkLib.getKind()) {
  case LibraryKind::Library: {
    llvm::SmallString<32> stem;
    if (llvm::sys::path::has_extension(path.str())) {
      stem = std::move(path);
    } else {
      // FIXME: Try the appropriate extension for the current platform?
      stem = "lib";
      stem += path;
      stem += LTDL_SHLIB_EXT;
    }

    // Try user-provided library search paths first.
    for (auto &libDir : searchPathOpts.LibrarySearchPaths) {
      path = libDir;
      llvm::sys::path::append(path, stem.str());
      success = loadRuntimeLib(path);
      if (success)
        break;
    }

    // Let loadRuntimeLib determine the best search paths.
    if (!success)
      success = loadRuntimeLib(stem);

    // If that fails, try our runtime library path.
    if (!success)
      success = loadRuntimeLib(stem, searchPathOpts.RuntimeLibraryPath);
    break;
  }
  case LibraryKind::Framework: {
    // If we have a framework, mangle the name to point to the framework
    // binary.
    llvm::SmallString<64> frameworkPart{std::move(path)};
    frameworkPart += ".framework";
    llvm::sys::path::append(frameworkPart, linkLib.getName());

    // Try user-provided framework search paths first; frameworks contain
    // binaries as well as modules.
    for (auto &frameworkDir : searchPathOpts.FrameworkSearchPaths) {
      path = frameworkDir;
      llvm::sys::path::append(path, frameworkPart.str());
      success = loadRuntimeLib(path);
      if (success)
        break;
    }

    // If that fails, let loadRuntimeLib search for system frameworks.
    if (!success)
      success = loadRuntimeLib(frameworkPart);
    break;
  }
  }

  return success;
}

bool swift::immediate::tryLoadLibraries(ArrayRef<LinkLibrary> LinkLibraries,
                                        SearchPathOptions SearchPathOpts,
                                        DiagnosticEngine &Diags) {
  SmallVector<bool, 4> LoadedLibraries;
  LoadedLibraries.append(LinkLibraries.size(), false);

  // Libraries are not sorted in the topological order of dependencies, and we
  // don't know the dependencies in advance.  Try to load all libraries until
  // we stop making progress.
  bool HadProgress;
  do {
    HadProgress = false;
    for (unsigned i = 0; i != LinkLibraries.size(); ++i) {
      if (!LoadedLibraries[i] &&
          tryLoadLibrary(LinkLibraries[i], SearchPathOpts)) {
        LoadedLibraries[i] = true;
        HadProgress = true;
      }
    }
  } while (HadProgress);

  return std::all_of(LoadedLibraries.begin(), LoadedLibraries.end(),
                     [](bool Value) { return Value; });
}

static void linkerDiagnosticHandlerNoCtx(const llvm::DiagnosticInfo &DI) {
  if (DI.getSeverity() != llvm::DS_Error)
    return;

  std::string MsgStorage;
  {
    llvm::raw_string_ostream Stream(MsgStorage);
    llvm::DiagnosticPrinterRawOStream DP(Stream);
    DI.print(DP);
  }
  llvm::errs() << "Error linking swift modules\n";
  llvm::errs() << MsgStorage << "\n";
}

static void linkerDiagnosticHandler(const llvm::DiagnosticInfo &DI,
                                    void *Context) {
  // This assert self documents our precondition that Context is always
  // nullptr. It seems that parts of LLVM are using the flexibility of having a
  // context. We don't really care about this.
  assert(Context == nullptr && "We assume Context is always a nullptr");

  return linkerDiagnosticHandlerNoCtx(DI);
}

bool swift::immediate::linkLLVMModules(llvm::Module *Module,
                                       std::unique_ptr<llvm::Module> SubModule
                            // TODO: reactivate the linker mode if it is
                            // supported in llvm again. Otherwise remove the
                            // commented code completely.
                            /*, llvm::Linker::LinkerMode LinkerMode */)
{
  llvm::LLVMContext &Ctx = SubModule->getContext();
  auto OldHandler = Ctx.getDiagnosticHandler();
  void *OldDiagnosticContext = Ctx.getDiagnosticContext();
  Ctx.setDiagnosticHandler(linkerDiagnosticHandler, nullptr);
  bool Failed = llvm::Linker::linkModules(*Module, std::move(SubModule));
  Ctx.setDiagnosticHandler(OldHandler, OldDiagnosticContext);

  return !Failed;
}

bool swift::immediate::IRGenImportedModules(
    CompilerInstance &CI,
    llvm::Module &Module,
    llvm::SmallPtrSet<swift::Module *, 8> &ImportedModules,
    SmallVectorImpl<llvm::Function*> &InitFns,
    IRGenOptions &IRGenOpts,
    const SILOptions &SILOpts) {
  swift::Module *M = CI.getMainModule();

  // Perform autolinking.
  SmallVector<LinkLibrary, 4> AllLinkLibraries(IRGenOpts.LinkLibraries);
  auto addLinkLibrary = [&](LinkLibrary linkLib) {
    AllLinkLibraries.push_back(linkLib);
  };

  M->forAllVisibleModules({}, /*includePrivateTopLevel=*/true,
                          [&](Module::ImportedModule import) {
    import.second->collectLinkLibraries(addLinkLibrary);
  });

  // Hack to handle thunks eagerly synthesized by the Clang importer.
  swift::Module *prev = nullptr;
  for (auto external : CI.getASTContext().ExternalDefinitions) {
    swift::Module *next = external->getModuleContext();
    if (next == prev)
      continue;
    next->collectLinkLibraries(addLinkLibrary);
    prev = next;
  }

  tryLoadLibraries(AllLinkLibraries, CI.getASTContext().SearchPathOpts,
                   CI.getDiags());

  ImportedModules.insert(M);
  if (!CI.hasSourceImport())
    return false;

  // IRGen the modules this module depends on. This is only really necessary
  // for imported source, but that's a very convenient thing to do in -i mode.
  // FIXME: Crawling all loaded modules is a hack.
  // FIXME: And re-doing SILGen, SIL-linking, SIL diagnostics, and IRGen is
  // expensive, because it's not properly being limited to new things right now.
  bool hadError = false;
  for (auto &entry : CI.getASTContext().LoadedModules) {
    swift::Module *import = entry.second;
    if (!ImportedModules.insert(import).second)
      continue;

    std::unique_ptr<SILModule> SILMod = performSILGeneration(import,
                                                             CI.getSILOptions());
    performSILLinking(SILMod.get());
    if (runSILDiagnosticPasses(*SILMod)) {
      hadError = true;
      break;
    }

    // FIXME: We shouldn't need to use the global context here, but
    // something is persisting across calls to performIRGeneration.
    auto SubModule = performIRGeneration(IRGenOpts, import,
                                         std::move(SILMod),
                                         import->getName().str(),
                                         getGlobalLLVMContext());

    if (CI.getASTContext().hadError()) {
      hadError = true;
      break;
    }

    if (!linkLLVMModules(&Module, std::move(SubModule)
                         // TODO: reactivate the linker mode if it is
                         // supported in llvm again. Otherwise remove the
                         // commented code completely.
                         /*, llvm::Linker::DestroySource */)) {
      hadError = true;
      break;
    }

    // FIXME: This is an ugly hack; need to figure out how this should
    // actually work.
    SmallVector<char, 20> NameBuf;
    StringRef InitFnName = (import->getName().str() + ".init").toStringRef(NameBuf);
    llvm::Function *InitFn = Module.getFunction(InitFnName);
    if (InitFn)
      InitFns.push_back(InitFn);
  }

  return hadError;
}

int swift::RunImmediately(CompilerInstance &CI, const ProcessCmdLine &CmdLine,
                          IRGenOptions &IRGenOpts, const SILOptions &SILOpts) {
  ASTContext &Context = CI.getASTContext();

  // IRGen the main module.
  auto *swiftModule = CI.getMainModule();
  // FIXME: We shouldn't need to use the global context here, but
  // something is persisting across calls to performIRGeneration.
  auto ModuleOwner = performIRGeneration(IRGenOpts, swiftModule,
                                         CI.takeSILModule(),
                                         swiftModule->getName().str(),
                                         getGlobalLLVMContext());
  auto *Module = ModuleOwner.get();

  if (Context.hadError())
    return -1;

  // Load libSwiftCore to setup process arguments.
  //
  // This must be done here, before any library loading has been done, to avoid
  // racing with the static initializers in user code.
  auto stdlib = loadSwiftRuntime(Context.SearchPathOpts.RuntimeLibraryPath);
  if (!stdlib) {
    CI.getDiags().diagnose(SourceLoc(),
                           diag::error_immediate_mode_missing_stdlib);
    return -1;
  }

  // Setup interpreted process arguments.
  using ArgOverride = void (*)(const char **, int);
#if defined(_WIN32)
  auto module = static_cast<HMODULE>(stdlib);
  auto emplaceProcessArgs = static_cast<ArgOverride>(
    GetProcAddress(module, "_swift_stdlib_overrideUnsafeArgvArgc"));
  if (emplaceProcessArgs == nullptr)
    return -1;
#else
  auto emplaceProcessArgs
          = (ArgOverride)dlsym(stdlib, "_swift_stdlib_overrideUnsafeArgvArgc");
  if (dlerror())
    return -1;
#endif

  SmallVector<const char *, 32> argBuf;
  for (size_t i = 0; i < CmdLine.size(); ++i) {
    argBuf.push_back(CmdLine[i].c_str());
  }
  argBuf.push_back(nullptr);

  (*emplaceProcessArgs)(argBuf.data(), CmdLine.size());

  SmallVector<llvm::Function*, 8> InitFns;
  llvm::SmallPtrSet<swift::Module *, 8> ImportedModules;
  if (IRGenImportedModules(CI, *Module, ImportedModules, InitFns,
                           IRGenOpts, SILOpts))
    return -1;

  llvm::PassManagerBuilder PMBuilder;
  PMBuilder.OptLevel = 2;
  PMBuilder.Inliner = llvm::createFunctionInliningPass(200);

  // Build the ExecutionEngine.
  llvm::EngineBuilder builder(std::move(ModuleOwner));
  std::string ErrorMsg;
  llvm::TargetOptions TargetOpt;
  std::string CPU;
  std::vector<std::string> Features;
  std::tie(TargetOpt, CPU, Features)
    = getIRTargetOptions(IRGenOpts, swiftModule->getASTContext());
  builder.setRelocationModel(llvm::Reloc::PIC_);
  builder.setTargetOptions(TargetOpt);
  builder.setMCPU(CPU);
  builder.setMAttrs(Features);
  builder.setErrorStr(&ErrorMsg);
  builder.setEngineKind(llvm::EngineKind::JIT);
  llvm::ExecutionEngine *EE = builder.create();
  if (!EE) {
    llvm::errs() << "Error loading JIT: " << ErrorMsg;
    return -1;
  }

  DEBUG(llvm::dbgs() << "Module to be executed:\n";
        Module->dump());

  EE->finalizeObject();

  // Run the generated program.
  for (auto InitFn : InitFns) {
    DEBUG(llvm::dbgs() << "Running initialization function "
            << InitFn->getName() << '\n');
    EE->runFunctionAsMain(InitFn, CmdLine, nullptr);
  }

  DEBUG(llvm::dbgs() << "Running static constructors\n");
  EE->runStaticConstructorsDestructors(false);
  DEBUG(llvm::dbgs() << "Running main\n");
  llvm::Function *EntryFn = Module->getFunction("main");
  return EE->runFunctionAsMain(EntryFn, CmdLine, nullptr);
}
<|endoftext|>