Engine.h

#ifndef ROSE_Partitioner2_Engine_H
#define ROSE_Partitioner2_Engine_H

#include <BinaryLoader.h>
#include <Disassembler.h>
#include <FileSystem.h>
#include <Partitioner2/Function.h>
#include <Partitioner2/Partitioner.h>
#include <Partitioner2/Utility.h>
#include <Progress.h>
#include <Sawyer/DistinctList.h>

namespace Rose {
namespace BinaryAnalysis {
namespace Partitioner2 {

class ROSE_DLL_API Engine {
public:
    struct Settings {
        LoaderSettings loader;                          
        DisassemblerSettings disassembler;              
        PartitionerSettings partitioner;                
        EngineSettings engine;                          
        AstConstructionSettings astConstruction;        
    private:
        friend class boost::serialization::access;

        template<class S>
        void serialize(S &s, unsigned version) {
            s & loader & disassembler & partitioner & engine & astConstruction;
        }
    };

    //                                  Internal data structures
private:
    // Engine callback for handling instructions added to basic blocks.  This is called when a basic block is discovered,
    // before it's attached to a partitioner, so it shouldn't really be modifying any state in the engine, but rather only
    // preparing the basic block to be processed.
    class BasicBlockFinalizer: public BasicBlockCallback {
        typedef Sawyer::Container::Map<rose_addr_t /*target*/, std::vector<rose_addr_t> /*sources*/> WorkList;
    public:
        static Ptr instance() { return Ptr(new BasicBlockFinalizer); }
        virtual bool operator()(bool chain, const Args &args) ROSE_OVERRIDE;
    };
    
    // Basic blocks that need to be worked on next. These lists are adjusted whenever a new basic block (or placeholder) is
    // inserted or erased from the CFG.
    class BasicBlockWorkList: public CfgAdjustmentCallback {
        Sawyer::Container::DistinctList<rose_addr_t> pendingCallReturn_;   // blocks that might need an E_CALL_RETURN edge
        Sawyer::Container::DistinctList<rose_addr_t> processedCallReturn_; // call sites whose may-return was indeterminate
        Sawyer::Container::DistinctList<rose_addr_t> finalCallReturn_;     // indeterminate call sites awaiting final analysis
        Sawyer::Container::DistinctList<rose_addr_t> undiscovered_;        // undiscovered basic block list (last-in-first-out)
        Engine *engine_;                                                   // engine to which this callback belongs
    protected:
        explicit BasicBlockWorkList(Engine *engine): engine_(engine) {}
    public:
        typedef Sawyer::SharedPointer<BasicBlockWorkList> Ptr;
        static Ptr instance(Engine *engine) { return Ptr(new BasicBlockWorkList(engine)); }
        virtual bool operator()(bool chain, const AttachedBasicBlock &args) ROSE_OVERRIDE;
        virtual bool operator()(bool chain, const DetachedBasicBlock &args) ROSE_OVERRIDE;
        Sawyer::Container::DistinctList<rose_addr_t>& pendingCallReturn() { return pendingCallReturn_; }
        Sawyer::Container::DistinctList<rose_addr_t>& processedCallReturn() { return processedCallReturn_; }
        Sawyer::Container::DistinctList<rose_addr_t>& finalCallReturn() { return finalCallReturn_; }
        Sawyer::Container::DistinctList<rose_addr_t>& undiscovered() { return undiscovered_; }
        void moveAndSortCallReturn(const Partitioner&);
    };

    // A work list providing constants from instructions that are part of the CFG.
    class CodeConstants: public CfgAdjustmentCallback {
    public:
        typedef Sawyer::SharedPointer<CodeConstants> Ptr;

    private:
        std::set<rose_addr_t> toBeExamined_;            // instructions waiting to be examined
        std::set<rose_addr_t> wasExamined_;             // instructions we've already examined
        rose_addr_t inProgress_;                        // instruction that is currently in progress
        std::vector<rose_addr_t> constants_;            // constants for the instruction in progress

    protected:
        CodeConstants(): inProgress_(0) {}

    public:
        static Ptr instance() { return Ptr(new CodeConstants); }

        // Possibly insert more instructions into the work list when a basic block is added to the CFG
        virtual bool operator()(bool chain, const AttachedBasicBlock &attached) ROSE_OVERRIDE;

        // Possibly remove instructions from the worklist when a basic block is removed from the CFG
        virtual bool operator()(bool chain, const DetachedBasicBlock &detached) ROSE_OVERRIDE;

