revng-revng/set.cpp

/// \file set.cpp
/// \brief Simple Expression Tracker pass implementation
/// This file is composed by three main parts: the OperationsStack
/// implementation, the SET algorithm and the SET pass

// LLVM includes
#include "llvm/IR/Instruction.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"

// Local includes
#include "debug.h"
#include "revamb.h"
#include "ir-helpers.h"
#include "osra.h"
#include "jumptargetmanager.h"
#include "set.h"

using namespace llvm;
using std::make_pair;

/// \brief Stack to keep track of the operations generating a specific value
///
/// The OperationsStacks offers the following features:
///
/// * it doesn't insert more than once an item (to avoid loops)
/// * it can traverse the stack from top to bottom to produce a value and, if
///   required register it with the JumpTargetManager
/// * cut the stack to a certain height
/// * manage the lifetime of orphan instruction it contains
class OperationsStack {
public:
  OperationsStack(JumpTargetManager *JTM,
                  const DataLayout &DL) : JTM(JTM), DL(DL) { }

  ~OperationsStack() {
    reset(false);
  }

  void explore(Constant *NewOperand);
  uint64_t materialize(Constant *NewOperand);

  void reset(bool Reliable) {
    // Delete all the temporary instructions we created
    for (Instruction *I : Operations)
      if (I->getParent() == nullptr)
        delete I;

    Operations.clear();
    OperationsSet.clear();
    IsReliable = Reliable;
  }

  void registerPCs() const {
    for (auto Pair : PCs)
      JTM->getBlockAt(Pair.first, Pair.second);
  }

  void cut(unsigned Height) {
    assert(Height <= Operations.size());
    while (Height != Operations.size()) {
      Instruction *Op = Operations.back();
      auto It = OperationsSet.find(Op);
      if (It != OperationsSet.end())
        OperationsSet.erase(It);
      else if (isa<BinaryOperator>(Op)) {
        // It's not in OperationsSet, it might a binary instruction where we
        // forced one operand to be constant, or an instruction generated from a
        // constant unary expression
        unsigned FreeOpIndex = isa<Constant>(Op->getOperand(0)) ? 1 : 0;
        auto *FreeOp = cast<Instruction>(Op->getOperand(FreeOpIndex));
        auto It = OperationsSet.find(FreeOp);
        assert(It != OperationsSet.end());
        OperationsSet.erase(It);
      }

      // We have the ownership of instruction without parent
      if (Op->getParent() == nullptr)
        delete Op;

      Operations.pop_back();
    }
  }

  bool insertIfNew(Instruction *I) {
    return insertIfNew(I, I);
  }

  bool insertIfNew(Instruction *I, Instruction *Ref) {
    if (OperationsSet.find(Ref) == OperationsSet.end()) {
      Operations.push_back(I);
      OperationsSet.insert(Ref);
      return true;
    }

    // If the given instruction doesn't have a parent we take ownership of it
    if (I->getParent() == nullptr)
      delete I;

    return false;
  }

  void insert(Instruction *I) {
    Operations.push_back(I);
  }

  unsigned height() const { return Operations.size(); }
  bool empty() const { return height() == 0; }

private:
  JumpTargetManager *JTM;
  const DataLayout &DL;

  std::vector<Instruction *> Operations;
  std::set<Instruction *> OperationsSet;
  std::set<std::pair<uint64_t, bool>> PCs;

  bool IsReliable;
};

