serializer.cpp

// ============================================================================
// serializer.cpp
//
// Graph-only execution-graph builder + time-aware ordering edges + DOT export.
//
// Build:
//   g++ -std=c++20 -O2 -Wall -Wextra serializer.cpp -o serializer
//
// Run:
//   ./serializer
//   ./serializer > serializer.dot
//   dot -Tpng serializer.dot -o serializer.png
// ============================================================================

#include <algorithm>
#include <cstdint>
#include <exception>
#include <iostream>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

using NodeId = uint32_t;

// ============================================================================
// TimeExpr: time expressions for pulse start times
// ============================================================================
struct TimeExpr {
  struct Abs   { uint64_t t; };
  struct Const { int64_t  dt; };
  struct EndOf { NodeId ref; int64_t off; };

  struct Add {
    std::shared_ptr<TimeExpr> a;
    std::shared_ptr<TimeExpr> b;
  };

  using TimeExprType = std::variant<Abs, Const, EndOf, Add>;
  TimeExprType v;

  static TimeExpr AbsT(uint64_t t) { return TimeExpr{Abs{t}}; }
  static TimeExpr ConstT(int64_t dt){ return TimeExpr{Const{dt}}; }
  static TimeExpr EndOfT(NodeId ref, int64_t off=0){ return 
    TimeExpr{EndOf{ref, off}}; 
  }
  static TimeExpr AddT(TimeExpr a, TimeExpr b) {
    TimeExpr e;
    e.v = Add{std::make_shared<TimeExpr>(std::move(a)),
              std::make_shared<TimeExpr>(std::move(b))};
    return e;
  }
};

static std::string to_string(const TimeExpr& t) {
  struct V {
    std::string operator()(const TimeExpr::Abs& x)   const { 
      return "Abs(" + std::to_string(x.t) + ")"; 
    }
    std::string operator()(const TimeExpr::Const& x) const { 
      return "Const(" + std::to_string(x.dt) + ")"; 
    }
    std::string operator()(const TimeExpr::EndOf& x) const {
      return "EndOf(" + std::to_string(x.ref) + ")+" + std::to_string(x.off);
    }
    std::string operator()(const TimeExpr::Add& x) const {
      return "Add(" + to_string(*x.a) + "," + to_string(*x.b) + ")";
    }
  };
  return std::visit(V{}, t.v);
}

// ============================================================================
// IR Node: pulse specification and timing information
// ============================================================================
struct PulseSpec {
  std::string name;
  uint8_t channel;
  uint64_t dur;
  float amp = 0.0f;
  float phase = 0.0f;
};

struct Node {
  NodeId id{};
  PulseSpec pulse{};
  TimeExpr time{TimeExpr::AbsT(0)};

  /// Explicit user-intended dependencies
  std::vector<NodeId> deps{};

  /// Injected: per-channel time ordering
  std::vector<NodeId> time_deps{};

  /// Solved: start and end times
  std::optional<uint64_t> t_start{};
  std::optional<uint64_t> t_end{};
};

// ============================================================================
// ExecutionGraph: the graph of pulses and their dependencies
// ============================================================================
struct ExecutionGraph {
  std::vector<Node> nodes;
  std::unordered_map<NodeId, size_t> idx;

  void reindex() {
    idx.clear();
    idx.reserve(nodes.size());
    for (size_t i = 0; i < nodes.size(); ++i) idx[nodes[i].id] = i;
  }

  Node& get(NodeId id) { return nodes.at(idx.at(id)); }
  const Node& get(NodeId id) const { return nodes.at(idx.at(id)); }
};

static inline void for_each_dep(const Node& n, const auto& f) {
  for (NodeId d : n.deps)      f(d);
  for (NodeId d : n.time_deps) f(d);
}

// ============================================================================
// Topo Sort: topological sort of the execution graph
// This is a vanilla Kahn's algorithm for topological sort.
// ============================================================================
static std::vector<NodeId> topo_sort(const ExecutionGraph& g) {
  std::unordered_map<NodeId, int> indeg;
  std::unordered_map<NodeId, std::vector<NodeId>> out;

  indeg.reserve(g.nodes.size());
  out.reserve(g.nodes.size());

  for (auto& n : g.nodes) indeg[n.id] = 0;

  for (auto& n : g.nodes) {
    for_each_dep(n, [&](NodeId d) {
      out[d].push_back(n.id);
      indeg[n.id] += 1;
    });
  }

  std::vector<NodeId> q;
  q.reserve(g.nodes.size());
  for (auto& [id, deg] : indeg) if (deg == 0) q.push_back(id);

  std::vector<NodeId> order;
  order.reserve(g.nodes.size());

  while (!q.empty()) {
    NodeId x = q.back();
    q.pop_back();
    order.push_back(x);
    for (NodeId y : out[x]) {
      if (--indeg[y] == 0) q.push_back(y);
    }
  }

  if (order.size() != g.nodes.size())
    throw std::runtime_error("Cycle detected in dependency graph.");

  return order;
}