        // Return the next available constant if any.
        Sawyer::Optional<rose_addr_t> nextConstant(const Partitioner &partitioner);
    };

    //                                  Data members
private:
    Settings settings_;                                 // Settings for the partitioner.
    SgAsmInterpretation *interp_;                       // interpretation set by loadSpecimen
    BinaryLoader *binaryLoader_;                        // how to remap, link, and fixup
    Disassembler *disassembler_;                        // not ref-counted yet, but don't destroy it since user owns it
    MemoryMap::Ptr map_;                                // memory map initialized by load()
    BasicBlockWorkList::Ptr basicBlockWorkList_;        // what blocks to work on next
    CodeConstants::Ptr codeFunctionPointers_;           // generates constants that are found in instruction ASTs
    Progress::Ptr progress_;                            // optional progress reporting

    //                                  Constructors
public:
    Engine()
        : interp_(NULL), binaryLoader_(NULL), disassembler_(NULL), basicBlockWorkList_(BasicBlockWorkList::instance(this)),
        progress_(Progress::instance()) {
        init();
    }

    explicit Engine(const Settings &settings)
        : settings_(settings), interp_(NULL), binaryLoader_(NULL), disassembler_(NULL),
        basicBlockWorkList_(BasicBlockWorkList::instance(this)), progress_(Progress::instance()) {
        init();
    }

    virtual ~Engine() {}

    //                                  The very top-level use case
public:
    SgAsmBlock* frontend(int argc, char *argv[],
                         const std::string &purpose, const std::string &description);
    virtual SgAsmBlock* frontend(const std::vector<std::string> &args,
                                 const std::string &purpose, const std::string &description);
    //                                  Basic top-level steps
public:
    void reset();

    Sawyer::CommandLine::ParserResult parseCommandLine(int argc, char *argv[],
                                                       const std::string &purpose, const std::string &description) /*final*/;
    virtual Sawyer::CommandLine::ParserResult parseCommandLine(const std::vector<std::string> &args,
                                                               const std::string &purpose, const std::string &description);
    virtual SgAsmInterpretation* parseContainers(const std::vector<std::string> &fileNames);
    SgAsmInterpretation* parseContainers(const std::string &fileName) /*final*/;
    virtual MemoryMap::Ptr loadSpecimens(const std::vector<std::string> &fileNames = std::vector<std::string>());
    MemoryMap::Ptr loadSpecimens(const std::string &fileName) /*final*/;
    virtual Partitioner partition(const std::vector<std::string> &fileNames = std::vector<std::string>());
    Partitioner partition(const std::string &fileName) /*final*/;
    SgAsmBlock* buildAst(const std::vector<std::string> &fileNames = std::vector<std::string>()) /*final*/;
    SgAsmBlock* buildAst(const std::string &fileName) /*final*/;
    //                                  Command-line parsing
    //
    // top-level: parseCommandLine
public:
    virtual Sawyer::CommandLine::SwitchGroup loaderSwitches();

    virtual Sawyer::CommandLine::SwitchGroup disassemblerSwitches();

    virtual Sawyer::CommandLine::SwitchGroup partitionerSwitches();

    virtual Sawyer::CommandLine::SwitchGroup engineSwitches();

    virtual Sawyer::CommandLine::SwitchGroup astConstructionSwitches();

    static std::string specimenNameDocumentation();

    virtual Sawyer::CommandLine::Parser commandLineParser(const std::string &purpose, const std::string &description);

    virtual void checkSettings();

    //                                  Container parsing
    //
    // top-level: parseContainers
public:
    virtual bool isNonContainer(const std::string&);

    virtual bool areContainersParsed() const;

    //                                  Load specimens
    //
    // top-level: loadSpecimens
public:
    virtual bool areSpecimensLoaded() const;

    virtual BinaryLoader *obtainLoader(BinaryLoader *hint=NULL);

    virtual void loadContainers(const std::vector<std::string> &fileNames);

    virtual void loadNonContainers(const std::vector<std::string> &names);

    virtual void adjustMemoryMap();

    MemoryMap::Ptr memoryMap() const /*final*/ { return map_; }
    virtual void memoryMap(const MemoryMap::Ptr &m) { map_ = m; }
    //                                  Disassembler
public:
    virtual Disassembler* obtainDisassembler(Disassembler *hint=NULL);
    //                                  Partitioner high-level functions
    //
    // top-level: partition
public:
    virtual void checkCreatePartitionerPrerequisites() const;

    virtual Partitioner createBarePartitioner();

    virtual Partitioner createGenericPartitioner();

