Remove batch splitting system

examples/main_pod_points.rs

@@ -8,15 +8,15 @@
 //!
 //! Run in real mode: `cargo run --release --example main_pod_points`
 //! Run in mock mode: `cargo run --release --example main_pod_points -- --mock`
-use std::env;
+use std::{collections::HashMap, env};
 
 use pod2::{
     backends::plonky2::{
         basetypes::DEFAULT_VD_SET, mainpod::Prover, mock::mainpod::MockProver,
         primitives::ec::schnorr::SecretKey, signer::Signer,
     },
     frontend::{MainPodBuilder, Operation, SignedDictBuilder},
-    lang::parse,
+    lang::load_module,
     middleware::{MainPodProver, Params, VDSet},
 };
 
@@ -88,10 +88,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         game_pk = game_pk,
     );
     println!("# custom predicate batch:{}", input);
-    let batch = parse(&input, &params, &[])?
-        .first_batch()
-        .expect("Expected batch")
-        .clone();
+    let module = load_module(&input, "points_module", &params, &HashMap::new())?;
+    let batch = module.batch.clone();
     let points_pred = batch.predicate_ref_by_name("points").unwrap();
     let over_9000_pred = batch.predicate_ref_by_name("over_9000").unwrap();
 

src/backends/plonky2/mainpod/mod.rs

