Fix MultiPodBuilder external input accounting and scheduling

src/frontend/multi_pod/cost.rs

@@ -46,6 +46,9 @@ impl From<&CustomPredicateRef> for CustomPredicateId {
 ///   op kind.
 #[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct OperationCost {
+    /// True for a statement copied directly from an external input POD.
+    #[serde(default)]
+    pub is_external_opening: bool,
     /// `Contains`-family proofs on a shallow tree (depth <=
     /// `max_depth_small`). May occupy either a small or medium state slot.
     pub merkle_proofs_small: usize,
@@ -247,6 +250,9 @@ impl OperationCost {
                 NativeOperation::PublicKeyFromEntries => {
                     cost.public_key = 1;
                 }
+                NativeOperation::OpenInputStatement => {
+                    cost.is_external_opening = true;
+                }
                 // Zero-cost ops (no per-POD resource consumed).
                 NativeOperation::None
                 | NativeOperation::EqualFromEntries
@@ -259,8 +265,6 @@ impl OperationCost {
                 | NativeOperation::ProductFromEntries
                 | NativeOperation::MaxFromEntries
                 | NativeOperation::HashFromEntries
-                // Tracked separately by the partitioner.
-                | NativeOperation::OpenInputStatement
                 // Syntactic sugar variants (lowered before proving).
                 | NativeOperation::GtEqFromEntries
                 | NativeOperation::GtFromEntries
@@ -361,6 +365,16 @@ mod tests {
         assert_eq!(pk.signed_by, 0);
     }
 
+    #[test]
+    fn open_input_is_marked_for_early_scheduling() {
+        let cost = OperationCost::from_operation(
+            &native_op(NativeOperation::OpenInputStatement, OperationAux::None),
+            &Params::default(),
+        );
+
+        assert!(cost.is_external_opening);
+    }
+
     /// Contains on a tree at depth `<= max_depth_small` is small-eligible;
     /// at depth `> max_depth_small` it must use a medium slot.
     #[test]

src/frontend/multi_pod/deps.rs

@@ -4,16 +4,14 @@ use std::collections::HashMap;
 
 use crate::{
     frontend::{Operation, OperationArg},
-    middleware::{
-        Hash, InputPodOpenStatement, NativeOperation, OperationAux, OperationType, Statement,
-    },
+    middleware::{InputPodOpenStatement, NativeOperation, OperationAux, OperationType, Statement},
 };
 
 /// Reference to a statement sourced from an external input POD.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct ExternalDependency {
-    /// Hash of the external POD containing `statement` in its public set.
-    pub pod_hash: Hash,
+    /// Position of the concrete source POD in the builder's input list.
+    pub pod_index: usize,
     /// The statement value itself.
     pub statement: Statement,
 }
@@ -36,14 +34,13 @@ pub struct DependencyGraph {
 }
 
 impl DependencyGraph {
-    /// Build the dependency graph from parallel `statements` and
-    /// `operations` arrays (where `operations[i]` produces `statements[i]`)
-    /// plus a `statement -> pod hash` map for recognising external
-    /// references.
+    /// Build dependencies from aligned statements and operations, where
+    /// `operations[i]` produces `statements[i]`, while retaining each external
+    /// reference's source POD.
     pub fn build(
         statements: &[Statement],
         operations: &[Operation],
-        external_pod_statements: &HashMap<Statement, Hash>,
+        external_pod_statements: &HashMap<Statement, usize>,
     ) -> Self {
         let mut statement_deps = Vec::with_capacity(statements.len());
 
@@ -85,9 +82,9 @@ impl DependencyGraph {
                         // it. Fall through to the external lookup below.
                     }
 
-                    if let Some(&pod_hash) = external_pod_statements.get(dep_stmt) {
+                    if let Some(&pod_index) = external_pod_statements.get(dep_stmt) {
                         deps.push(StatementSource::External(ExternalDependency {
-                            pod_hash,
+                            pod_index,
                             statement: dep_stmt.clone(),
                         }));
                     } else {
@@ -105,11 +102,11 @@ impl DependencyGraph {
             // accounts for the input-tree slot and the external-pod reference.
             if let OperationType::Native(NativeOperation::OpenInputStatement) = op.0 {
                 if let OperationAux::OpenInputStatement(InputPodOpenStatement {
-                    sts_root, ..
+                    pod_index, ..
                 }) = &op.2
                 {
                     deps.push(StatementSource::External(ExternalDependency {
-                        pod_hash: *sts_root,
+                        pod_index: *pod_index,
                         statement: statements[idx].clone(),
                     }));
                 }

src/frontend/multi_pod/diagnostics.rs

@@ -16,7 +16,8 @@ use std::{
 
 use super::{
     cost::{OperationCost, ResourceTotals},
-    shape::{AbstractDep, InputShape, OutputShape},
+    partition::SegmentChecker,
+    shape::{input_pod_slots, AbstractDep, InputShape, OutputShape},
 };
 use crate::middleware::Params;
 
@@ -198,7 +199,8 @@ pub struct PodUtilization {
     pub num_statements: usize,
     pub resources: Vec<UtilizationRow>,
     pub imports: UtilizationRow,
-    pub external_pods: UtilizationRow,
+    /// Number of input-POD slots occupied. See [`input_pod_slots`].
+    pub input_pod_slots: UtilizationRow,
     /// Statements this POD contributes to a Merkle tree. For intermediate
     /// PODs: locally-proved statements consumed downstream. For the output
     /// POD: the fresh tree size (`|output_public_indices|`). The limit is
@@ -215,9 +217,8 @@ pub struct SolutionBreakdown {
 }
 
 impl SolutionBreakdown {
-    /// Build a breakdown from an [`InputShape`] and its [`OutputShape`].
-    /// Re-derives per-POD imports and external-pod references from the
-    /// dep graph and the partition; both are pure functions of the inputs.
+    /// Computes per-POD resource use from an [`InputShape`] and its
+    /// [`OutputShape`].
     pub fn from_solution(input: &InputShape, output: &OutputShape) -> Self {
         let n = input.num_statements();
         let pod_count = output.pod_count;
@@ -298,25 +299,20 @@ impl SolutionBreakdown {
                     }
                 }
 
-                // Statement-table cap: each `OpenInputStatement` op
-                // produces a statement in the POD's statement table, so
-                // the "total statements" row reflects local statements
-                // PLUS chain and external imports: the same number
-                // `segment_feasible_with` checks against
-                // `max_statements`.
+                // Only chain imports add statements beyond assigned nodes.
                 let total_imports = chain_imports.len() + external_imports.len();
                 if let Some(row) = resources.iter_mut().find(|r| r.name == "total statements") {
-                    row.used += total_imports;
+                    row.used += chain_imports.len();
                 }
 
                 let imports_row = UtilizationRow {
                     name: "tree imports",
                     used: total_imports,
                     limit: input.params.max_open_input_statement_ops,
                 };
-                let external_row = UtilizationRow {
-                    name: "external pods",
-                    used: external_pods.len(),
+                let input_pods_row = UtilizationRow {
+                    name: "input pods",
+                    used: input_pod_slots(pod_idx == 0, external_pods.len()),
                     limit: input.params.max_input_pods,
                 };
                 let publishes_row = UtilizationRow {
@@ -331,7 +327,7 @@ impl SolutionBreakdown {
                     num_statements: num_stmts,
                     resources,
                     imports: imports_row,
-                    external_pods: external_row,
+                    input_pod_slots: input_pods_row,
                     publishes: publishes_row,
                 }
             })
@@ -364,7 +360,7 @@ impl fmt::Display for SolutionBreakdown {
             for row in
                 pod.resources
                     .iter()
-                    .chain([&pod.imports, &pod.external_pods, &pod.publishes])
+                    .chain([&pod.imports, &pod.input_pod_slots, &pod.publishes])
             {
                 if row.used > 0 {
                     let pct = if row.limit > 0 {
@@ -477,12 +473,13 @@ pub fn diagnose_failure(input: &InputShape) -> Option<CapViolation> {
     // doesn't fit; if it doesn't fit on its own (Unsplittable above didn't
     // trigger but the segment-relative caps overflow), surface the
     // violation that broke the camel's back.
+    let mut checker = SegmentChecker::new(&ordering, input);
     let mut segment_start = 0_usize;
     let mut segment_index = 0_usize;
     let mut pos = 0_usize;
     while pos < n {
         let next_pos = pos + 1;
-        if super::partition::segment_feasible(&ordering, input, segment_start, next_pos) {
+        if checker.is_feasible(segment_start, next_pos, false) {
             pos = next_pos;
             continue;
         }
@@ -547,10 +544,20 @@ fn identify_overflow(
     let s = ordering[a];
     let c = &input.costs[s];
 
+    let external_pods: HashSet<usize> = input.dep_edges[s]
+        .iter()
+        .filter_map(|dep| match dep {
+            AbstractDep::External { pod, .. } => Some(*pod),
+            AbstractDep::Internal(_) => None,
+        })
+        .collect();
+    let input_pods = input_pod_slots(segment_index == 0, external_pods.len());
+
     let state = &params.containers.state_ops;
     let transition = &params.containers.transition_ops;
     let categories: &[(&'static str, usize, usize)] = &[
         ("total statements", 1, params.max_statements),
+        ("input pods", input_pods, params.max_input_pods),
         (
             "merkle proofs (small)",
             c.merkle_proofs_small,
@@ -632,6 +639,58 @@ mod tests {
         assert_eq!(bottleneck.min_pods(), Some(3)); // ceil(6/2)
     }
 
+    /// Reports the predecessor slot when another resource cap moves a
+    /// statement from the first POD to a later POD.
+    #[test]
+    fn diagnose_failure_names_the_input_pod_cap() {
+        let params = Params {
+            max_signed_by_ops: 1,
+            ..Params::default()
+        };
+        let two_external_pods = vec![
+            AbstractDep::External {
+                pod: 0,
+                statement: 0,
+            },
+            AbstractDep::External {
+                pod: 1,
+                statement: 1,
+            },
+        ];
+        let signed_by_cost = OperationCost {
+            signed_by: 1,
+            ..OperationCost::default()
+        };
+        let input = InputShape {
+            costs: vec![signed_by_cost.clone(), signed_by_cost],
+            dep_edges: vec![two_external_pods.clone(), two_external_pods],
+            output_public_indices: vec![1],
+            num_external_pods: 2,
+            statement_pod: vec![0, 1],
+            params,
+        };
+
+        assert!(
+            partition::partition(&input).is_none(),
+            "the second POD cannot hold two external pods alongside the chain slot"
+        );
+        match diagnose_failure(&input).expect("failure must be diagnosed") {
+            CapViolation::Resource {
+                offending_stmt,
+                category,
+                used,
+                max_allowed,
+                ..
+            } => {
+                assert_eq!(category, "input pods");
+                assert_eq!(offending_stmt, 1);
+                assert_eq!(used, 3);
+                assert_eq!(max_allowed, 2);
+            }
+            other => panic!("expected an input-pod cap violation, got {}", other),
+        }
+    }
+
     #[test]
     fn solution_breakdown_reports_per_pod_utilisation() {
         let params = Params {