Graph.h

#pragma once

#include <momemta/config.h>
#include <momemta/Configuration.h>
#include <momemta/Module.h>

#include <ExecutionPath.h>

#include <map>
#include <string>
#include <vector>

#include <boost/config.hpp>
#include <boost/functional/hash.hpp>
#include <boost/graph/adjacency_list.hpp>

#ifdef DEBUG_TIMING
#include <chrono>
#endif

namespace momemta {

// Graph definitions

struct Vertex {
    uint32_t id;
    std::string name; // Unique name of the module
    std::string type; // Module type
    momemta::ModuleList::value_type def;
    Configuration::ModuleDecl decl;
};

struct Edge {
    std::string description;
    bool virt; 
    InputTag tag; 
};

typedef boost::adjacency_list<boost::listS, boost::listS, boost::bidirectionalS, Vertex, Edge> Graph;
typedef boost::graph_traits<Graph>::vertex_descriptor vertex_t;
typedef boost::graph_traits<Graph>::edge_descriptor edge_t;

class ComputationGraph {
public:
    void initialize(PoolPtr pool);

    // Interface to module methods
    void configure();
    void beginIntegration();
    Module::Status execute();
    void endIntegration();
    void finish();

#ifdef DEBUG_TIMING
    void logTimings() const;
#endif

    void setNDimensions(size_t n);

    size_t getNDimensions() const;

    void addDecl(const boost::uuids::uuid& path, const Configuration::ModuleDecl& decl);
    const std::vector<boost::uuids::uuid>& getPaths() const;
    const std::vector<Configuration::ModuleDecl>& getDecls(const boost::uuids::uuid& path) const;

private:
    std::vector<boost::uuids::uuid> sorted_execution_paths;
    std::unordered_map<
            boost::uuids::uuid,
            std::vector<Configuration::ModuleDecl>,
            boost::hash<boost::uuids::uuid>
    > module_decls;

    std::vector<ModulePtr> modules;

    size_t n_dimensions; 

#ifdef DEBUG_TIMING
    std::unordered_map<Module *, std::chrono::high_resolution_clock::duration> module_timings;
#endif
};

class ComputationGraphBuilder {
public:

    ComputationGraphBuilder(
            const momemta::ModuleList& available_modules,
            const Configuration& configuration
    );

    std::shared_ptr<ComputationGraph> build();

    void exportGraph(const std::string& output) const;

private:
    void prune_graph();
    void sort_graph();
    void validate();

    const momemta::ModuleList& available_modules;
    const Configuration& configuration;

    Graph g;
    std::unordered_map<std::string, vertex_t> vertices;

    std::list<vertex_t> sorted_vertices; // filled only if sort_graph is called

    bool graph_exportable = false;
};

}