ParallelPartitioner.h

#ifndef ROSE_BinaryAnalysis_Partitioner2_ParallelPartitioner_H
#define ROSE_BinaryAnalysis_Partitioner2_ParallelPartitioner_H
#include <featureTests.h>
#if defined(ROSE_ENABLE_BINARY_ANALYSIS) && __cplusplus >= 201103L

#include <Rose/BinaryAnalysis/InstructionCache.h>
#include <Rose/Progress.h>
#include <queue>
#include <Rose/BinaryAnalysis/Partitioner2/BasicTypes.h>
#include <Rose/BinaryAnalysis/Partitioner2/Semantics.h>
#include <Sawyer/Attribute.h>
#include <Sawyer/Graph.h>
#include <Sawyer/IntervalSetMap.h>
#include <Sawyer/Message.h>
#include <Sawyer/SharedObject.h>

namespace Rose {
namespace BinaryAnalysis {
namespace Partitioner2 {
namespace Experimental {
namespace ParallelPartitioner {

class Partitioner;

// Settings and configuration, including miscellaneous basic types.

enum class Accuracy {
    LOW,                                                
    HIGH,                                               
    DEFAULT,                                            
};

using FunctionReasons = BitFlags<SgAsmFunction::FunctionReason>;

struct Settings {
    size_t maxAnalysisBBlockSize = 20; 
    Accuracy successorAccuracy = Accuracy::LOW; 
    Accuracy functionCallDetectionAccuracy = Accuracy::LOW; 
    size_t minHoleSearch = 8; 
    SemanticMemoryParadigm semanticMemoryParadigm = MAP_BASED_MEMORY; 
};

// Diagnostics

extern Sawyer::Message::Facility mlog;

// Used internally. See Rose::Diagnostics::initialize
void initDiagnostics();

// Cached information

template<class V, class K>
class CachedItem {
public:
    using Value = V;
    using Key = K;
    using OptionalValue = Sawyer::Optional<Value>;

private:
    mutable SAWYER_THREAD_TRAITS::Mutex mutex_; // protects all following data members
    Sawyer::Optional<Key> key_;
    OptionalValue value_;

public:
    void set(const Key &key, const Value &value) {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        key_ = key;
        value_ = value;
    }

    bool maybeSet(const Key &key, const Value &value) {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        if (key_ && *key_ == key) {
            return false;
        } else {
            key_ = key;
            value_ = value;
            return true;
        }
    }

    OptionalValue get(const Key &key) const {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        if (key_ && *key_ == key)
            return value_;
        return {};
    }

    OptionalValue take(const Key &key) {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        if (key_ && *key_ == key) {
            auto retval = value_;
            key_.reset();
            value_.reset();
            return retval;
        }
        return {};
    }

    void reset() {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        key_.reset();
        value_.reset();
    }

    bool exists(const Key &key) const {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        return key_ && *key_ == key;
    }
};

// Same as above, but no keys
template<class V>
class CachedItem<V, void> {
public:
    using Value = V;
    using Key = void;
    using OptionalValue = Sawyer::Optional<Value>;

private:
    mutable SAWYER_THREAD_TRAITS::Mutex mutex_;
    OptionalValue value_;

public:
    void set(const Value &value) {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        value_ = value;
    }

    bool maybeSet(const Value &value) {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        if (value_) {
            return false;
        } else {
            value_ = value;
            return true;
        }
    }

    OptionalValue get() const {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        return value_;
    }

    OptionalValue take() const {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        auto retval = value_;
        value_ = Sawyer::Nothing();
        return retval;
    }

    void reset() {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        value_.reset();
    }

    bool exists() const {
        SAWYER_THREAD_TRAITS::LockGuard lock(mutex_);
        return value_;
    }
};

template<class Cache>
class Borrowed {
    Cache &cache_;
    typename Cache::Key key_;
    typename Cache::OptionalValue value_;
    bool canceled_;

public:
    Borrowed(Cache &cache, typename Cache::Key key)
        : cache_(cache), key_(key), value_(cache_.take(key)), canceled_(false) {}

    Borrowed(Cache &cache, typename Cache::Key key, const typename Cache::Value &value)
        : cache_(cache), key_(key), value_(value), canceled_(false) {}

    ~Borrowed() {
        rtn();
    }

    void cancel() {
        canceled_ = true;
    }

    void rtn() {
        if (!canceled_ && value_ && !cache_.maybeSet(key_, *value_)) {
            //mlog[Diagnostics::WARN] <<"cache borrow contention detected: key=" <<StringUtility::addrToString(key_) <<"\n";
        }
        value_.reset();
    }