@@ -817,7 +817,10 @@ impl Pod for MainPod {
 
 #[cfg(test)]
 pub mod tests {
-    use std::{any::Any, collections::HashSet};
+    use std::{
+        any::Any,
+        collections::{HashMap, HashSet},
+    };
 
     use num::{BigUint, One};
 
@@ -836,7 +839,7 @@ pub mod tests {
         frontend::{
             self, literal, CustomPredicateBatchBuilder, MainPodBuilder, StatementTmplBuilder as STB,
         },
-        lang::parse,
+        lang::load_module,
         middleware::{
             self, containers::Set, CustomPredicateRef, NativePredicate as NP, Signer as _,
             DEFAULT_VD_LIST, DEFAULT_VD_SET,
@@ -1165,21 +1168,20 @@ pub mod tests {
     #[test]
     fn test_undetermined_values() {
         let params = Default::default();
-        let batch = parse(
+        let module = load_module(
             r#"
             two_equal(x,y,z) = OR(
                 Equal(x,y)
                 Equal(y,z)
                 Equal(x,z)
             )
             "#,
+            "test",
             &params,
-            &[],
+            &HashMap::new(),
         )
-        .unwrap()
-        .first_batch()
-        .unwrap()
-        .clone();
+        .unwrap();
+        let batch = module.batch.clone();
         let mut builder = MainPodBuilder::new(&params, &DEFAULT_VD_SET);
         let cpr = CustomPredicateRef { batch, index: 0 };
         let eq_st = builder.priv_op(frontend::Operation::eq(1, 1)).unwrap();

src/examples/custom.rs

@@ -1,10 +1,8 @@
-use std::sync::Arc;
-
-use hex::ToHex;
+use std::{collections::HashMap, sync::Arc};
 
 use crate::{
     frontend::{PodRequest, Result},
-    lang::parse,
+    lang::{load_module, parse_request, Module},
     middleware::{CustomPredicateBatch, Params},
 };
 
@@ -32,11 +30,8 @@ pub fn eth_dos_batch(params: &Params) -> Result<Arc<CustomPredicateBatch>> {
             eth_dos_ind(src, dst, distance)
         )
         "#;
-    let batch = parse(input, params, &[])
-        .expect("lang parse")
-        .first_batch()
-        .expect("Expected batch")
-        .clone();
+    let module = load_module(input, "eth_dos", params, &HashMap::new()).expect("lang parse");
+    let batch = module.batch.clone();
     println!("a.0. {}", batch.predicates()[0]);
     println!("a.1. {}", batch.predicates()[1]);
     println!("a.2. {}", batch.predicates()[2]);
@@ -46,17 +41,22 @@ pub fn eth_dos_batch(params: &Params) -> Result<Arc<CustomPredicateBatch>> {
 
 pub fn eth_dos_request() -> Result<PodRequest> {
     let batch = eth_dos_batch(&Params::default())?;
-    let batch_id = batch.id().encode_hex::<String>();
-    let input = format!(
-        r#"
-        use batch _, _, _, eth_dos from 0x{batch_id}
+    let eth_dos_module = Arc::new(Module::new(batch, HashMap::new()));
+
+    let mut available_modules = HashMap::new();
+    available_modules.insert("eth_dos_mod".to_string(), eth_dos_module);
+
+    let input = r#"
+        use module eth_dos_mod
         REQUEST(
-            eth_dos(src, dst, distance)
+            eth_dos_mod::eth_dos(src, dst, distance)
         )
-        "#,
-    );
-    let parsed = parse(&input, &Params::default(), &[batch])?;
-    Ok(parsed.request)
+        "#;
+    Ok(parse_request(
+        input,
+        &Params::default(),
+        &available_modules,
+    )?)
 }
 
 #[cfg(test)]

src/examples/mod.rs

@@ -1,6 +1,9 @@
 pub mod custom;
 
-use std::{collections::HashSet, sync::LazyLock};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::LazyLock,
+};
 
 use custom::eth_dos_batch;
 use num::BigUint;
@@ -12,7 +15,7 @@ use crate::{
     frontend::{
         MainPod, MainPodBuilder, Operation, PodRequest, Result, SignedDict, SignedDictBuilder,
     },
-    lang::parse,
+    lang::parse_request,
     middleware::{
         self, containers::Set, hash_values, CustomPredicateRef, Params, Predicate, PublicKey,
         Signer as _, Statement, StatementArg, TypedValue, VDSet, Value,
@@ -90,8 +93,7 @@ pub fn zu_kyc_pod_request(gov_signer: &Value, pay_signer: &Value) -> Result<PodR
     )
     "#,
     );
-    let parsed = parse(&input, &Params::default(), &[])?;
-    Ok(parsed.request)
+    Ok(parse_request(&input, &Params::default(), &HashMap::new())?)
 }
 
 // ETHDoS

src/frontend/mod.rs

@@ -797,6 +797,7 @@ impl MainPodCompiler {
 
 #[cfg(test)]
 pub mod tests {
+    use std::collections::HashMap;
 
     use num::BigUint;
 
@@ -812,7 +813,7 @@ pub mod tests {
             tickets_pod_full_flow, zu_kyc_pod_builder, zu_kyc_pod_request,
             zu_kyc_sign_dict_builders, EthDosHelper, MOCK_VD_SET,
         },
-        lang::parse,
+        lang::load_module,
         middleware::{
             containers::{Array, Set},
             Signer as _, Value,
@@ -1381,11 +1382,8 @@ pub mod tests {
             Equal(b, 5)
         )
         "#;
-        let batch = parse(input, &params, &[])
-            .unwrap()
-            .first_batch()
-            .unwrap()
-            .clone();
+        let module = load_module(input, "test", &params, &HashMap::new()).unwrap();
+        let batch = module.batch.clone();
         let pred_test = batch.predicate_ref_by_name("Test").unwrap();
 
         // Try to build with wrong type in 1st arg
@@ -1433,11 +1431,8 @@ pub mod tests {
             c(6, 3)
         )
         "#;
-        let batch = parse(input, &params, &[])
-            .unwrap()
-            .first_batch()
-            .unwrap()
-            .clone();
+        let module = load_module(input, "test", &params, &HashMap::new()).unwrap();
+        let batch = module.batch.clone();
         let pred_test = batch.predicate_ref_by_name("Test").unwrap();
 
         let mut builder = MainPodBuilder::new(&params, vd_set);
@@ -1458,11 +1453,8 @@ pub mod tests {
             c(6, 3)
         )
         "#;
-        let batch = parse(input, &params, &[])
-            .unwrap()
-            .first_batch()
-            .unwrap()
-            .clone();
+        let module = load_module(input, "test", &params, &HashMap::new()).unwrap();
+        let batch = module.batch.clone();
         let pred_test = batch.predicate_ref_by_name("Test").unwrap();
 
         let mut builder = MainPodBuilder::new(&params, vd_set);
@@ -1500,12 +1492,11 @@ pub mod tests {
         "#;
 
         // Parse and batch the predicate (this handles splitting internally)
-        let parsed = parse(input, &params, &[])?;
-        let batches = &parsed.custom_batches;
+        let module = load_module(input, "test", &params, &HashMap::new())?;
 
         // Verify it was split
-        assert!(batches.split_chain("large_pred").is_some());
-        let chain_info = batches.split_chain("large_pred").unwrap();
+        assert!(module.split_chains.contains_key("large_pred"));
+        let chain_info = module.split_chains.get("large_pred").unwrap();
         assert_eq!(chain_info.chain_pieces.len(), 2);
         assert_eq!(chain_info.real_statement_count, 6);
 
@@ -1537,10 +1528,10 @@ pub mod tests {
         let statements = vec![st_a, st_b, st_c, st_d, st_e, st_f];
 
         // Use apply_predicate (primary API) to automatically wire the split chain
-        let result = batches.apply_predicate(&mut builder, "large_pred", statements, true)?;
+        let result = module.apply_predicate(&mut builder, "large_pred", statements, true)?;
 
         // The result should be a valid statement
-        let predicate = batches.predicate_ref_by_name("large_pred").unwrap();
+        let predicate = module.predicate_ref_by_name("large_pred").unwrap();
         match &result {
             Statement::Custom(pred_ref, _) => {
                 assert_eq!(pred_ref, &predicate);

src/frontend/multi_pod/mod.rs

@@ -539,7 +539,7 @@ mod tests {
         dict,
         examples::MOCK_VD_SET,
         frontend::{Operation as FrontendOp, SignedDictBuilder},
-        lang::parse,
+        lang::load_module,
     };
 
     #[test]
@@ -666,18 +666,17 @@ mod tests {
 
         // pred_a accepts a Contains statement
         // pred_b accepts a pred_a statement (Custom statement from pred_a)
-        let parsed = parse(
+        let module = load_module(
             r#"
             pred_a(X) = AND(Contains(X, "k", 1))
             pred_b(X) = AND(pred_a(X))
             "#,
+            "test",
             &params,
-            &[],
+            &HashMap::new(),
         )
-        .expect("parse predicates");
-        let batch = parsed
-            .first_batch()
-            .expect("parse predicates should have a batch");
+        .expect("load module");
+        let batch = &module.batch;
 
         let mut builder = MultiPodBuilder::new(&params, vd_set);
 
@@ -1398,20 +1397,19 @@ mod tests {
         let vd_set = &*MOCK_VD_SET;
 
         // Chain of predicates: each accepts the output of the previous
-        let parsed = parse(
+        let module = load_module(
             r#"
             pred_a(X) = AND(Contains(X, "k", 1))
             pred_b(X) = AND(pred_a(X))
             pred_c(X) = AND(pred_b(X))
             pred_d(X) = AND(pred_c(X))
             "#,
+            "test",
             &params,
-            &[],
+            &HashMap::new(),
         )
-        .expect("parse predicates");
-        let batch = parsed
-            .first_batch()
-            .expect("parse predicates should have a batch");
+        .expect("load module");
+        let batch = &module.batch;
 
         let mut builder = MultiPodBuilder::new(&params, vd_set);
 
@@ -1526,7 +1524,7 @@ mod tests {
         // pred_a takes TWO custom statement arguments (b_out and c_out)
         // pred_b and pred_c each take a Contains
         // Note: AND clauses are newline-separated, not comma-separated
-        let parsed = parse(
+        let module = load_module(
             r#"
             pred_b(X) = AND(Contains(X, "k", 1))
             pred_c(X) = AND(Contains(X, "k", 1))
@@ -1535,13 +1533,12 @@ mod tests {
                 pred_c(Y)
             )
             "#,
+            "test",
             &params,
-            &[],
+            &HashMap::new(),
         )
-        .expect("parse predicates");
-        let batch = parsed
-            .first_batch()
-            .expect("parse predicates should have a batch");
+        .expect("load module");
+        let batch = &module.batch;
 
         let mut builder = MultiPodBuilder::new(&params, vd_set);