uint64_t OperationsStack::materialize(Constant *NewOperand) {
  for (Instruction *I : make_range(Operations.rbegin(), Operations.rend())) {
    if (auto *Load = dyn_cast<LoadInst>(I)) {
      // OK, we've got a load, let's see if the load address is
      // constant
      assert(NewOperand != nullptr && !isa<UndefValue>(NewOperand));

      if (Load->getType()->isIntegerTy()) {
        unsigned Size = Load->getType()->getPrimitiveSizeInBits() / 8;
        assert(Size != 0);
        NewOperand = JTM->readConstantInt(NewOperand, Size);
      } else if (Load->getType()->isPointerTy()) {
        NewOperand = JTM->readConstantPointer(NewOperand, Load->getType());
      } else {
        assert(false);
      }

      if (NewOperand == nullptr)
        break;
    } else if (auto *Call = dyn_cast<CallInst>(I)) {
      Function *Callee = Call->getCalledFunction();
      assert(Callee != nullptr && Callee->getIntrinsicID() == Intrinsic::bswap);
      (void) Callee;

      uint64_t Value = NewOperand->getUniqueInteger().getLimitedValue();

      Type *T = NewOperand->getType();
      if (T->isIntegerTy(16))
        Value = ByteSwap_16(Value);
      else if (T->isIntegerTy(32))
        Value = ByteSwap_32(Value);
      else if (T->isIntegerTy(64))
        Value = ByteSwap_64(Value);
      else
        llvm_unreachable("Unexpected type");

      NewOperand = ConstantInt::get(T, Value);
    } else {
      // Replace non-const operand with NewOperand
      std::vector<Constant *> Operands;
      bool NonConstFound = false;
      (void) NonConstFound;

      for (Value *Op : I->operand_values()) {
        if (auto *Const = dyn_cast<Constant>(Op)) {
          Operands.push_back(Const);
        } else {
          assert(!NonConstFound);
          NonConstFound = true;
          Operands.push_back(NewOperand);
        }
      }

      NewOperand = ConstantFoldInstOperands(I->getOpcode(),
                                            I->getType(),
                                            Operands,
                                            DL);
      assert(NewOperand != nullptr);
    }
  }

  // We made it, mark the value to be explored
  if (NewOperand != nullptr) {
    assert(!isa<UndefValue>(NewOperand));
    return getZExtValue(NewOperand, DL);
  }

  return 0;
}

void OperationsStack::explore(Constant *NewOperand) {
  uint64_t PC = materialize(NewOperand);

  if (PC != 0 && JTM->isInterestingPC(PC))
    PCs.insert({ PC, IsReliable });
}

/// \brief Simple Expression Tracker implementation
class SET {

public:
  SET(Function &F,
      JumpTargetManager *JTM,
      OSRAPass *OSRA,
      std::set<BasicBlock *> *Visited) :
    DL(F.getParent()->getDataLayout()),
    JTM(JTM),
    OS(JTM, DL),
    F(F),
    OSRA(OSRA),
    Visited(Visited) { }

  /// \brief Run the Simple Expression Tracker on F
  bool run();

private:
  /// \brief Enqueue all the store seen by the Start load instruction
  void enqueueStores(LoadInst *Start);
  /// \brief Process V
  /// \return a boolean indicating whether V has been handled properly and a new
  ///         Value from which SET should proceed
  std::pair<bool, Value *> handleInstruction(Instruction *Target, Value *V);

private:
  const unsigned MaxDepth = 3;
  const DataLayout &DL;
  JumpTargetManager *JTM;
  OperationsStack OS;
  Function& F;
  OSRAPass *OSRA;
  std::set<BasicBlock *> *Visited;
  std::vector<std::pair<Value *, unsigned>> WorkList;
};

void SET::enqueueStores(LoadInst *Start) {
  unsigned InitialHeight = OS.height();
  auto *Destination = Start->getPointerOperand();
  std::stack<std::pair<Instruction *, unsigned>> ToExplore;
  std::set<BasicBlock *> Visited;
  ToExplore.push(make_pair(Start, 0));

  Instruction *I = Start;

  while (!ToExplore.empty()) {
    unsigned Depth;
    std::tie(I, Depth) = ToExplore.top();
    ToExplore.pop();

    auto *BB = I->getParent();
    if (Visited.find(BB) != Visited.end())
      continue;

    Visited.insert(BB);
    BasicBlock::reverse_iterator It(make_reverse_iterator(I));
    BasicBlock::reverse_iterator Begin(BB->rend());

    bool Found = false;
    for (; It != Begin; It++) {
      if (auto *Store = dyn_cast<StoreInst>(&*It)) {
        if (Store->getPointerOperand() == Destination) {
          auto NewPair = make_pair(Store->getValueOperand(), InitialHeight);
          if (contains(WorkList, NewPair))
            WorkList.push_back(NewPair);
          Found = true;
          break;
        }
      }
    }

    // If we haven't find a store, proceed recursively in the predecessors
    if (!Found && Depth < MaxDepth)
      for (BasicBlock *Predecessor : make_range(pred_begin(BB), pred_end(BB)))
        if (Predecessor != JTM->dispatcher() && !Predecessor->empty())
          ToExplore.push(make_pair(&*Predecessor->rbegin(), Depth + 1));
  }

}