// ============================================================================
// Time Solve: evaluate time expressions and solve for start and end times
// ============================================================================
static std::optional<uint64_t> eval_timeexpr(const TimeExpr& t, 
                                             const ExecutionGraph& g) {
  struct V {
    const ExecutionGraph& g;

    std::optional<uint64_t> operator()(const TimeExpr::Abs& x) const { 
      return x.t; 
    }
    std::optional<uint64_t> operator()(const TimeExpr::Const&) const { 
      return std::nullopt; 
    }

    std::optional<uint64_t> operator()(const TimeExpr::EndOf& x) const {
      const Node& ref = g.get(x.ref);
      if (!ref.t_end) return std::nullopt;
      int64_t t64 = (int64_t)(*ref.t_end) + x.off;
      if (t64 < 0) throw std::runtime_error(
        "Negative time from EndOf(): " + std::to_string(t64)
      );
      return (uint64_t)t64;
    }

    std::optional<uint64_t> operator()(const TimeExpr::Add& x) const {
      auto a = eval_timeexpr(*x.a, g);
      auto b = eval_timeexpr(*x.b, g);

      auto const_val = [&](const TimeExpr& e)->std::optional<int64_t>{
        if (auto* c = std::get_if<TimeExpr::Const>(&e.v)) return c->dt;
        return std::nullopt;
      };

      if (a && b) return *a + *b;

      if (a && !b) {
        auto cb = const_val(*x.b);
        if (!cb) return std::nullopt;
        int64_t t64 = (int64_t)(*a) + *cb;
        if (t64 < 0) throw std::runtime_error(
          "Negative time from Add(): " + std::to_string(t64));
        return (uint64_t)t64;
      }
      if (!a && b) {
        auto ca = const_val(*x.a);
        if (!ca) return std::nullopt;
        int64_t t64 = (int64_t)(*b) + *ca;
        if (t64 < 0) throw std::runtime_error(
          "Negative time from Add(): " + std::to_string(t64));
        return (uint64_t)t64;
      }
      return std::nullopt;
    }
  };

  return std::visit(V{g}, t.v);
}

static void solve_times(ExecutionGraph& g, const std::vector<NodeId>& topo) {
  for (NodeId id : topo) {
    Node& n = g.get(id);

    auto t_expr = eval_timeexpr(n.time, g);
    if (!t_expr) {
      throw std::runtime_error(
        "Cannot evaluate time expr for node " + std::to_string(n.id) +
        ": " + to_string(n.time)
      );
    }

    uint64_t t_deps = 0;
    for (NodeId d : n.deps) {
      const Node& dep = g.get(d);
      if (!dep.t_end) throw std::runtime_error(
        "Internal error: dep end not solved.");
      t_deps = std::max(t_deps, *dep.t_end);
    }

    uint64_t t_start = std::max(*t_expr, t_deps);
    uint64_t t_end   = t_start + n.pulse.dur;

    n.t_start = t_start;
    n.t_end   = t_end;
  }
}

// ============================================================================
// AddChannelTimeDependencies: inject per-channel time ordering
// edges between pulses on the same channel.
// ============================================================================
static void add_channel_time_dependencies(ExecutionGraph& g) {
  for (auto& n : g.nodes) n.time_deps.clear();

  std::unordered_map<uint8_t, std::vector<NodeId>> by_ch;
  by_ch.reserve(g.nodes.size());

  for (auto& n : g.nodes) {
    if (!n.t_start || !n.t_end)
      throw std::runtime_error(
        "Internal error: add_channel_time_dependencies before solve_times()"
      );
    by_ch[n.pulse.channel].push_back(n.id);
  }

  for (auto& [ch, ids] : by_ch) {
    std::sort(ids.begin(), ids.end(), [&](NodeId a, NodeId b) {
      const Node& A = g.get(a);
      const Node& B = g.get(b);
      if (*A.t_start != *B.t_start) return *A.t_start < *B.t_start;
      return A.id < B.id;
    });

    for (size_t i = 1; i < ids.size(); ++i) {
      g.get(ids[i]).time_deps.push_back(ids[i - 1]);
    }

    (void)ch;
  }
}

// ============================================================================
// dump_graph: dump the execution graph to a text file
// ============================================================================
static void dump_graph(const ExecutionGraph& g) {
  std::cout << "ExecutionGraph (" << g.nodes.size() << " nodes)\n";
  for (auto& n : g.nodes) {
    std::cout << "  Node " << n.id
              << " " << n.pulse.name
              << " ch=" << int(n.pulse.channel)
              << " dur=" << n.pulse.dur
              << " time=" << to_string(n.time)
              << " start=" << (n.t_start ? std::to_string(*n.t_start) : "?")
              << " end=" << (n.t_end ? std::to_string(*n.t_end) : "?")
              << " deps=[";
    for (size_t i = 0; i < n.deps.size(); ++i) {
      std::cout << n.deps[i] << (i + 1 < n.deps.size() ? "," : "");
    }
    std::cout << "] time_deps=[";
    for (size_t i = 0; i < n.time_deps.size(); ++i) {
      std::cout << n.time_deps[i] << (i + 1 < n.time_deps.size() ? "," : "");
    }
    std::cout << "]\n";
  }
  std::cout << "\n";
}

