graph.cpp

#include "graph.h"

Node *Graph::addNode(std::unique_ptr<Node> node) {
  Node *raw = node.get();
  nodes_.push_back(std::move(node));
  auto out = raw->output_names();
  for (auto &tensor : out) {
    if (!value_map_.emplace(tensor, raw).second) {
      throw std::runtime_error("Duplicate producer for value: " + tensor);
    }
  }
  return raw;
}

void Graph::constructEdges() {
  for (auto &node : nodes_) {
    for (auto &tensor : node->input_names()) {
      auto it = value_map_.find(tensor);
      if (it != value_map_.end()) {
        auto parent = it->second;
        node->addParent(parent);
        parent->addChild(node.get());
      }
    }
  }
  built = true;
}

void Graph::topoSort() {
  if (!built) {
    constructEdges();
  }

  std::unordered_map<Node *, int> indeg;
  std::queue<Node *> q;
  indeg.reserve(nodes_.size());

  for (auto &node : nodes_) {
    Node *n = node.get();
    indeg[n] = n->parents().size();
    if (indeg[n] == 0) {
      q.push(n);
    }
  }

  sorted_nodes_.clear();
  sorted_nodes_.reserve(nodes_.size());
  while (!q.empty()) {
    Node *n = q.front();
    q.pop();
    sorted_nodes_.push_back(n);
    for (auto m : n->children()) {
      if (--indeg[m] == 0) {
        q.push(m);
      }
    }
  }
  if (sorted_nodes_.size() != nodes_.size()) {
    throw std::runtime_error(
        "Topological sort failed: graph has a cycle or bad edges.");
  }
  for (Node *n : sorted_nodes_)
    n->pending_dependencies().store(n->parents().size());
}