bool SET::run() {
  for (BasicBlock& BB : make_range(F.begin(), F.end())) {

    if (Visited->find(&BB) != Visited->end())
      continue;
    Visited->insert(&BB);

    for (Instruction& Instr : BB) {
      assert(Instr.getParent() == &BB);

      auto *Store = dyn_cast<StoreInst>(&Instr);
      auto *Load = dyn_cast<LoadInst>(&Instr);
      bool IsStore = Store != nullptr;
      bool IsPCStore = IsStore && JTM->isPCReg(Store->getPointerOperand());

      // Keep this for future use
      bool IsLoad = false && Load != nullptr;
      if ((!IsStore && !IsLoad)
          || (IsPCStore
              && isa<ConstantInt>(Store->getValueOperand()))
          || (IsLoad
              && (isa<GlobalVariable>(Load->getPointerOperand())
                  || isa<AllocaInst>(Load->getPointerOperand()))))
        continue;

      // Operations is a stack of ConstantInt uses in a BinaryOperator
      // TODO: hardcoded
      OS.reset(/* IsPCStore */ false);
      assert(WorkList.empty());
      if (IsStore)
        WorkList.push_back(make_pair(Store->getValueOperand(), 0));
      else
        WorkList.push_back(make_pair(Load->getPointerOperand(), 0));

      std::set<Value *> Visited;

      while (!WorkList.empty()) {
        unsigned Height;
        Value *V;
        std::tie(V, Height) = WorkList.back();
        WorkList.pop_back();

        if (Visited.find(V) != Visited.end())
          continue;
        Visited.insert(V);

        // Discard operations we no longer need
        OS.cut(Height);

        while (V != nullptr) {
          // TODO: use this
          bool Handled;
          std::tie(Handled, V) = handleInstruction(&Instr, V);
        }
      }
    }
  }

  OS.registerPCs();

  return false;
}

char SETPass::ID = 0;

void SETPass::getAnalysisUsage(AnalysisUsage &AU) const {
  if (UseOSRA)
    AU.addRequired<OSRAPass>();
}

bool SETPass::runOnFunction(Function &F) {
  OSRAPass *OSRA = getAnalysisIfAvailable<OSRAPass>();

  SET SimpleExpressionTracker(F, JTM, OSRA, Visited);
  return SimpleExpressionTracker.run();
}

