SEGCallSite.h

/*
 * SEGCallSite.h
 *
 * Qingkai
 *
 * This node is to model a call site.
 */
 
#ifndef IR_SEG_SEGCALLSITE_H
#define IR_SEG_SEGCALLSITE_H
 
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/CallSite.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/Value.h>
 
#include <map>
#include <vector>
 
#include "IR/LLVMWrapper/FunctionWrapper.h"
#include "IR/SEG/SEGCallSitePseudoInputNode.h"
#include "IR/SEG/SymbolicExprGraph.h"
 
using namespace llvm;
 
class SEGCallSiteOutputNode;
class SEGCallSiteCommonOutputNode;
class SEGCallSitePseudoOutputNode;
class SEGCallSiteArgumentSummaryNode;
class PersistedSEGCallSite;
 
typedef struct SEGCallSiteInputStruct {
  SEGOperandNode *InputNode;
  size_t InputIndex;
  bool IsCommonInput;
 
  SEGCallSiteInputStruct(SEGOperandNode *N, size_t I, bool C)
      : InputNode(N), InputIndex(I), IsCommonInput(C) {}
} SEGCallSiteInputStruct;
 
typedef struct SideEffectIO {
  std::vector<SEGCallSiteInputStruct> PseudoInputs;
  std::vector<SEGCallSitePseudoOutputNode *> PseudoOutputs;
  bool IsBackEdge = false;
  bool IsExpanded = false;
} SideEffectIO;
 
namespace FalconPlus {
class IntraFalcon;
}
 
class SEGCallSite : public SEGSiteBase {
public:
  SEGCallSite(PersistedSEGCallSite *Site, SymbolicExprGraph *SEG)
      : SEGSiteBase(Site, SEG) {}
 
  virtual void assembleSEGObject(std::map<int, SEGObject *> &FuncSEGObjMap);
 
private:
  std::set<const SEGNodeBase *> CommonInputNodes;
 
  std::vector<SEGCallSiteInputStruct> CommonInputs;
  SEGCallSiteCommonOutputNode *CommonOutput = nullptr;
 
  std::map<WrappedFunction *, SideEffectIO, CmpPointerOfFunction> SideEffectMap;
 
  // Input and output summary nodes
  // Index 0 means the ap-depth larger than considered
  // Index from 1 to InputSummaryNode.size(), is the summary nodes for the
  // corresponding ap-depth
  // TODO: Summary outputs are not enabled now
  kvec<SEGCallSiteArgumentSummaryNode *> InputSummaryNodes;
  kvec<SEGOperandNode *> OutputSummaryNodes;
 
  SEGCallSite(Instruction *CS, SymbolicExprGraph *G);
 
  friend class IntraFalcon;
  friend class FalconPlus::IntraFalcon;
  friend class MantaIntraFalcon;
  friend class SymbolicExprGraph;
  friend class SymbolicExprGraphBuilder;
  friend class OCValueFlowBuilder;
 
private:
  void addPseudoOutput(Function *, SEGCallSitePseudoOutputNode *Node);
 
  void addPseudoOutput(WrappedFunction *WrappedCallee,
                       SEGCallSitePseudoOutputNode *Node);
 
  void addPseudoInput(Function *, SEGOperandNode *Node);
 
  void addPseudoInput(WrappedFunction *WrappedCallee, SEGOperandNode *Node);
 
  void setCommonOutput(SEGCallSiteCommonOutputNode *Node);
 
  void addCommonInput(SEGOperandNode *Node);
 
  void addCallee(Function *Callee) {
    if (getCalledFunction()) {
      assert(getCalledFunction() == Callee);
    }
    auto *WrappedCallee = WrappedValue::findWrappedFunction(Callee);
    SideEffectMap[WrappedCallee];
  }
 
  // set input/output summary node with ap-depth
  void setInputSummaryNode(int APDepth,
                           SEGCallSiteArgumentSummaryNode *SummaryNode);
  void setOutputSummaryNode(int APDepth, SEGOperandNode *SummaryNode);
 
public:
  void setBackEdge(WrappedFunction *Callee) {
    SideEffectMap[Callee].IsBackEdge = true;
  }
 
  bool isBackEdge(WrappedFunction *Callee) const {
    if (SideEffectMap.count(Callee)) {
      return SideEffectMap.at(Callee).IsBackEdge;
    }
    return false;
  }
 