    const typename Cache::OptionalValue& get() const {
        return value_;
    }
};

template<class Cache>
Borrowed<Cache> borrow(Cache &cache, typename Cache::Key key) {
    return Borrowed<Cache>(cache, key);
}

template<class Cache>
Borrowed<Cache> borrow(Cache &cache, typename Cache::Key key, const typename Cache::Value &value) {
    return Borrowed<Cache>(cache, key, value);
}

// Information stored at each vertex of the CFG.

class InsnInfo: boost::noncopyable {
public:
    using Ptr = std::shared_ptr<InsnInfo>;      
    using List = std::vector<Ptr>;              
    struct Cached: boost::noncopyable {
        CachedItem<bool, uint64_t> isFunctionCall; 
        CachedItem<Semantics::RiscOperatorsPtr, uint64_t> semantics;  
    };

private:
    InstructionPtr ast_;                        // possibly evicted AST rooted at an SgAsmInstruction node.
    const rose_addr_t va_;                      // starting address accessible even for evicted instruction
    Cached cached_;                             // cached items computed elsewhere

    mutable SAWYER_THREAD_TRAITS::Mutex mutex_; // protects all following data members
    size_t size_;                               // instruction size in bytes even for evicted instruction
    bool wasDecoded_;                           // has any DecodeInstruction work completed?
    FunctionReasons functionReasons_;          // is this insn the entry point of a function?

    // Cached items associated with basic blocks. The key is the hash of the basic block instruction addresses.

public:
    explicit InsnInfo(const InstructionPtr &insn)
        : ast_(insn), va_(insn->get_address()), size_(insn->get_size()), wasDecoded_(false) {
        ASSERT_not_null(insn);
        ASSERT_require(size_ > 0);
    }

    explicit InsnInfo(rose_addr_t va)
        : va_(va), size_(0), wasDecoded_(false) {}

    rose_addr_t address() const {
        return va_;
    }

    Sawyer::Optional<size_t> size() const;

    Sawyer::Optional<AddressInterval> hull() const;

    FunctionReasons functionReasons() const;
    void functionReasons(FunctionReasons);
    void insertFunctionReasons(FunctionReasons);
    void eraseFunctionReasons(FunctionReasons);
    const Cached& cached() const { return cached_; }
    Cached& cached() { return cached_; }
    bool wasDecoded() const;
    void setDecoded();
    InstructionPtr ast() const;

    static bool addressOrder(const Ptr &a, const Ptr &b);

    static uint64_t hash(const List&);

private:
    friend class Partitioner;

    // Set the underlying instruction.
    //
    // This sets the underlying instruction if it isn't set already, and returns the underlying instruction. If the return value
    // is the same pointer as the argument then the instruction was set, otherwise it already had some other value. Also, the
    // @ref wasDecoded flag is set.  The instruction can be a valid instruction or a null pointer.
    //
    // Thread safety: This function is thread safe.
    InstructionPtr setAstMaybe(const InstructionPtr&);

};

// CFG edge information

class CfgEdge {
public:
    using EdgeTypes = BitFlags<EdgeType>;
    EdgeTypes types_;

public:
    CfgEdge()
        : types_(E_NORMAL) {}

    /*implicit*/ CfgEdge(EdgeType type)
        : types_(type) {}

    void merge(const CfgEdge &other) {
        types_.set(other.types_);
    }

    EdgeTypes types() const {
        return types_;
    }

    void types(EdgeTypes x) {
        types_.set(x);
    }

    friend std::ostream& operator<<(std::ostream&, const CfgEdge&);
};

// Control flow graph

// Lookup keys used internally for CFG vertices.
class InsnInfoKey {
    rose_addr_t va_;
public:
    /*implicit*/ InsnInfoKey(const InsnInfo::Ptr &insnInfo)
        : va_(insnInfo->address()) {}
    /*impilcit*/ InsnInfoKey(rose_addr_t va)
        : va_(va) {}

    // FIXME: we should be using an unordered_map for the index.
    bool operator<(const InsnInfoKey &other) const {
        return va_ < other.va_;
    }
};

using InsnCfg = Sawyer::Container::Graph<std::shared_ptr<InsnInfo>, CfgEdge, InsnInfoKey>;

// WorkItem -- base class for individual parallel units of work

class WorkItem {
public:
    enum class Priority {
        DiscoverInstruction = 0,
        NextUnusedVa = -100
    };

private:
    Partitioner &partitioner_;                  // partitioner that owns this work item
    Priority priority_;                         // primary sort key
    uint64_t sort_;                             // secondary sort key

protected:
    WorkItem(Partitioner &partitioner, Priority priority, uint64_t sort)
        : partitioner_(partitioner), priority_(priority), sort_(sort) {}