    virtual Partitioner createTunedPartitioner();

    virtual Partitioner createPartitionerFromAst(SgAsmInterpretation*);

    virtual Partitioner createPartitioner();

    virtual void runPartitionerInit(Partitioner&);

    virtual void runPartitionerRecursive(Partitioner&);

    virtual void runPartitionerFinal(Partitioner&);
    
    virtual void runPartitioner(Partitioner&);


    //                                  Partitioner mid-level functions
    //
    // These are the functions called by the partitioner high-level stuff.  These are sometimes overridden in subclasses,
    // although it is more likely that the high-level stuff is overridden.
public:
    virtual void labelAddresses(Partitioner&);

    virtual std::vector<DataBlock::Ptr> makeConfiguredDataBlocks(Partitioner&, const Configuration&);

    virtual std::vector<Function::Ptr> makeConfiguredFunctions(Partitioner&, const Configuration&);

    virtual std::vector<Function::Ptr> makeEntryFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeErrorHandlingFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeImportFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeExportFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeSymbolFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeContainerFunctions(Partitioner&, SgAsmInterpretation*);

    virtual std::vector<Function::Ptr> makeInterruptVectorFunctions(Partitioner&, const AddressInterval &vector);

    virtual std::vector<Function::Ptr> makeUserFunctions(Partitioner&, const std::vector<rose_addr_t>&);

    virtual void discoverBasicBlocks(Partitioner&);

    virtual Function::Ptr makeNextDataReferencedFunction(const Partitioner&, rose_addr_t &startVa /*in,out*/);

    virtual Function::Ptr makeNextCodeReferencedFunction(const Partitioner&);

    virtual std::vector<Function::Ptr> makeCalledFunctions(Partitioner&);

    virtual std::vector<Function::Ptr> makeNextPrologueFunction(Partitioner&, rose_addr_t startVa);

    virtual std::vector<Function::Ptr>
    makeFunctionFromInterFunctionCalls(Partitioner &partitioner, rose_addr_t &startVa /*in,out*/);

    virtual void discoverFunctions(Partitioner&);

    virtual std::set<rose_addr_t> attachDeadCodeToFunction(Partitioner&, const Function::Ptr&, size_t maxIterations=size_t(-1));

    virtual DataBlock::Ptr attachPaddingToFunction(Partitioner&, const Function::Ptr&);

    virtual std::vector<DataBlock::Ptr> attachPaddingToFunctions(Partitioner&);

    virtual size_t attachAllSurroundedCodeToFunctions(Partitioner&);
    
    virtual size_t attachSurroundedCodeToFunctions(Partitioner&);

    virtual void attachBlocksToFunctions(Partitioner&);

    virtual std::set<rose_addr_t> attachDeadCodeToFunctions(Partitioner&, size_t maxIterations=size_t(-1));

    virtual std::vector<DataBlock::Ptr> attachSurroundedDataToFunctions(Partitioner&);

    virtual void updateAnalysisResults(Partitioner&);


    //                                  Partitioner low-level functions
    //
    // These are functions that a subclass seldom overrides, and maybe even shouldn't override because of their complexity or
    // the way the interact with one another.
public:
    virtual bool makeNextCallReturnEdge(Partitioner&, boost::logic::tribool assumeCallReturns);

    virtual BasicBlock::Ptr makeNextBasicBlockFromPlaceholder(Partitioner&);

    virtual BasicBlock::Ptr makeNextBasicBlock(Partitioner&);


    //                                  Build AST
public:
    // Used internally by ROSE's ::frontend disassemble instructions to build the AST that goes under each SgAsmInterpretation.
    static void disassembleForRoseFrontend(SgAsmInterpretation*);