  void setExpanded(WrappedFunction *Callee) {
    SideEffectMap[Callee].IsExpanded = true;
  }
 
  bool isExpanded(WrappedFunction *Callee) const {
    if (Callee && SideEffectMap.count(Callee)) {
      return SideEffectMap.at(Callee).IsExpanded;
    }
    return false;
  }
 
  virtual PersistedSEGObject *createPersistedObject() const;
 
  const SEGOperandNode *getCommonInput(size_t Index) const;
 
  const SEGOperandNode *getPseudoInput(Function *, size_t Index) const;
 
  const SEGOperandNode *getPseudoInput(WrappedFunction *, size_t Index) const;
 
  const SEGCallSiteCommonOutputNode *getCommonOutput() const;
 
  const SEGCallSitePseudoOutputNode *getPseudoOutput(Function *,
                                                     size_t Index) const;
 
  const SEGCallSitePseudoOutputNode *getPseudoOutput(WrappedFunction *,
                                                     size_t Index) const;
 
  bool hasCommonOutput() const;
 
  bool isCommonInput(const SEGNodeBase *Node) const {
    return this->CommonInputNodes.count(Node);
  }
 
  size_t getNumOutputs(Function *) const;
 
  size_t getNumPseudoOutputs(Function *) const;
 
  size_t getNumInputs(Function *) const;
 
  size_t getNumCommonInputs() const;
 
  size_t getNumPseudoInputs(Function *) const;
 
  size_t getInputIndex(const SEGNodeBase *Input) const;
 
  bool hasCallee(Function *Func) const {
    auto *WrappedFunc = WrappedValue::findWrappedFunction(Func);
    return SideEffectMap.count(WrappedFunc);
  }
 
  Function *getCalledFunction() const {
    return this->getLLVMCallSite().getCalledFunction();
  }
 
  WrappedFunction *getWrappedCalledFunction() const {
    return this->getWrappedLLVMCallSite().getCalledFunction();
  }
 
  static bool classof(const SEGObject *N) {
    return N->getKind() == SEGOBJK_CallSite;
  }
 
  /*==------------------------Iterators--------------------==*/
 
  class CalleeIterator {
  private:
    std::map<WrappedFunction *, SideEffectIO,
             CmpPointerOfFunction>::const_iterator It;
 
  public:
    CalleeIterator(bool End, const std::map<WrappedFunction *, SideEffectIO,
                                            CmpPointerOfFunction> &M) {
      if (End) {
        It = M.end();
      } else {
        It = M.begin();
      }
    }
 
    CalleeIterator(const CalleeIterator &CIt) : It(CIt.It) {}
 
    CalleeIterator &operator=(const CalleeIterator &CIt) {
      if (this != &CIt) {
        this->It = CIt.It;
      }
      return *this;
    }
 
    CalleeIterator operator++(int) {
      CalleeIterator Old(*this);
      ++(*this);
      return Old;
    }
 
    CalleeIterator &operator++() {
      It++;
      return *this;
    }
 
    WrappedFunction *operator*() { return It->first; }
 
    bool operator==(const CalleeIterator &CIt) { return It == CIt.It; }
 
    bool operator!=(const CalleeIterator &CIt) { return It != CIt.It; }
  };
 
  class InputIterator {
  private:
    size_t I;
 
    const std::vector<SEGCallSiteInputStruct> &CommonIns;
    const std::vector<SEGCallSiteInputStruct> &PseudoIns;
 
  public:
    InputIterator(bool End, const std::vector<SEGCallSiteInputStruct> &CI,
                  const std::vector<SEGCallSiteInputStruct> &PI)
        : I(End ? CI.size() + PI.size() : 0), CommonIns(CI), PseudoIns(PI) {}
 
    InputIterator(const InputIterator &It)
        : I(It.I), CommonIns(It.CommonIns), PseudoIns(It.PseudoIns) {}
 
    InputIterator operator++(int) {
      InputIterator Old(*this);
      ++(*this);
      return Old;
    }
 
    InputIterator &operator++() {
      I++;
      return *this;
    }
 
    const SEGCallSiteInputStruct &operator*() {
      return I < CommonIns.size() ? CommonIns[I]
                                  : PseudoIns[I - CommonIns.size()];
    }
 
