Method partition_with_tree

crates/kernel/src/shard.rs:231–258

Like [`Self::partition`], but also returns the **bisection tree** — the binary tree of min-cut splits whose leaves are the shard ids. Reusing this as the proof-aggregation tree (rather than an arbitrary flat fold) discharges each cross-shard assumption at the lowest common ancestor of the shards that share it; the per-bisection min-cut means sibling subtrees have the thinnest possible interface, s

(
    &self,
    num_shards: usize,
    epsilon: f64,
  )

  /// shards that share it; the per-bisection min-cut means sibling subtrees have
  /// the thinnest possible interface, so most discharge happens low in the tree.
  pub fn partition_with_tree(
    &self,
    num_shards: usize,
    epsilon: f64,
  ) -> (Vec<u32>, AggNode) {
    let n = self.num_vertices();
    if num_shards <= 1 {
      return (vec![0u32; n], AggNode::Leaf(0)); // everything in shard 0
    }
    // Cap each block's balance weight at the ideal per-shard heartbeats
    // (total / num_shards). This keeps balancing heartbeat-aware while a balanced
    // split is achievable (few shards), but stops a single oversized *atomic*
    // block from skewing every split toward it once there are many shards. The
    // resulting heartbeat imbalance is accepted; the goal is `num_shards`
    // **non-empty** parallel shards with minimal cross-shard ingress.
    let total_hb: u128 = self.vweight.iter().map(|&w| u128::from(w)).sum();
    let cap =
      (total_hb / num_shards as u128).max(1).min(u128::from(u64::MAX)) as u64;
    let sub = SubHyper::full(self, cap);
    // Atomic assignment buffer: independent subtrees own disjoint block sets, so
    // they can be partitioned on separate threads without their writes aliasing.
    let shard_of: Vec<AtomicU32> = (0..n).map(|_| AtomicU32::new(0)).collect();
    let prog = PartitionProgress::new(num_shards);
    let tree = rec_bisect(&sub, num_shards, 0, epsilon, &shard_of, &prog);
    prog.done();
    let assignment = shard_of.into_iter().map(AtomicU32::into_inner).collect();
    (assignment, tree)
  }
}

// ============================================================================