std::pair<bool, Value *> SET::handleInstruction(Instruction *Target, Value *V) {
  // Setting handled to true makes return sucessfully but doesn't prevent
  // checking OSRA results
  bool Handled = false;

  if (auto *C = dyn_cast<ConstantInt>(V)) {
    // We reached the end of the path, materialize the value
    OS.explore(C);
    return { true, nullptr };
  } else if (auto *BinOp = dyn_cast<BinaryOperator>(V)) {

    // Append a reference to the operation to the Operations stack
    Use& FirstOp = BinOp->getOperandUse(0);
    Use& SecondOp = BinOp->getOperandUse(1);

    if (isa<ConstantInt>(FirstOp.get()) || isa<ConstantInt>(SecondOp.get())) {
      assert(!(isa<ConstantInt>(FirstOp.get())
               && isa<ConstantInt>(SecondOp.get())));
      bool FirstConstant = isa<ConstantInt>(FirstOp.get());

      // Add to the operations stack the constant one and proceed with the other
      if (OS.insertIfNew(BinOp))
        return { true, FirstConstant ? SecondOp.get() : FirstOp.get() };
    } else if (OSRA != nullptr) {
      Constant *ConstantOp = nullptr;
      Value *FreeOp = nullptr;
      Type *Int64 = Type::getInt64Ty(F.getParent()->getContext());
      std::tie(ConstantOp, FreeOp) = OSRA->identifyOperands(BinOp, Int64, DL);

      if (FreeOp == nullptr && ConstantOp != nullptr) {
        // The operation has been folded
        OS.explore(ConstantOp);
        Handled = true;
      } else if (FreeOp != nullptr && ConstantOp != nullptr) {
        // We were able to identify a constant operand
        unsigned FreeOpIndex = BinOp->getOperand(0) == FreeOp ? 0 : 1;

        // Note: the lifetime of the cloned instruction is managed by the
        //       OperationsStack
        Instruction *Clone = BinOp->clone();
        Clone->setOperand(1 - FreeOpIndex, ConstantOp);
        // This is a dirty trick to keep track of the original
        // instruction
        Clone->setOperand(FreeOpIndex, BinOp);

        // TODO: this might leave to infinte loops
        if (OS.insertIfNew(Clone, BinOp))
          return { true, BinOp->getOperandUse(FreeOpIndex).get() };
      }
    }
  } else if (auto *Load = dyn_cast<LoadInst>(V)) {
    auto *Pointer = Load->getPointerOperand();

    // If we're loading a global or local variable, look for the last write to
    // that variable, otherwise see if it's a load from a constant address which
    // points to a constant memory area
    if (isa<GlobalVariable>(Pointer) || isa<AllocaInst>(Pointer)) {
      enqueueStores(Load);
      Handled = true;
    } else {
      if (OS.insertIfNew(Load))
        return { true, Pointer };
    }
  } else if (auto *Unary = dyn_cast<UnaryInstruction>(V)) {
    if (OS.insertIfNew(Unary))
      return { true, Unary->getOperand(0) };
  } else if (auto *Expression = dyn_cast<ConstantExpr>(V)) {
    if (Expression->getNumOperands() == 1) {
      auto *ExprAsInstr = Expression->getAsInstruction();
      OS.insert(ExprAsInstr);
      return { true, Expression->getOperand(0) };
    }
  } else if (auto *Call = dyn_cast<CallInst>(V)) {
    Function *Callee = Call->getCalledFunction();
    if (Callee != nullptr && Callee->getIntrinsicID() == Intrinsic::bswap) {
      OS.insert(Call);
      return { true, Call->getArgOperand(0) };
    }
  } // End of the switch over instruction type

  // We don't know how to proceed, but we can still check if the current
  // instruction is associated with a suitable OSR
  if (OSRA != nullptr && !Handled && !OS.empty()) {
    const OSRAPass::OSR *O = OSRA->getOSR(V);
    if (O == nullptr)
      return { false, nullptr };

    using CI = ConstantInt;
    Type *Int64 = IntegerType::get(F.getParent()->getContext(), 64);
    if (O->isConstant()) {
      // If it's just a single constant, use it
      OS.explore(CI::get(Int64, O->base()));
      return { true, nullptr };
    } else if (!O->boundedValue()->isTop()
               && !O->boundedValue()->isBottom()
               && O->boundedValue()->isSingleRange()) {
      // We have a limited range, let's use it all

      // Perform a preliminary check that whole range fits into the executable
      // area
      Constant *MinConst, *MaxConst;
      std::tie(MinConst, MaxConst) = O->boundaries(Int64, DL);
      uint64_t Min = getZExtValue(MinConst, DL);
      uint64_t Max = getZExtValue(MaxConst, DL);
      uint64_t Step = O->factor();

      // TODO: the O->size() threshold is pretty arbitrary, the best solution
      //       here is probably restore it to int64_t::max(), assert if it's
      //       larger than 10000 and only apply it to store to memory, pc and
      //       maybe other registers (lr?)
      auto MaterializedMin = OS.materialize(MinConst);
      auto MaterializedMax = OS.materialize(MaxConst);
      auto MaterializedStep = OS.materialize(CI::get(Int64, Step));
      if (!JTM->isExecutableRange(MaterializedMin, MaterializedMax)
          || !JTM->isInstructionAligned(MaterializedStep)
          || O->size() >= 10000)
        return { false, nullptr };

      if (O->size() > 1000)
        dbg << "Warning: " << O->size() << " jump targets added\n";

      DBG("osrjts", dbg << "Adding " << std::dec << O->size()
          << " jump targets from 0x"
          << std::hex << JTM->getPC(Target).first << "\n");

      // Note: addition and comparison for equality are all sign-safe
      // operations, no need to use Constants in this case.
      // TODO: switch to a super-elegant iterator
      for (uint64_t Position = Min; Position != Max; Position += Step)
        OS.explore(CI::get(Int64, Position));
      OS.explore(CI::get(Int64, Max));
      return { true, nullptr };
    }
  }

  return { Handled, nullptr };
}