    //                                  Settings and properties
public:
    const Settings& settings() const /*final*/ { return settings_; }
    Settings& settings() /*final*/ { return settings_; }
    bool exitOnError() const /*final*/ { return settings_.engine.exitOnError; }
    virtual void exitOnError(bool b) { settings_.engine.exitOnError = b; }
    Progress::Ptr progress() const /*final*/ { return progress_; }
    virtual void progress(const Progress::Ptr &progress) { progress_ = progress; }
    SgAsmInterpretation* interpretation() const /*final*/ { return interp_; }
    virtual void interpretation(SgAsmInterpretation *interp) { interp_ = interp; }
    BinaryLoader* binaryLoader() const /*final*/ { return binaryLoader_; }
    virtual void binaryLoader(BinaryLoader *loader) { binaryLoader_ = loader; }
    size_t deExecuteZerosThreshold() const /*final*/ { return settings_.loader.deExecuteZerosThreshold; }
    virtual void deExecuteZerosThreshold(size_t n) { settings_.loader.deExecuteZerosThreshold = n; }
    size_t deExecuteZerosLeaveAtFront() const /*final*/ { return settings_.loader.deExecuteZerosLeaveAtFront; }
    virtual void deExecuteZerosLeaveAtFront(size_t n) { settings_.loader.deExecuteZerosLeaveAtFront = n; }
    size_t deExecuteZerosLeaveAtBack() const /*final*/ { return settings_.loader.deExecuteZerosLeaveAtBack; }
    virtual void deExecuteZerosLeaveAtBack(size_t n) { settings_.loader.deExecuteZerosLeaveAtBack = n; }
    MemoryDataAdjustment memoryDataAdjustment() const /*final*/ { return settings_.loader.memoryDataAdjustment; }
    virtual void memoryDataAdjustment(MemoryDataAdjustment x) { settings_.loader.memoryDataAdjustment = x; }
    bool memoryIsExecutable() const /*final*/ { return settings_.loader.memoryIsExecutable; }
    virtual void memoryIsExecutable(bool b) { settings_.loader.memoryIsExecutable = b; }
    Disassembler *disassembler() const /*final*/ { return disassembler_; }
    virtual void disassembler(Disassembler *d) { disassembler_ = d; }
    const std::string& isaName() const /*final*/ { return settings_.disassembler.isaName; }
    virtual void isaName(const std::string &s) { settings_.disassembler.isaName = s; }
    const std::vector<rose_addr_t>& startingVas() const /*final*/ { return settings_.partitioner.startingVas; }
    std::vector<rose_addr_t>& startingVas() /*final*/ { return settings_.partitioner.startingVas; }
    bool usingSemantics() const /*final*/ { return settings_.partitioner.base.usingSemantics; }
    virtual void usingSemantics(bool b) { settings_.partitioner.base.usingSemantics = b; }
    SemanticMemoryParadigm semanticMemoryParadigm() const /*final*/ { return settings_.partitioner.semanticMemoryParadigm; }
    virtual void semanticMemoryParadigm(SemanticMemoryParadigm p) { settings_.partitioner.semanticMemoryParadigm = p; }
    bool followingGhostEdges() const /*final*/ { return settings_.partitioner.followingGhostEdges; }
    virtual void followingGhostEdges(bool b) { settings_.partitioner.followingGhostEdges = b; }
    bool discontiguousBlocks() const /*final*/ { return settings_.partitioner.discontiguousBlocks; }
    virtual void discontiguousBlocks(bool b) { settings_.partitioner.discontiguousBlocks = b; }
    bool findingFunctionPadding() const /*final*/ { return settings_.partitioner.findingFunctionPadding; }
    virtual void findingFunctionPadding(bool b) { settings_.partitioner.findingFunctionPadding = b; }
    bool findingThunks() const /*final*/ { return settings_.partitioner.findingThunks; }
    virtual void findingThunks(bool b) { settings_.partitioner.findingThunks = b; }
    bool splittingThunks() const /*final*/ { return settings_.partitioner.splittingThunks; }
    virtual void splittingThunks(bool b) { settings_.partitioner.splittingThunks = b; }
    bool findingDeadCode() const /*final*/ { return settings_.partitioner.findingDeadCode; }
    virtual void findingDeadCode(bool b) { settings_.partitioner.findingDeadCode = b; }
    rose_addr_t peScramblerDispatcherVa() const /*final*/ { return settings_.partitioner.peScramblerDispatcherVa; }
    virtual void peScramblerDispatcherVa(rose_addr_t va) { settings_.partitioner.peScramblerDispatcherVa = va; }
    bool findingIntraFunctionCode() const /*final*/ { return settings_.partitioner.findingIntraFunctionCode; }
    virtual void findingIntraFunctionCode(bool b) { settings_.partitioner.findingIntraFunctionCode = b; }
    bool findingIntraFunctionData() const /*final*/ { return settings_.partitioner.findingIntraFunctionData; }
    virtual void findingIntraFunctionData(bool b) { settings_.partitioner.findingIntraFunctionData = b; }
    const AddressInterval& interruptVector() const /*final*/ { return settings_.partitioner.interruptVector; }
    virtual void interruptVector(const AddressInterval &i) { settings_.partitioner.interruptVector = i; }
    bool doingPostAnalysis() const /*final*/ { return settings_.partitioner.doingPostAnalysis; }
    virtual void doingPostAnalysis(bool b) { settings_.partitioner.doingPostAnalysis = b; }
    bool doingPostFunctionMayReturn() const /*final*/ { return settings_.partitioner.doingPostFunctionMayReturn; }
    virtual void doingPostFunctionMayReturn(bool b) { settings_.partitioner.doingPostFunctionMayReturn = b; }
    bool doingPostFunctionStackDelta() const /*final*/ { return settings_.partitioner.doingPostFunctionStackDelta; }
    virtual void doingPostFunctionStackDelta(bool b) { settings_.partitioner.doingPostFunctionStackDelta = b; }
    bool doingPostCallingConvention() const /*final*/ { return settings_.partitioner.doingPostCallingConvention; }
    virtual void doingPostCallingConvention(bool b) { settings_.partitioner.doingPostCallingConvention = b; }
    bool doingPostFunctionNoop() const /*final*/ { return settings_.partitioner.doingPostFunctionNoop; }
    virtual void doingPostFunctionNoop(bool b) { settings_.partitioner.doingPostFunctionNoop = b; }
    FunctionReturnAnalysis functionReturnAnalysis() const /*final*/ { return settings_.partitioner.functionReturnAnalysis; }
    virtual void functionReturnAnalysis(FunctionReturnAnalysis x) { settings_.partitioner.functionReturnAnalysis = x; }
    bool findingInterFunctionCalls() const /*final*/ { return settings_.partitioner.findingInterFunctionCalls; }
    virtual void findingInterFunctionCalls(bool b) { settings_.partitioner.findingInterFunctionCalls = b; }
    bool findingDataFunctionPointers() const /*final*/ { return settings_.partitioner.findingDataFunctionPointers; }
    virtual void findingDataFunctionPointers(bool b) { settings_.partitioner.findingDataFunctionPointers = b; }
    bool findingCodeFunctionPointers() const /*final*/ { return settings_.partitioner.findingCodeFunctionPointers; }
    virtual void findingCodeFunctionPointers(bool b) { settings_.partitioner.findingCodeFunctionPointers = b; }
    bool checkingCallBranch() const /*final*/ { return settings_.partitioner.base.checkingCallBranch; }
    virtual void checkingCallBranch(bool b) { settings_.partitioner.base.checkingCallBranch = b; }
    bool basicBlockSemanticsAutoDrop() const /*final*/ { return settings_.partitioner.base.basicBlockSemanticsAutoDrop; }
    void basicBlockSemanticsAutoDrop(bool b) { settings_.partitioner.base.basicBlockSemanticsAutoDrop = b; }
    const std::vector<std::string>& configurationNames() /*final*/ const { return settings_.engine.configurationNames; }
    std::vector<std::string>& configurationNames() /*final*/ { return settings_.engine.configurationNames; }
    bool namingConstants() const /*final*/ { return settings_.partitioner.namingConstants; }
    virtual void namingConstants(bool b) { settings_.partitioner.namingConstants = b; }
    bool namingStrings() const /*final*/ { return settings_.partitioner.namingStrings; }
    virtual void namingStrings(bool b) { settings_.partitioner.namingStrings = b; }
    bool demangleNames() const /*final*/ { return settings_.partitioner.demangleNames; }
    virtual void demangleNames(bool b) { settings_.partitioner.demangleNames = b; }
    bool astAllowEmptyGlobalBlock() const /*final*/ { return settings_.astConstruction.allowEmptyGlobalBlock; }
    virtual void astAllowEmptyGlobalBlock(bool b) { settings_.astConstruction.allowEmptyGlobalBlock = b; }
    bool astAllowFunctionWithNoBasicBlocks() const /*final*/ {
        return settings_.astConstruction.allowFunctionWithNoBasicBlocks;
    }
    virtual void astAllowFunctionWithNoBasicBlocks(bool b) {
        settings_.astConstruction.allowFunctionWithNoBasicBlocks = b;
    }
    bool astAllowEmptyBasicBlock() const /*final*/ { return settings_.astConstruction.allowEmptyBasicBlocks; }
    virtual void astAllowEmptyBasicBlock(bool b) { settings_.astConstruction.allowEmptyBasicBlocks = b; }
    bool astCopyAllInstructions() const /*final*/ { return settings_.astConstruction.copyAllInstructions; }
    virtual void astCopyAllInstructions(bool b) { settings_.astConstruction.copyAllInstructions = b; }
    //                                  Internal stuff
private:
    void init();
};

} // namespace
} // namespace
} // namespace

#endif