longfellow 0.1.0

data/vendor/longfellow-zk/lib/circuits/compiler/schedule.h

@@ -0,0 +1,435 @@
+// Copyright 2026 Google LLC.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef PRIVACY_PROOFS_ZK_LIB_CIRCUITS_COMPILER_SCHEDULE_H_
+#define PRIVACY_PROOFS_ZK_LIB_CIRCUITS_COMPILER_SCHEDULE_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <memory>
+#include <vector>
+
+#include "algebra/compare.h"
+#include "arrays/affine.h"
+#include "circuits/compiler/node.h"
+#include "sumcheck/circuit.h"
+#include "sumcheck/equad.h"
+#include "sumcheck/quad.h"
+#include "sumcheck/quad_builder.h"
+#include "util/ceildiv.h"
+#include "util/panic.h"
+
+namespace proofs {
+template <class Field>
+class Scheduler {
+  using Elt = typename Field::Elt;
+  using nodeinfo = NodeInfoF<Field>;
+  using node = NodeF<Field>;
+  using size_t_for_storage = term::size_t_for_storage;
+  using quad_corner_t = typename Quad<Field>::quad_corner_t;
+
+  const Field& f_;
+  const std::vector<node>& nodes_;
+
+ public:
+  size_t nwires_;
+  size_t nquad_terms_;
+  size_t nwires_overhead_;
+
+  Scheduler(const std::vector<node>& nodes, const Field& f)
+      : f_(f),
+        nodes_(nodes),
+        nwires_(0),
+        nquad_terms_(0),
+        nwires_overhead_(0) {}
+
+  std::unique_ptr<Circuit<Field>> mkcircuit(const std::vector<Elt>& constants,
+                                            size_t depth_ub, size_t nc) {
+    std::unique_ptr<Circuit<Field>> c = std::make_unique<Circuit<Field>>();
+
+    // number of layers and copies
+    c->nl = depth_ub - 1;  // depth 0 = input nodes, not a "layer"
+    c->nc = nc;
+    c->logc = lg(nc);
+
+    auto lnodes = order_by_layer(constants, depth_ub);
+
+    // TODO [matteof 2025-03-12] ASSIGN_WIRE_IDS() renames LNODES in
+    // order to sort it and assign LNODES[].DESIRED_WIRE ID.  Then it
+    // throws away the renamed LNODES.  Then FILL_LAYERS() renames
+    // LNODES again in order to produce the final quad.  It would be
+    // better to produce the quad directly in ASSIGN_WIRE_IDS().  Punt
+    // for now, this is just a performance optimization of the
+    // compiler anyway.
+    //
+    assign_wire_ids(lnodes);
+    fill_layers(c.get(), depth_ub, lnodes);
+
+    return c;
+  }
+
+ private:
+  // per-layer representation of nodes and terms
+  struct lterm {
+    Elt k;
+    quad_corner_t lop0, lop1;
+  };
+  struct lnode {
+    quad_corner_t desired_wire_id;
+
+    // Copy wires are forced to be distinct from wires in the
+    // original dag, in order to avoid ambiguity in renaming.
+    //
+    // Copy wires are always of the form 1*op, which doesn't
+    // normally appear in the dag because the algebraic simplifier
+    // reduces it to op.  However, one can in theory create such
+    // a node by judicious use of linear().  Rather than
+    // trying to figure out which circuits one is not allowed
+    // to write, it seems simpler to just handle this case
+    // uniformly.
+    bool is_copy_wire;
+
+    std::vector<lterm> lterms;
+
+    lnode(quad_corner_t desired_wire_id, bool is_copy_wire,
+          const std::vector<lterm>& lterms)
+        : desired_wire_id(desired_wire_id),
+          is_copy_wire(is_copy_wire),
+          lterms(lterms) {}
+  };
+
+  quad_corner_t lop_of_op_at_depth(
+      const std::vector<std::vector<quad_corner_t>>& lop, size_t op,
+      size_t d) const {
+    const node& n = nodes_.at(op);
+    return lop.at(op).at(d - n.info.depth);
+  }
+
+  // Convert the DAG of nodes into a layered dag of lnodes.
+  std::vector<std::vector<lnode>> order_by_layer(
+      const std::vector<Elt>& constants, size_t depth_ub) {
+    // The source DAG is indexed by NODES_[OP].
+    // The destination dag uses a two-dimensional indexing
+    // scheme LNODES[D][LOP], where D is the depth.
+
+    // A single value NODES_[OP] may be replicated multiple times in
+    // LNODES.  The mapping is maintained in array LOPS such that
+    // LOPS[OP][D - D0] contains the LOP index of node OP at depth D.
+    // D0 is the depth at which NODES_[OP] is first computed, and
+    // there is no point in storing LOPS[OP] for D < D0.
+
+    std::vector<std::vector<lnode>> lnodes(depth_ub);
+    std::vector<std::vector<quad_corner_t>> lops(nodes_.size());
+
+    nwires_overhead_ = 0;
+
+    for (size_t op = 0; op < nodes_.size(); ++op) {
+      const auto& n = nodes_[op];
+      const nodeinfo& nfo = n.info;
+      if (nfo.is_needed && !n.zero()) {
+        size_t d = nfo.depth;
+
+        // Allocate the LOP at depth D
+        quad_corner_t lop = quad_corner_t(lnodes.at(d).size());
+        lops.at(op).push_back(lop);
+
+        // create a LOPS entry for depth D
+        /*scope*/ {
+          std::vector<lterm> lterms;
+          for (const auto& t : n.terms) {
+            lterm lt = {
+                .k = constants.at(t.ki),
+                .lop0 = lop_of_op_at_depth(lops, t.op0, d - 1),
+                .lop1 = lop_of_op_at_depth(lops, t.op1, d - 1),
+            };
+            lterms.push_back(lt);
+          }
+          lnodes.at(d).push_back(lnode(nfo.desired_wire_id(d, depth_ub),
+                                       /*is_copy_wire=*/false, lterms));
+        }
+
+        // create copy wires
+        for (d = nfo.depth + 1; d < nfo.max_needed_depth; ++d) {
+          quad_corner_t lop_dm1 = lop;
+
+          // allocate the LOP at depth D
+          lop = quad_corner_t(lnodes.at(d).size());
+          lops.at(op).push_back(lop);
+
+          std::vector<lterm> lterms;
+
+          // Insert a multiplication by one of the layer
+          // at the previous layer.
+          lterm lt = {
+              .k = f_.one(),
+              .lop0 = quad_corner_t(0),
+              .lop1 = lop_dm1,
+          };
+          lterms.push_back(lt);
+          lnodes.at(d).push_back(lnode(nfo.desired_wire_id(d, depth_ub),
+                                       /*is_copy_wire=*/true, lterms));
+          ++nwires_overhead_;
+        }  // for copy wires
+      }  // if needed
+    }  // for OP
+
+    return lnodes;
+  }
+
+  //------------------------------------------------------------
+  // canonical assignment of wire ids
+  //------------------------------------------------------------
+  //
+  // The canonicalization order is a matter of convention.
+  // We make some arbitrary choices that appear to interact
+  // better with ZSTD compression.  The label [ARBITRARY CHOICE]
+  // denotes all places in the code where this occurs.
+  //
+  class renamed_lterm {
+   public:
+    Elt k_;
+    quad_corner_t rlop0_, rlop1_;
+
+    // [ARBITRARY CHOICE] Consistent with corner::canonicalize() in
+    // sumcheck/quad.h
+    renamed_lterm(const Elt& k, quad_corner_t rlop0, quad_corner_t rlop1)
+        : k_(k),
+          rlop0_(std::min<quad_corner_t>(rlop0, rlop1)),
+          rlop1_(std::max<quad_corner_t>(rlop0, rlop1)) {}
+
+    bool operator==(const renamed_lterm& y) const {
+      return rlop0_ == y.rlop0_ && rlop1_ == y.rlop1_ && k_ == y.k_;
+    }
+
+    // canonical order
+    static bool compare(const renamed_lterm& a, const renamed_lterm& b,
+                        const Field& F) {
+      if (a.rlop0_ < b.rlop0_) return true;
+      if (a.rlop0_ > b.rlop0_) return false;
+      if (a.rlop1_ < b.rlop1_) return true;
+      if (a.rlop1_ > b.rlop1_) return false;
+      return elt_less_than(a.k_, b.k_, F);
+    }
+  };
+
+  class renamed_lnode {
+   public:
+    quad_corner_t desired_wire_id_;
+    quad_corner_t original_wire_index_;
+    bool is_copy_wire_;
+    std::vector<renamed_lterm> rlterms_;
+
+    renamed_lnode(quad_corner_t desired_wire_id,
+                  quad_corner_t original_wire_index, bool is_copy_wire,
+                  const std::vector<renamed_lterm>& rlterms)
+        : desired_wire_id_(desired_wire_id),
+          original_wire_index_(original_wire_index),
+          is_copy_wire_(is_copy_wire),
+          rlterms_(rlterms) {}
+
+    bool operator==(const renamed_lnode& y) const {
+      if (is_copy_wire_ != y.is_copy_wire_) return false;
+      if (rlterms_.size() != y.rlterms_.size()) return false;
+      size_t l = rlterms_.size();
+      for (size_t i = 0; i < l; ++i) {
+        if (!(rlterms_[i] == y.rlterms_[i])) return false;
+      }
+      return true;
+    }
+
+    // canonical order
+    static bool compare(const renamed_lnode& ra, const renamed_lnode& rb,
+                        const Field& F) {
+      // Defined before undefined.  This choice is mandated by the
+      // fact that the range of defined wire id's starts at 0.
+      if (ra.desired_wire_id_ != nodeinfo::kWireIdUndefined) {
+        if (rb.desired_wire_id_ != nodeinfo::kWireIdUndefined) {
+          return ra.desired_wire_id_ < rb.desired_wire_id_;
+        } else {
+          return true;
+        }
+      } else {
+        if (rb.desired_wire_id_ != nodeinfo::kWireIdUndefined) {
+          return false;
+        }
+        // else both undefined
+      }
+
+      // [ARBITRARY CHOICE] Lexicographic order on the reverse of the
+      // terms array.  This seems to compress much better than
+      // the normal lexicographic order.
+      for (size_t ia = ra.rlterms_.size(), ib = rb.rlterms_.size();
+           ia-- > 0 && ib-- > 0;) {
+        const renamed_lterm& rlta = ra.rlterms_[ia];
+        const renamed_lterm& rltb = rb.rlterms_[ib];
+        if (renamed_lterm::compare(rlta, rltb, F)) return true;
+        if (renamed_lterm::compare(rltb, rlta, F)) return false;
+      }
+
+      // [ARBITRARY CHOICE] If the common suffixes are the same, the
+      // shorter terms come first.
+      if (ra.rlterms_.size() < rb.rlterms_.size()) return true;
+      if (ra.rlterms_.size() > rb.rlterms_.size()) return false;
+
+      // Nodes that were in the original dag come first.
+      if (!ra.is_copy_wire_ && rb.is_copy_wire_) return true;
+      if (!rb.is_copy_wire_ && ra.is_copy_wire_) return false;
+
+      // equal, i.e., not less-than
+      return false;
+    }
+  };
+
+  template <class T>
+  bool uniq(const std::vector<T>& sorted) {
+    for (size_t i = 0; i + 1 < sorted.size(); ++i) {
+      if (sorted[i] == sorted[i + 1]) return false;
+    }
+    return true;
+  }
+
+  void assign_wire_ids(std::vector<std::vector<lnode>>& lnodes) {
+    // all inputs are expected to be defined already
+    assert_all_desired_wire_id_defined(lnodes.at(0));
+
+    for (size_t d = 1; d < lnodes.size(); ++d) {
+      const std::vector<lnode>& lnodes_at_dm1 = lnodes.at(d - 1);
+      const std::vector<lnode>& lnodes_at_d = lnodes.at(d);
+
+      // Create a renamed clone of LNODES_AT_D, in which all
+      // the LOP's are mapped to their desired wire id's
+      // at the previous layer.  We use different types
+      // to avoid any possibility of confusion.
+      std::vector<renamed_lnode> renamed_at_d;
+
+      quad_corner_t original_wire_index(0);
+      for (const lnode& ln : lnodes_at_d) {
+        std::vector<renamed_lterm> rlterms;
+
+        // rename all terms
+        rlterms.reserve(ln.lterms.size());
+        for (const lterm& lt : ln.lterms) {
+          rlterms.push_back(renamed_lterm(
+              lt.k,
+              lnodes_at_dm1.at(static_cast<size_t>(lt.lop0)).desired_wire_id,
+              lnodes_at_dm1.at(static_cast<size_t>(lt.lop1)).desired_wire_id));
+        }
+
+        // canonicalize the terms order
+        std::sort(rlterms.begin(), rlterms.end(),
+                  [&](const renamed_lterm& a, const renamed_lterm& b) {
+                    return renamed_lterm::compare(a, b, f_);
+                  });
+
+        // Terms must be unique, otherwise the canonicalization is
+        // ill-defined.  Uniqueness is guaranteed by the algebraic
+        // simplifier, but assert it for good measure.
+        check(uniq(rlterms), "rlterms not unique");
+
+        renamed_at_d.push_back(renamed_lnode(
+            ln.desired_wire_id, original_wire_index, ln.is_copy_wire, rlterms));
+        ++original_wire_index;
+      }
+
+      check(renamed_at_d.size() == lnodes_at_d.size(),
+            "renamed_at_d.size() == lnodes_at_d.size()");
+
+      std::sort(renamed_at_d.begin(), renamed_at_d.end(),
+                [&](const renamed_lnode& a, const renamed_lnode& b) {
+                  return renamed_lnode::compare(a, b, f_);
+                });
+
+      // Nodes must be unique, otherwise the canonicalization is
+      // ill-defined.
+      check(uniq(renamed_at_d), "renamed_at_d not unique");
+
+      quad_corner_t wid(0);
+      std::vector<lnode>& wlnodes_at_d = lnodes.at(d);
+
+      for (const renamed_lnode& ln : renamed_at_d) {
+        lnode& lnpi =
+            wlnodes_at_d.at(static_cast<size_t>(ln.original_wire_index_));
+        if (lnpi.desired_wire_id != nodeinfo::kWireIdUndefined) {
+          // We must have computed the same wire id
+          check(wid == lnpi.desired_wire_id, "wid == lnpi.desired_wire_id");
+        } else {
+          lnpi.desired_wire_id = wid;
+        }
+        wid++;
+      }
+    }
+  }
+
+  void assert_all_desired_wire_id_defined(const std::vector<lnode>& layer) {
+    for (const auto& ln : layer) {
+      check(ln.desired_wire_id != nodeinfo::kWireIdUndefined,
+            "ln.desired_wire_id != kWireIdUndefined");
+    }
+  }
+
+  void fill_layers(Circuit<Field>* c, size_t depth_ub,
+                   const std::vector<std::vector<lnode>>& lnodes) {
+    check(depth_ub == lnodes.size(), "depth_ub == lnodes.size()");
+
+    corner_t nv = corner_t(lnodes.at(depth_ub - 1).size());
+
+    nwires_ = nv;
+    c->nv = nv;
+    c->logv = lg(nv);
+
+    // d-- > 1 (not 0) because depth 0 denotes input nodes, not a layer.
+    // Sumcheck counts layers starting from the output, hence the loop
+    // counts downwards.
+    for (size_t d = depth_ub; d-- > 1;) {
+      corner_t nw =
+          corner_t(lnodes.at(d - 1).size());  // inputs[d] == outputs[d-1]
+      nwires_ += nw;
+      c->l.push_back(
+          Layer<Field>{.nw = nw,
+                       .logw = lg(nw),
+                       .quad = mkquad(lnodes.at(d), lnodes.at(d - 1))});
+    }
+  }
+
+  std::unique_ptr<const Quad<Field>> mkquad(
+      const std::vector<lnode>& lnodes0,  // wires at this layer
+      const std::vector<lnode>& lnodes1   // wires at the previous layer
+  ) {
+    size_t nterms0 = 0;
+    for (const auto& ln0 : lnodes0) {
+      nterms0 += ln0.lterms.size();
+    }
+    nquad_terms_ += nterms0;
+
+    auto S = std::make_unique<EQuad<Field>>(nterms0);
+    size_t i = 0;
+    for (const auto& ln0 : lnodes0) {
+      for (const auto& lt : ln0.lterms) {
+        S->ec_[i++] = typename EQuad<Field>::ecorner{
+            .g = ln0.desired_wire_id,
+            .h = {lnodes1.at(static_cast<size_t>(lt.lop0)).desired_wire_id,
+                  lnodes1.at(static_cast<size_t>(lt.lop1)).desired_wire_id},
+            .v = lt.k};
+      }
+    }
+    S->canonicalize(f_);
+    return QuadBuilder<Field>::compress(S.get(), f_);
+  }
+};
+
+}  // namespace proofs
+
+#endif  // PRIVACY_PROOFS_ZK_LIB_CIRCUITS_COMPILER_SCHEDULE_H_