    bool operator==(const InputIterator &It) { return I == It.I; }
 
    bool operator!=(const InputIterator &It) { return I != It.I; }
  };
 
  class OutputIterator {
  private:
    size_t I;
 
    SEGCallSiteCommonOutputNode *CommonOut;
    const std::vector<SEGCallSitePseudoOutputNode *> &PseudoOuts;
 
  public:
    OutputIterator(bool End, SEGCallSiteCommonOutputNode *CO,
                   const std::vector<SEGCallSitePseudoOutputNode *> &POs)
        : I(End ? (POs.size() + (CO ? 1 : 0)) : 0), CommonOut(CO),
          PseudoOuts(POs) {}
 
    OutputIterator(const OutputIterator &It)
        : I(It.I), CommonOut(It.CommonOut), PseudoOuts(It.PseudoOuts) {}
 
    OutputIterator operator++(int) {
      OutputIterator Old(*this);
      ++(*this);
      return Old;
    }
 
    OutputIterator &operator++() {
      I++;
      return *this;
    }
 
    const SEGCallSiteOutputNode *operator*() {
      if (CommonOut) {
        if (I == 0) {
          return (const SEGCallSiteOutputNode *)CommonOut;
        } else {
          assert(I >= 1);
          return (const SEGCallSiteOutputNode *)PseudoOuts[I - 1];
        }
      } else {
        return (const SEGCallSiteOutputNode *)PseudoOuts[I];
      }
    }
 
    bool operator==(const OutputIterator &It) { return I == It.I; }
 
    bool operator!=(const OutputIterator &It) { return I != It.I; }
  };
 
  CalleeIterator callee_begin() const {
    return CalleeIterator(false, SideEffectMap);
  }
 
  CalleeIterator callee_end() const {
    return CalleeIterator(true, SideEffectMap);
  }
 
  size_t callee_size() const { return SideEffectMap.size(); }
 
  InputIterator input_begin(WrappedFunction *Callee) const {
    auto It = SideEffectMap.find(Callee);
    assert(It != SideEffectMap.end());
    auto &PseudoInputsVec = It->second.PseudoInputs;
    return InputIterator(false, CommonInputs, PseudoInputsVec);
  }
 
  InputIterator input_begin(Function *Callee) const {
    auto *WrappedCallee = WrappedValue::findWrappedFunction(Callee);
    return input_begin(WrappedCallee);
  }
 
  InputIterator input_end(WrappedFunction *Callee) const {
    auto It = SideEffectMap.find(Callee);
    assert(It != SideEffectMap.end());
    auto &PseudoInputsVec = It->second.PseudoInputs;
    return InputIterator(true, CommonInputs, PseudoInputsVec);
  }
 
  InputIterator input_end(Function *Callee) const {
    auto *WrappedCallee = WrappedValue::findWrappedFunction(Callee);
    return input_end(WrappedCallee);
  }
 
  std::vector<SEGCallSiteInputStruct>::const_iterator
  common_input_begin() const {
    return CommonInputs.begin();
  }
 
  std::vector<SEGCallSiteInputStruct>::const_iterator common_input_end() const {
    return CommonInputs.end();
  }
 
  OutputIterator output_begin(WrappedFunction *Callee) const {
    auto It = SideEffectMap.find(Callee);
    assert(It != SideEffectMap.end());
    auto &PseudoOutputsVec = It->second.PseudoOutputs;
    return OutputIterator(false, CommonOutput, PseudoOutputsVec);
  }
 
  OutputIterator output_begin(Function *Callee) const {
    auto *WrappedCallee = WrappedValue::findWrappedFunction(Callee);
    return output_begin(WrappedCallee);
  }
 
  OutputIterator output_end(WrappedFunction *Callee) const {
    auto It = SideEffectMap.find(Callee);
    assert(It != SideEffectMap.end());
    auto &PseudoOutputsVec = It->second.PseudoOutputs;
    return OutputIterator(true, CommonOutput, PseudoOutputsVec);
  }
 
  OutputIterator output_end(Function *Callee) const {
    auto *WrappedCallee = WrappedValue::findWrappedFunction(Callee);
    return output_end(WrappedCallee);
  }
};
 
#endif