    virtual ~WorkItem() {}

    Partitioner& partitioner() const {
        return partitioner_;
    }

    Priority priority() const {
        return priority_;
    }

public:
    bool operator<(const WorkItem&) const;

public:
    virtual void run() = 0;

    virtual std::string title() const = 0;
};

// DecodeInstruction -- a unit of work

class DecodeInstruction: public WorkItem {
protected:
    rose_addr_t insnVa;                     // starting address of the instruction

public:
    DecodeInstruction(Partitioner &partitioner, rose_addr_t va)
        : WorkItem(partitioner, WorkItem::Priority::DiscoverInstruction,
#if 0
                   // Process instructions in order of their address. This might be better for caching other data because once
                   // we evict it we might never need it again. On the other hand, it seems to increase contention for things
                   // like basic block semantics since it's more likely that two threads will be working near each other in the
                   // specimen virtual memory.
                   va
#else
                   // Process instructions in random order. This seems to reduce contention for things like basic block semantics and
                   // improve performance by about 10% at this early stage of development.
                   Sawyer::fastRandomIndex(UNLIMITED)
#endif
                   ), insnVa(va) {}


    std::string title() const override {
        return "decode insn " + StringUtility::addrToString(insnVa);
    }

    void run() override;
};

// NextUnusedRegion -- a unit of work

class NextUnusedRegion: public WorkItem {
protected:
    AddressInterval where;                      // what part of the address space to search

public:
    NextUnusedRegion(Partitioner &partitioner, const AddressInterval &where)
        : WorkItem(partitioner, WorkItem::Priority::NextUnusedVa, where ? where.least() : uint64_t(0)), where(where) {}

    std::string title() const override {
        return "scan unused va within " + StringUtility::addrToString(where);
    }

    void run() override;
};

// Scheduler -- schedules units of work for parallel execution

class WorkItemSorter {
public:
    bool operator()(const std::shared_ptr<WorkItem>&, const std::shared_ptr<WorkItem>&) const;
};

class Scheduler final {
public:
    using Item = std::shared_ptr<WorkItem>; 
    using Container = std::vector<Item>;
    using Queue = std::priority_queue<Item, Container, WorkItemSorter>; 
private:
    mutable SAWYER_THREAD_TRAITS::Mutex mutex_; // protects all following data members
    Queue queue_;

public:
    void insert(const Item&);

    bool isEmpty() const;

    Item next();

    void reportStatus(std::ostream&) const;
};

// Worker -- top level function for worker threads

class Worker {
public:
    void operator()(std::shared_ptr<WorkItem> work, Scheduler&) {
        work->run();
    }
};

// Address usage map

using Aum = Sawyer::Container::IntervalSetMap<AddressInterval, AddressSet>;

// Partitioner

class Partitioner: public Sawyer::Attribute::Storage<>, public Sawyer::SharedObject, boost::noncopyable {
    Settings settings_;
    Scheduler scheduler_;
    std::shared_ptr<InstructionCache> insnCache_;
    Progress::Ptr progress_;                            // reports percent of memory disassembled
    rose_addr_t nExeVas_;                               // number of executable addresses in memory map

    mutable SAWYER_THREAD_TRAITS::Mutex mutex_;         // protects all following data members
    bool isRunning_;                                    // true while "run" is called
    InsnCfg insnCfg_;                                   // the global control flow graph
    Aum aum_;                                           // address usage map

public:
    //--------------------------------------------------------------------------------------------------------------------
    // Constructors and initialization
    //--------------------------------------------------------------------------------------------------------------------

    Partitioner(const MemoryMap::Ptr &memory, const Disassembler::BasePtr &decoder, const Settings &settings = Settings());

    const Settings& settings() const { return settings_; }
    Settings& settings() { return settings_; }
    static Accuracy choose(Accuracy a, Accuracy b);

    //--------------------------------------------------------------------------------------------------------------------
    // Status
    //--------------------------------------------------------------------------------------------------------------------
public:
    Progress::Ptr progress() const;

    void statusReports();

    //--------------------------------------------------------------------------------------------------------------------
    // Functions related to memory
    //--------------------------------------------------------------------------------------------------------------------
public:

    MemoryMap::Ptr memoryMap() const;

    size_t nDecodedAddresses() const;

    AddressIntervalSet unusedExecutableVas(AddressInterval where) const;

    //--------------------------------------------------------------------------------------------------------------------
    // Functions related to instructions
    //--------------------------------------------------------------------------------------------------------------------
public:

    InstructionPtr decodeInstruction(rose_addr_t);

    InsnInfo::Ptr makeInstruction(rose_addr_t);
    InsnInfo::Ptr makeInstruction(rose_addr_t, const InstructionPtr&);
    InsnInfo::Ptr existingInstruction(rose_addr_t);

    InstructionPtr existingInstructionAst(rose_addr_t);

    InstructionCache& instructionCache() const;

    struct LockInCache {
        std::vector<LockedInstruction> locks;                   
        std::vector<SgAsmInstruction*> insns;                   
    };

    LockInCache lockInCache(const InsnInfo::List&);

    struct CreateLinkedCfgVertices {
        bool createdSource = false;                             
        bool createdTarget = false;                             
        bool createdEdge = false;                               
        InsnInfo::Ptr source;                                   
        InsnInfo::Ptr target;                                   
    };

    CreateLinkedCfgVertices createLinkedCfgVertices(rose_addr_t sourceVa, rose_addr_t targetVa, const CfgEdge&);

    void scheduleDecodeInstruction(rose_addr_t insnVa);

    void scheduleNextUnusedRegion(const AddressInterval &where);

    //--------------------------------------------------------------------------------------------------------------------
    // Basic blocks
    //--------------------------------------------------------------------------------------------------------------------
public:
    InsnInfo::List basicBlockEndingAt(rose_addr_t, size_t maxInsns = UNLIMITED) const;

    InsnInfo::List basicBlockContaining(rose_addr_t) const;

    Borrowed<CachedItem<Semantics::RiscOperatorsPtr, uint64_t>> basicBlockSemantics(const InsnInfo::List&);

    std::vector<SymbolicExpression::Ptr> computeSuccessors(const InsnInfo::List&, Accuracy accuracy);
    std::vector<SymbolicExpression::Ptr> computeSuccessors(rose_addr_t insnVa, Accuracy accuracy);
    AddressSet computedConcreteSuccessors(rose_addr_t insnVa, Accuracy accuracy);

    bool isFunctionCall(rose_addr_t insnVa, Accuracy accuracy);

private:
    // Reads the instruction pointer register from the sepcified state, which must be non-null and have a non-null current
    // state, and splits it into a list of symbolic expressions. The most common return value is a single concrete address
    // which occurs for almost all non-branching instructions (such as ADD, NOP, etc.). The second most common case is two
    // concrete addresses from a conditional branch instruction. Another common case is that one of the successors will be a
    // free variable to represent that we don't know the successors (as usually happens with x86 RET and other computed
    // branches).
    std::vector<SymbolicExpression::Ptr> splitSuccessors(const InstructionSemantics::BaseSemantics::RiscOperatorsPtr &ops);

    //--------------------------------------------------------------------------------------------------------------------
    // Functions to run things.
    //--------------------------------------------------------------------------------------------------------------------
public:

    void run(size_t maxWorkers);

    bool isRunning() const;

private:
    // Used internally to change the isRunning value. This function is thread safe.
    void isRunning(bool b);

    //--------------------------------------------------------------------------------------------------------------------
    // Post processing functions. The following functons can only be called when the partitioner is not running (i.e.,
    // only when isRunning returns false).
    //--------------------------------------------------------------------------------------------------------------------
public:

    const InsnCfg& insnCfg() const;
    InsnCfg& insnCfg();
    void printInsnCfg(std::ostream&) const;

    void dumpInsnCfg(std::ostream&, const Rose::BinaryAnalysis::Partitioner2::Partitioner&) const;

    std::vector<InsnInfo::List> allBasicBlocks() const;

    std::map<rose_addr_t /*insn*/, rose_addr_t /*bb*/> calculateInsnToBbMap() const;

    static bool addressOrder(const InsnInfo::List &a, const InsnInfo::List &b);

    std::map<rose_addr_t /*funcVa*/, AddressSet /*insnVas*/> assignFunctions();

    void transferResults(const Rose::BinaryAnalysis::Partitioner2::PartitionerPtr &out);

    rose_addr_t remap();
};


} // namespace
} // namespace
} // namespace
} // namespace
} // namespace
#endif
#endif