static void emit_dot(const ExecutionGraph& g, std::ostream& os) {
  os << "digraph ExecutionGraph {\n";
  os << "  rankdir=LR;\n";
  os << "  node [shape=box, fontsize=10];\n";
  os << "  edge [fontsize=9];\n\n";

  for (const auto& n : g.nodes) {
    std::string label = n.pulse.name;
    label += "\\n";
    label += "id=" + std::to_string(n.id);
    label += " ch=" + std::to_string(int(n.pulse.channel));
    label += " dur=" + std::to_string(n.pulse.dur);

    if (n.t_start && n.t_end) {
      label += "\\n[" + std::to_string(*n.t_start) + "," + std::to_string(*n.t_end) + ")";
    } else {
      label += "\\n[?,?)";
    }

    os << "  n" << n.id << " [label=\"" << label << "\"];\n";
  }
  os << "\n";

  // explicit deps: solid
  os << "  // deps (solid)\n";
  for (const auto& n : g.nodes) {
    for (NodeId d : n.deps) {
      os << "  n" << d << " -> n" << n.id << " [style=solid,label=\"dep\"];\n";
    }
  }

  os << "\n";

  // time ordering deps: dashed
  os << "  // time_deps (dashed)\n";
  for (const auto& n : g.nodes) {
    for (NodeId d : n.time_deps) {
      os << "  n" << d << " -> n" << n.id << " [style=dashed,label=\"time\"];\n";
    }
  }

  os << "}\n";
}

// ============================================================================
// ProgramBuilder: the API for building the execution graph
// ============================================================================
struct PulseHandle { NodeId id; };

class ProgramBuilder {
public:
  struct PulseBuilder {
    ProgramBuilder& pb;
    PulseHandle h;

    PulseHandle at(uint64_t t_abs) { 
      pb.set_time(h, TimeExpr::AbsT(t_abs)); 
      return h; 
    }
    PulseHandle at(TimeExpr te) { 
      pb.set_time(h, std::move(te));        
      return h; 
    }

    PulseHandle after(PulseHandle dep, int64_t offset) {
      pb.add_dep(h, dep);
      pb.set_time(h, TimeExpr::EndOfT(dep.id, offset));
      return h;
    }

    PulseHandle depends_on(PulseHandle dep) { pb.add_dep(h, dep); return h; }
  };

  PulseBuilder pulse(std::string name, uint8_t ch, uint64_t dur, float amp=0.0f, 
                     float phase=0.0f) {
    Node n;
    n.id = next_id_++;
    n.pulse = PulseSpec{std::move(name), ch, dur, amp, phase};
    n.time = TimeExpr::AbsT(0);
    g_.nodes.push_back(std::move(n));
    g_.reindex();
    return PulseBuilder{*this, PulseHandle{g_.nodes.back().id}};
  }

  TimeExpr end_of(PulseHandle h) { return TimeExpr::EndOfT(h.id, 0); }
  TimeExpr add(TimeExpr a, int64_t dt) { 
    return TimeExpr::AddT(std::move(a), TimeExpr::ConstT(dt)); 
  }

  void compile_graph() {
    g_.reindex();
    auto order = topo_sort(g_);
    solve_times(g_, order);
    add_channel_time_dependencies(g_);

    g_.reindex();
    /// (void)topo_sort(g_);
  }

  const ExecutionGraph& graph() const { return g_; }

private:
  void set_time(PulseHandle h, TimeExpr te) { 
    g_.get(h.id).time = std::move(te); 
  }
  void add_dep(PulseHandle h, PulseHandle dep) { 
    g_.get(h.id).deps.push_back(dep.id); 
  }

  ExecutionGraph g_;
  NodeId next_id_ = 1;
};

// ============================================================================
// Subroutine Example: a subroutine that emits pulses and dependencies
// ============================================================================
struct InitHandles { PulseHandle p0, p1, cond; };

InitHandles qubit_init(ProgramBuilder& pb, uint64_t t0, PulseHandle meas) {
  InitHandles out;
  out.p0   = pb.pulse("init_p0",   1, 40).at(t0);
  out.p1   = pb.pulse("init_p1",   1, 40).at(t0 + 30);
  out.cond = pb.pulse("init_cond", 1, 40).after(meas, 200);
  return out;
}

int main() {
  try {
    ProgramBuilder pb;
    auto meas = pb.pulse("meas", 0, 25).at(100);
    auto h1 = qubit_init(pb, 200, meas);
    auto h2 = qubit_init(pb, 400, meas);
    pb.pulse("post_init", 2, 30).after(h1.cond, 50);

    pb.compile_graph();
    dump_graph(pb.graph());

    emit_dot(pb.graph(), std::cerr);

    return 0;
  } catch (const std::exception& e) {
    std::cerr << "ERROR: " << e.what() << "\n";
    return 1;
  }
}