Improved predicate splitting (#445)

* Multi-batch splitting

* Invoke split predicates by name, passing in full argument list

* Reorder batches to prevent failure of forward references where possible

* Rename APIs for clarity

* Simplify example

* Add more docs

* Review updates

* Remove duplicate code

* Comment topological sort algorithm

Cargo.toml

@@ -34,6 +34,7 @@ rand = "0.8.5"
 hashbrown = { version = "0.14.3", default-features = false, features = ["serde"] }
 pest = "2.8.0"
 pest_derive = "2.8.0"
+petgraph = "0.6"
 directories = { version = "6.0.0", optional = true }
 minicbor-serde = { version = "0.5.0", features = ["std"], optional = true }
 serde_bytes = "0.11"

examples/main_pod_points.rs

@@ -88,7 +88,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         game_pk = game_pk,
     );
     println!("# custom predicate batch:{}", input);
-    let batch = parse(&input, &params, &[])?.custom_batch;
+    let batch = parse(&input, &params, &[])?
+        .first_batch()
+        .expect("Expected 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

@@ -1179,7 +1179,9 @@ pub mod tests {
             &[],
         )
         .unwrap()
-        .custom_batch;
+        .first_batch()
+        .unwrap()
+        .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

@@ -32,7 +32,11 @@ pub fn eth_dos_batch(params: &Params) -> Result<Arc<CustomPredicateBatch>> {
             eth_dos_ind(src, dst, distance)
         )
         "#;
-    let batch = parse(input, params, &[]).expect("lang parse").custom_batch;
+    let batch = parse(input, params, &[])
+        .expect("lang parse")
+        .first_batch()
+        .expect("Expected batch")
+        .clone();
     println!("a.0. {}", batch.predicates[0]);
     println!("a.1. {}", batch.predicates[1]);
     println!("a.2. {}", batch.predicates[2]);

src/frontend/error.rs

@@ -71,6 +71,12 @@ impl From<crate::lang::LangError> for Error {
     }
 }
 
+impl From<crate::lang::MultiOperationError> for Error {
+    fn from(value: crate::lang::MultiOperationError) -> Self {
+        Error::custom(value.to_string())
+    }
+}
+
 impl Debug for Error {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         std::fmt::Display::fmt(self, f)

src/frontend/mod.rs

@@ -1390,7 +1390,11 @@ pub mod tests {
             Equal(b, 5)
         )
         "#;
-        let batch = parse(input, &params, &[]).unwrap().custom_batch;
+        let batch = parse(input, &params, &[])
+            .unwrap()
+            .first_batch()
+            .unwrap()
+            .clone();
         let pred_test = batch.predicate_ref_by_name("Test").unwrap();
 
         // Try to build with wrong type in 1st arg
@@ -1414,4 +1418,83 @@ pub mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_apply_predicate_e2e() -> Result<()> {
+        // End-to-end test of apply_predicate with MockProver
+        // Tests a split predicate being applied through the full pipeline
+        let params = Params::default();
+        let vd_set = &*MOCK_VD_SET;
+
+        // Create a predicate that will split (6 Equal statements)
+        // The predicate checks that values at different keys are equal to specific literals
+        let input = r#"
+            large_pred(A) = AND(
+                Equal(A["a"], 1)
+                Equal(A["b"], 2)
+                Equal(A["c"], 3)
+                Equal(A["d"], 4)
+                Equal(A["e"], 5)
+                Equal(A["f"], 6)
+            )
+        "#;
+
+        // Parse and batch the predicate (this handles splitting internally)
+        let parsed = parse(input, &params, &[])?;
+        let batches = &parsed.custom_batches;
+
+        // Verify it was split
+        assert!(batches.split_chain("large_pred").is_some());
+        let chain_info = batches.split_chain("large_pred").unwrap();
+        assert_eq!(chain_info.chain_pieces.len(), 2);
+        assert_eq!(chain_info.real_statement_count, 6);
+
+        // Create a signed dict with the required entries
+        let mut signed_builder = SignedDictBuilder::new(&params);
+        signed_builder.insert("a", 1);
+        signed_builder.insert("b", 2);
+        signed_builder.insert("c", 3);
+        signed_builder.insert("d", 4);
+        signed_builder.insert("e", 5);
+        signed_builder.insert("f", 6);
+        let signer = Signer(SecretKey(1u32.into()));
+        let signed_dict = signed_builder.sign(&signer)?;
+
+        // Build the main pod
+        let mut builder = MainPodBuilder::new(&params, vd_set);
+        builder.pub_op(Operation::dict_signed_by(&signed_dict))?;
+
+        // Create 6 Equal statements (one for each predicate constraint) in original order
+        // Each proves that signed_dict["x"] = n, matching the Equal(A["x"], n) template
+        let st_a = builder.priv_op(Operation::eq((&signed_dict, "a"), 1))?;
+        let st_b = builder.priv_op(Operation::eq((&signed_dict, "b"), 2))?;
+        let st_c = builder.priv_op(Operation::eq((&signed_dict, "c"), 3))?;
+        let st_d = builder.priv_op(Operation::eq((&signed_dict, "d"), 4))?;
+        let st_e = builder.priv_op(Operation::eq((&signed_dict, "e"), 5))?;
+        let st_f = builder.priv_op(Operation::eq((&signed_dict, "f"), 6))?;
+
+        // Pass statements in original declaration order
+        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)?;
+
+        // The result should be a valid statement
+        let predicate = batches.predicate_ref_by_name("large_pred").unwrap();
+        match &result {
+            Statement::Custom(pred_ref, _) => {
+                assert_eq!(pred_ref, &predicate);
+            }
+            _ => panic!("Expected Statement::Custom, got {:?}", result),
+        }
+
+        // Prove with MockProver
+        let prover = MockProver {};
+        let pod = builder.prove(&prover)?;
+
+        // Verify the pod
+        pod.pod.verify()?;
+
+        Ok(())
+    }
 }

src/lang/error.rs

@@ -22,6 +22,9 @@ pub enum LangError {
 
     #[error("Lowering error: {0}")]
     Lowering(Box<LoweringError>),
+
+    #[error("Batching error: {0}")]
+    Batching(Box<BatchingError>),
 }
 
 /// Validation errors from frontend AST validation
@@ -90,14 +93,6 @@ pub enum ValidationError {
 /// Lowering errors from frontend AST lowering to middleware
 #[derive(Debug, thiserror::Error)]
 pub enum LoweringError {
-    #[error("Too many custom predicates in batch '{batch_name}': {count} exceeds limit of {max}{}", if *.original_count != *.count { format!(" (started with {} predicates before automatic splitting)", original_count) } else { String::new() })]
-    TooManyPredicates {
-        batch_name: String,
-        count: usize,
-        max: usize,
-        original_count: usize,
-    },
-
     #[error("Too many statements in predicate '{predicate}': {count} exceeds limit of {max}")]
     TooManyStatements {
         predicate: String,
@@ -127,6 +122,9 @@ pub enum LoweringError {
     #[error("Splitting error: {0}")]
     Splitting(#[from] SplittingError),
 
+    #[error("Batching error: {0}")]
+    Batching(#[from] BatchingError),
+
     #[error("Cannot lower document with validation errors")]
     ValidationErrors,
 }
@@ -235,6 +233,13 @@ fn format_public_args_at_split_error(
     msg
 }
 
+/// Batching errors from multi-batch packing
+#[derive(Debug, thiserror::Error)]
+pub enum BatchingError {
+    #[error("Internal batching error: {message}")]
+    Internal { message: String },
+}
+
 /// Splitting errors from predicate splitting
 #[derive(Debug, thiserror::Error)]
 pub enum SplittingError {
@@ -271,13 +276,6 @@ pub enum SplittingError {
         max_allowed: usize,
         suggestion: Option<Box<RefactorSuggestion>>,
     },
-
-    #[error("Too many predicates in chain for '{predicate}': {count} exceeds batch limit of {max_allowed}")]
-    TooManyPredicatesInChain {
-        predicate: String,
-        count: usize,
-        max_allowed: usize,
-    },
 }
 
 impl From<ParseError> for LangError {
@@ -303,3 +301,9 @@ impl From<LoweringError> for LangError {
         LangError::Lowering(Box::new(err))
     }
 }
+
+impl From<BatchingError> for LangError {
+    fn from(err: BatchingError) -> Self {
+        LangError::Batching(Box::new(err))
+    }
+}

src/lang/frontend_ast_lower.rs

@@ -1,38 +1,103 @@
 //! Lowering from frontend AST to middleware structures
 //!
 //! This module converts validated frontend AST to middleware data structures.
-//! Currently implements basic 1:1 conversion without automatic predicate splitting.
+//! Supports automatic predicate splitting and multi-batch packing.
 
 use std::{
     collections::{HashMap, HashSet},
     str::FromStr,
-    sync::Arc,
 };
 
 use crate::{
-    frontend::{BuilderArg, CustomPredicateBatchBuilder, StatementTmplBuilder},
+    frontend::{BuilderArg, StatementTmplBuilder},
     lang::{
         frontend_ast::*,
+        frontend_ast_batch::{self, PredicateBatches},
         frontend_ast_split,
         frontend_ast_validate::{PredicateKind, ValidatedAST},
     },
     middleware::{
-        self, containers, CustomPredicateBatch, IntroPredicateRef, NativePredicate, Params,
-        Predicate, PredicateOrWildcard, StatementTmpl as MWStatementTmpl,
+        self, containers, IntroPredicateRef, NativePredicate, Params, Predicate,
+        PredicateOrWildcard, StatementTmpl as MWStatementTmpl,
         StatementTmplArg as MWStatementTmplArg, Wildcard,
     },
 };
 
-/// Result of lowering: optional custom predicate batch and optional request
+// ============================================================================
+// Shared lowering utilities
+// ============================================================================
+// These functions convert AST types to middleware/builder types and are used
+// by both the request lowering (in this module) and predicate batching
+// (in frontend_ast_batch).
+
+/// Lower a literal value from AST to middleware Value.
+///
+/// This is a pure conversion that cannot fail.
+pub fn lower_literal(lit: &LiteralValue) -> middleware::Value {
+    match lit {
+        LiteralValue::Int(i) => middleware::Value::from(i.value),
+        LiteralValue::Bool(b) => middleware::Value::from(b.value),
+        LiteralValue::String(s) => middleware::Value::from(s.value.clone()),
+        LiteralValue::Raw(r) => middleware::Value::from(r.hash.hash),
+        LiteralValue::PublicKey(pk) => middleware::Value::from(pk.point),
+        LiteralValue::SecretKey(sk) => middleware::Value::from(sk.secret_key.clone()),
+        LiteralValue::Array(a) => {
+            let elements: Vec<_> = a.elements.iter().map(lower_literal).collect();
+            let array = containers::Array::new(elements);
+            middleware::Value::from(array)
+        }
+        LiteralValue::Set(s) => {
+            let elements: std::collections::HashSet<_> =
+                s.elements.iter().map(lower_literal).collect();
+            let set = containers::Set::new(elements);
+            middleware::Value::from(set)
+        }
+        LiteralValue::Dict(d) => {
+            let pairs: std::collections::HashMap<_, _> = d
+                .pairs
+                .iter()
+                .map(|pair| {
+                    let key = middleware::Key::from(pair.key.value.as_str());
+                    let value = lower_literal(&pair.value);
+                    (key, value)
+                })
+                .collect();
+            let dict = containers::Dictionary::new(pairs);
+            middleware::Value::from(dict)
+        }
+    }
+}
+
+/// Lower a statement argument from AST to BuilderArg.
+///
+/// This is a pure conversion that cannot fail.
+pub fn lower_statement_arg(arg: &StatementTmplArg) -> BuilderArg {
+    match arg {
+        StatementTmplArg::Literal(lit) => {
+            let value = lower_literal(lit);
+            BuilderArg::Literal(value)
+        }
+        StatementTmplArg::Wildcard(id) => BuilderArg::WildcardLiteral(id.name.clone()),
+        StatementTmplArg::AnchoredKey(ak) => {
+            let key_str = match &ak.key {
+                AnchoredKeyPath::Bracket(s) => s.value.clone(),
+                AnchoredKeyPath::Dot(id) => id.name.clone(),
+            };
+            BuilderArg::Key(ak.root.name.clone(), key_str)
+        }
+    }
+}
+
+/// Result of lowering: optional custom predicate batches and optional request
 ///
 /// A Podlang file can contain:
-/// - Just custom predicates (batch: Some, request: None)
-/// - Just a request (batch: None, request: Some)
-/// - Both (batch: Some, request: Some)
-/// - Neither (batch: None, request: None) - just imports
+/// - Just custom predicates (batches: Some, request: None)
+/// - Just a request (batches: None, request: Some)
+/// - Both (batches: Some, request: Some)
+/// - Neither (batches: None, request: None) - just imports
 #[derive(Debug, Clone)]
 pub struct LoweredOutput {
-    pub batch: Option<Arc<CustomPredicateBatch>>,
+    pub batches: Option<PredicateBatches>,
     pub request: Option<crate::frontend::PodRequest>,
 }
 
@@ -60,71 +125,70 @@ pub fn lower(
 struct Lowerer<'a> {
     validated: ValidatedAST,
     params: &'a Params,
-    /// Map of predicate names to their index in the current batch (for split predicates)
-    batch_predicate_index: HashMap<String, usize>,
 }
 
 impl<'a> Lowerer<'a> {
     fn new(validated: ValidatedAST, params: &'a Params) -> Self {
-        Self {
-            validated,
-            params,
-            batch_predicate_index: HashMap::new(),
-        }
+        Self { validated, params }
     }
 
-    fn lower(mut self, batch_name: String) -> Result<LoweredOutput, LoweringError> {
-        // Lower custom predicates (if any)
-        let batch = self.lower_batch(batch_name)?;
+    fn lower(self, batch_name: String) -> Result<LoweredOutput, LoweringError> {
+        // Lower custom predicates (if any) - now supports multiple batches
+        let batches = self.lower_batches(batch_name)?;
 
-        // Lower request (if any) - pass batch so BatchSelf refs can be converted to Custom refs
-        let request = self.lower_request(batch.as_ref())?;
+        // Lower request (if any) - pass batches so refs can be resolved
+        let request = self.lower_request(batches.as_ref())?;
 
-        Ok(LoweredOutput { batch, request })
+        Ok(LoweredOutput { batches, request })
     }
 
-    fn lower_batch(
-        &mut self,
-        batch_name: String,
-    ) -> Result<Option<Arc<CustomPredicateBatch>>, LoweringError> {
+    fn lower_batches(&self, batch_name: String) -> Result<Option<PredicateBatches>, LoweringError> {
         // Extract and split custom predicates from document
-        let (custom_predicates, original_count) = self.extract_and_split_predicates()?;
+        let custom_predicates = self.extract_and_split_predicates()?;
 
         // If no custom predicates, return None
         if custom_predicates.is_empty() {
             return Ok(None);
         }
 
-        // Check batch size constraint
-        if custom_predicates.len() > self.params.max_custom_batch_size {
-            return Err(LoweringError::TooManyPredicates {
-                batch_name: batch_name.clone(),
-                count: custom_predicates.len(),
-                max: self.params.max_custom_batch_size,
-                original_count,
-            });
+        // Build map of imported predicates for batching
+        let imported_predicates = self.build_imported_predicates_map();
+
+        // Use the new batching module to pack predicates into batches
+        let batches = frontend_ast_batch::batch_predicates(
+            custom_predicates,
+            self.params,
+            &batch_name,
+            &imported_predicates,
+        )?;
+
+        Ok(Some(batches))
     }
 
-        // Build index of all predicates in the batch
-        for (idx, pred) in custom_predicates.iter().enumerate() {
-            self.batch_predicate_index
-                .insert(pred.name.name.clone(), idx);
+    fn build_imported_predicates_map(
+        &self,
+    ) -> HashMap<String, frontend_ast_batch::ImportedPredicateInfo> {
+        let symbols = self.validated.symbols();
+        let mut imported = HashMap::new();
+
+        for (name, info) in &symbols.predicates {
+            if let PredicateKind::BatchImported { batch, index } = &info.kind {
+                imported.insert(
+                    name.clone(),
+                    frontend_ast_batch::ImportedPredicateInfo {
+                        batch: batch.clone(),
+                        index: *index,
+                    },
+                );
+            }
         }
 
-        // Create custom predicate batch using builder
-        let mut cpb_builder =
-            CustomPredicateBatchBuilder::new(self.params.clone(), batch_name.clone());
-
-        for pred_def in &custom_predicates {
-            self.lower_custom_predicate(pred_def, &mut cpb_builder)?;
-        }
-
-        Ok(Some(cpb_builder.finish()))
+        imported
     }
 
     fn lower_request(
         &self,
-        batch: Option<&Arc<CustomPredicateBatch>>,
+        batches: Option<&PredicateBatches>,
     ) -> Result<Option<crate::frontend::PodRequest>, LoweringError> {
         let doc = self.validated.document();
 
@@ -141,43 +205,77 @@ impl<'a> Lowerer<'a> {
         // Build wildcard map from all wildcards used in the request statements
         let wildcard_map = self.build_request_wildcard_map(request_def);
 
-        // Lower each statement to a builder first
-        let mut statement_builders = Vec::new();
-        for stmt in &request_def.statements {
-            let stmt_builder = self.lower_statement_to_builder(stmt)?;
-            statement_builders.push(stmt_builder);
-        }
-
-        // Resolve builders to middleware statement templates
+        // Lower each statement to middleware templates, resolving predicates
         let mut request_templates = Vec::new();
-        for stmt_builder in statement_builders {
-            let mw_stmt =
-                self.resolve_request_statement_builder(stmt_builder, &wildcard_map, batch)?;
+        for stmt in &request_def.statements {
+            let mw_stmt = self.lower_request_statement(stmt, &wildcard_map, batches)?;
             request_templates.push(mw_stmt);
         }
 
         Ok(Some(crate::frontend::PodRequest::new(request_templates)))
     }
 
-    fn resolve_request_statement_builder(
+    fn lower_request_statement(
         &self,
-        stmt_builder: StatementTmplBuilder,
+        stmt: &StatementTmpl,
         wildcard_map: &HashMap<String, usize>,
-        batch: Option<&Arc<CustomPredicateBatch>>,
+        batches: Option<&PredicateBatches>,
     ) -> Result<MWStatementTmpl, LoweringError> {
-        // First desugar the builder
-        let desugared = stmt_builder.desugar();
+        // Enforce argument count limit for request statements
+        if stmt.args.len() > self.params.max_statement_args {
+            return Err(LoweringError::TooManyStatementArgs {
+                count: stmt.args.len(),
+                max: self.params.max_statement_args,
+            });
+        }
 
-        // Convert BatchSelf predicate to Custom if we have a batch
-        let mut predicate = desugared.predicate;
-        if let Some(batch_ref) = batch {
-            if let Predicate::BatchSelf(index) = predicate {
-                predicate = Predicate::Custom(middleware::CustomPredicateRef::new(
-                    batch_ref.clone(),
-                    index,
-                ));
+        let pred_name = &stmt.predicate.name;
+        let symbols = self.validated.symbols();
+
+        // Resolve predicate - for request statements, local custom predicates
+        // must be resolved to CustomPredicateRef (not BatchSelf)
+        let predicate = if let Ok(native) = NativePredicate::from_str(pred_name) {
+            Predicate::Native(native)
+        } else if let Some(info) = symbols.predicates.get(pred_name) {
+            match &info.kind {
+                PredicateKind::Native(np) => Predicate::Native(*np),
+                PredicateKind::Custom { .. } => {
+                    // Local custom predicates - resolve to CustomPredicateRef
+                    let batches = batches.ok_or_else(|| LoweringError::PredicateNotFound {
+                        name: pred_name.clone(),
+                    })?;
+                    let pred_ref = batches.predicate_ref_by_name(pred_name).ok_or_else(|| {
+                        LoweringError::PredicateNotFound {
+                            name: pred_name.clone(),
                         }
+                    })?;
+                    Predicate::Custom(pred_ref)
                 }
+                PredicateKind::BatchImported { batch, index } => {
+                    Predicate::Custom(middleware::CustomPredicateRef::new(batch.clone(), *index))
+                }
+                PredicateKind::IntroImported {
+                    name,
+                    verifier_data_hash,
+                } => Predicate::Intro(IntroPredicateRef {
+                    name: name.clone(),
+                    args_len: info.public_arity,
+                    verifier_data_hash: *verifier_data_hash,
+                }),
+            }
+        } else {
+            return Err(LoweringError::PredicateNotFound {
+                name: pred_name.clone(),
+            });
+        };
+
+        // Create a builder with the resolved predicate and desugar
+        let mut builder = StatementTmplBuilder::new(predicate);
+        for arg in &stmt.args {
+            let builder_arg = lower_statement_arg(arg);
+            builder = builder.arg(builder_arg);
+        }
+        let desugared = builder.desugar();
 
         // Convert BuilderArgs to StatementTmplArgs
         let mut mw_args = Vec::new();
@@ -202,7 +300,7 @@ impl<'a> Lowerer<'a> {
 
         Ok(MWStatementTmpl {
             // TODO: Support wildcard
-            pred_or_wc: PredicateOrWildcard::Predicate(predicate),
+            pred_or_wc: PredicateOrWildcard::Predicate(desugared.predicate),
             args: mw_args,
         })
     }
@@ -251,7 +349,7 @@ impl<'a> Lowerer<'a> {
 
     fn extract_and_split_predicates(
         &self,
-    ) -> Result<(Vec<CustomPredicateDef>, usize), LoweringError> {
+    ) -> Result<Vec<frontend_ast_split::SplitResult>, LoweringError> {
         let doc = self.validated.document();
         let predicates: Vec<CustomPredicateDef> = doc
             .items
@@ -262,182 +360,14 @@ impl<'a> Lowerer<'a> {
             })
             .collect();
 
-        let original_count = predicates.len();
-
         // Apply splitting to each predicate as needed
-        let mut split_predicates = Vec::new();
+        let mut split_results = Vec::new();
         for pred in predicates {
-            let chain = frontend_ast_split::split_predicate_if_needed(pred, self.params)?;
-            split_predicates.extend(chain);
+            let result = frontend_ast_split::split_predicate_if_needed(pred, self.params)?;
+            split_results.push(result);
         }
 
-        Ok((split_predicates, original_count))
-    }
-
-    fn lower_custom_predicate(
-        &self,
-        pred_def: &CustomPredicateDef,
-        cpb_builder: &mut CustomPredicateBatchBuilder,
-    ) -> Result<(), LoweringError> {
-        let name = pred_def.name.name.clone();
-
-        // Note: Constraint checking is handled by the splitting phase
-        // Predicates passed here should already be within limits
-
-        // Collect public and private argument names
-        let mut public_arg_names = Vec::new();
-        let mut private_arg_names = Vec::new();
-
-        for arg in &pred_def.args.public_args {
-            public_arg_names.push(arg.name.clone());
-        }
-
-        if let Some(private_args) = &pred_def.args.private_args {
-            for arg in private_args {
-                private_arg_names.push(arg.name.clone());
-            }
-        }
-
-        // Lower statements to builders
-        let mut statement_builders = Vec::new();
-        for stmt in &pred_def.statements {
-            let stmt_builder = self.lower_statement_to_builder(stmt)?;
-            statement_builders.push(stmt_builder);
-        }
-
-        // Convert to &str slices for builder API
-        let public_args_str: Vec<&str> = public_arg_names.iter().map(|s| s.as_str()).collect();
-        let private_args_str: Vec<&str> = private_arg_names.iter().map(|s| s.as_str()).collect();
-
-        // Add predicate to batch using builder
-        let conjunction = pred_def.conjunction_type == ConjunctionType::And;
-
-        cpb_builder
-            .predicate(
-                &name,
-                conjunction,
-                &public_args_str,
-                &private_args_str,
-                &statement_builders,
-            )
-            .map_err(|e| match e {
-                crate::frontend::Error::Middleware(mw_err) => LoweringError::Middleware(mw_err),
-                _ => LoweringError::InvalidArgumentType,
-            })?;
-
-        Ok(())
-    }
-
-    fn lower_statement_to_builder(
-        &self,
-        stmt: &StatementTmpl,
-    ) -> Result<StatementTmplBuilder, LoweringError> {
-        // Get predicate
-        let pred_name = &stmt.predicate.name;
-        let symbols = self.validated.symbols();
-
-        // Check for native predicates first
-        let predicate = if let Ok(native) = NativePredicate::from_str(pred_name) {
-            Predicate::Native(native)
-        } else if let Some(&index) = self.batch_predicate_index.get(pred_name) {
-            // References to other predicates in the same batch (including split chains)
-            Predicate::BatchSelf(index)
-        } else if let Some(info) = symbols.predicates.get(pred_name) {
-            match &info.kind {
-                PredicateKind::Native(np) => Predicate::Native(*np),
-                PredicateKind::Custom { index } => Predicate::BatchSelf(*index),
-                PredicateKind::BatchImported { batch, index } => {
-                    Predicate::Custom(middleware::CustomPredicateRef::new(batch.clone(), *index))
-                }
-                PredicateKind::IntroImported {
-                    name,
-                    verifier_data_hash,
-                } => Predicate::Intro(IntroPredicateRef {
-                    name: name.clone(),
-                    args_len: info.public_arity,
-                    verifier_data_hash: *verifier_data_hash,
-                }),
-            }
-        } else {
-            unreachable!("Predicate {} not found", pred_name);
-        };
-
-        // Check args count
-        if stmt.args.len() > self.params.max_statement_args {
-            return Err(LoweringError::TooManyStatementArgs {
-                count: stmt.args.len(),
-                max: self.params.max_statement_args,
-            });
-        }
-
-        // Convert AST args to BuilderArgs
-        let mut builder = StatementTmplBuilder::new(predicate);
-        for arg in &stmt.args {
-            let builder_arg = Self::lower_statement_arg_to_builder(arg)?;
-            builder = builder.arg(builder_arg);
-        }
-
-        // Return builder without calling .desugar() - that will happen later
-        Ok(builder)
-    }
-
-    fn lower_statement_arg_to_builder(arg: &StatementTmplArg) -> Result<BuilderArg, LoweringError> {
-        match arg {
-            StatementTmplArg::Literal(lit) => {
-                let value = Self::lower_literal(lit)?;
-                Ok(BuilderArg::Literal(value))
-            }
-            StatementTmplArg::Wildcard(id) => {
-                // For builder, we just need the wildcard name
-                Ok(BuilderArg::WildcardLiteral(id.name.clone()))
-            }
-            StatementTmplArg::AnchoredKey(ak) => {
-                let key_str = match &ak.key {
-                    AnchoredKeyPath::Bracket(s) => s.value.clone(),
-                    AnchoredKeyPath::Dot(id) => id.name.clone(),
-                };
-                Ok(BuilderArg::Key(ak.root.name.clone(), key_str))
-            }
-        }
-    }
-
-    fn lower_literal(lit: &LiteralValue) -> Result<middleware::Value, LoweringError> {
-        let value = match lit {
-            LiteralValue::Int(i) => middleware::Value::from(i.value),
-            LiteralValue::Bool(b) => middleware::Value::from(b.value),
-            LiteralValue::String(s) => middleware::Value::from(s.value.clone()),
-            LiteralValue::Raw(r) => middleware::Value::from(r.hash.hash),
-            LiteralValue::PublicKey(pk) => middleware::Value::from(pk.point),
-            LiteralValue::SecretKey(sk) => middleware::Value::from(sk.secret_key.clone()),
-            LiteralValue::Array(a) => {
-                let elements: Result<Vec<_>, _> =
-                    a.elements.iter().map(Self::lower_literal).collect();
-                let array = containers::Array::new(elements?);
-                middleware::Value::from(array)
-            }
-            LiteralValue::Set(s) => {
-                let elements: Result<Vec<_>, _> =
-                    s.elements.iter().map(Self::lower_literal).collect();
-                let set_values: std::collections::HashSet<_> = elements?.into_iter().collect();
-                let set = containers::Set::new(set_values);
-                middleware::Value::from(set)
-            }
-            LiteralValue::Dict(d) => {
-                let pairs: Result<Vec<(middleware::Key, middleware::Value)>, LoweringError> = d
-                    .pairs
-                    .iter()
-                    .map(|pair| {
-                        let key = middleware::Key::from(pair.key.value.as_str());
-                        let value = Self::lower_literal(&pair.value)?;
-                        Ok((key, value))
-                    })
-                    .collect();
-                let dict_map: std::collections::HashMap<_, _> = pairs?.into_iter().collect();
-                let dict = containers::Dictionary::new(dict_map);
-                middleware::Value::from(dict)
-            }
-        };
-        Ok(value)
+        Ok(split_results)
     }
 }
 
@@ -458,9 +388,16 @@ mod tests {
         lower(validated, params, "test_batch".to_string())
     }
 
-    // Helper to get the batch from the output (expecting it to exist)
-    fn expect_batch(output: &LoweredOutput) -> &Arc<CustomPredicateBatch> {
-        output.batch.as_ref().expect("Expected batch to be present")
+    // Helper to get the first batch from the output (expecting it to exist)
+    fn expect_batch(
+        output: &LoweredOutput,
+    ) -> &std::sync::Arc<crate::middleware::CustomPredicateBatch> {
+        output
+            .batches
+            .as_ref()
+            .expect("Expected batches to be present")
+            .first_batch()
+            .expect("Expected at least one batch")
     }
 
     #[test]
@@ -547,13 +484,20 @@ mod tests {
 
         let lowered = result.unwrap();
         // Should be automatically split into 2 predicates (my_pred and my_pred_1)
-        assert_eq!(expect_batch(&lowered).predicates().len(), 2);
+        let batches = lowered.batches.as_ref().expect("Expected batches");
+        assert_eq!(batches.total_predicate_count(), 2);
 
-        // First predicate should have 5 statements (4 + chain call)
-        assert_eq!(expect_batch(&lowered).predicates()[0].statements().len(), 5);
-
-        // Second predicate should have 2 statements
-        assert_eq!(expect_batch(&lowered).predicates()[1].statements().len(), 2);
+        // With topological sorting, my_pred_1 comes first (since my_pred depends on it)
+        // my_pred_1 has 2 statements
+        // my_pred has 5 statements (4 + chain call)
+        // Just verify we have the right total statement counts
+        let batch = batches.first_batch().unwrap();
+        let total_statements: usize = batch
+            .predicates()
+            .iter()
+            .map(|p| p.statements().len())
+            .sum();
+        assert_eq!(total_statements, 7); // 5 + 2 = 7 total statements
     }
 
     #[test]
@@ -642,108 +586,64 @@ mod tests {
     }
 
     #[test]
-    fn test_error_message_with_splitting() {
-        // Create a document with predicates that will exceed the batch limit after splitting
-        // We'll create 2 predicates with 4 statements each (max arity = 5)
-        // Each will NOT split individually, but together they exceed a small batch limit
+    fn test_multi_batch_packing() {
+        // Create more predicates than fit in a single batch
+        // With max_custom_batch_size = 4, 5 predicates should span 2 batches
         let input = r#"
-            pred1(A) = AND (
-                Equal(A["a"], 1)
-                Equal(A["b"], 2)
-            )
-            pred2(B) = AND (
-                Equal(B["c"], 3)
-                Equal(B["d"], 4)
-            )
+            pred1(A) = AND(Equal(A["a"], 1))
+            pred2(B) = AND(Equal(B["b"], 2))
+            pred3(C) = AND(Equal(C["c"], 3))
+            pred4(D) = AND(Equal(D["d"], 4))
+            pred5(E) = AND(Equal(E["e"], 5))
         "#;
 
-        // Use very restrictive params to force the error
-        let params = Params {
-            max_custom_batch_size: 1,
-            ..Default::default()
-        };
+        let params = Params::default(); // max_custom_batch_size = 4
 
         let result = parse_validate_and_lower(input, &params);
+        assert!(result.is_ok());
 
-        // Should fail with TooManyPredicates error
-        assert!(result.is_err());
-        let err = result.unwrap_err();
+        let lowered = result.unwrap();
+        let batches = lowered.batches.as_ref().expect("Expected batches");
 
-        if let LoweringError::TooManyPredicates {
-            count,
-            max,
-            original_count,
-            ..
-        } = err
-        {
-            assert_eq!(count, 2); // 2 predicates after splitting (no splitting occurred)
-            assert_eq!(max, 1);
-            assert_eq!(original_count, 2); // Started with 2 predicates
+        // Should have 2 batches
+        assert_eq!(batches.batch_count(), 2);
+        assert_eq!(batches.total_predicate_count(), 5);
 
-            // Error message should NOT mention splitting since no splitting occurred
-            let err_msg = format!("{}", err);
-            assert!(!err_msg.contains("before automatic splitting"));
-        } else {
-            panic!("Expected TooManyPredicates error, got: {:?}", err);
-        }
+        // First batch should have 4 predicates
+        assert_eq!(batches.batches()[0].predicates().len(), 4);
+        // Second batch should have 1 predicate
+        assert_eq!(batches.batches()[1].predicates().len(), 1);
     }
 
     #[test]
-    fn test_error_message_after_splitting() {
-        // Create TWO predicates that will EACH split into 2 predicates
-        // This tests the case where splitting causes the batch to be too large
-        // but no individual predicate chain exceeds the limit
+    fn test_split_chains_span_batches() {
+        // Create predicates that will split, plus additional predicates
+        // to force the split chains across batch boundaries
         let input = r#"
-            pred1(A) = AND (
-                Equal(A["a"], 1)
-                Equal(A["b"], 2)
-                Equal(A["c"], 3)
-                Equal(A["d"], 4)
-                Equal(A["e"], 5)
-                Equal(A["f"], 6)
-            )
-            pred2(B) = AND (
-                Equal(B["a"], 1)
-                Equal(B["b"], 2)
-                Equal(B["c"], 3)
-                Equal(B["d"], 4)
-                Equal(B["e"], 5)
-                Equal(B["f"], 6)
+            pred1(A) = AND(Equal(A["a"], 1))
+            pred2(B) = AND(Equal(B["b"], 2))
+            pred3(C) = AND(Equal(C["c"], 3))
+            large_pred(D) = AND(
+                Equal(D["a"], 1)
+                Equal(D["b"], 2)
+                Equal(D["c"], 3)
+                Equal(D["d"], 4)
+                Equal(D["e"], 5)
+                Equal(D["f"], 6)
             )
         "#;
 
-        // Use params where each predicate splits into 2, but total of 4 exceeds batch limit
-        let params = Params {
-            // Allow 3 predicates in batch
-            // Default max_custom_predicate_arity is 5, so each will split into 2 predicates
-            // Total: 2 original predicates -> 4 after splitting (exceeds limit of 3)
-            max_custom_batch_size: 3,
-            ..Default::default()
-        };
+        let params = Params::default();
 
         let result = parse_validate_and_lower(input, &params);
+        assert!(result.is_ok());
 
-        // Should fail with TooManyPredicates error
-        assert!(result.is_err());
-        let err = result.unwrap_err();
+        let lowered = result.unwrap();
+        let batches = lowered.batches.as_ref().expect("Expected batches");
 
-        if let LoweringError::TooManyPredicates {
-            count,
-            max,
-            original_count,
-            ..
-        } = err
-        {
-            assert_eq!(count, 4); // 4 predicates after splitting (2 from each)
-            assert_eq!(max, 3);
-            assert_eq!(original_count, 2); // Started with 2 predicates
-
-            // Error message SHOULD mention splitting since splitting occurred
-            let err_msg = format!("{}", err);
-            assert!(err_msg.contains("before automatic splitting"));
-            assert!(err_msg.contains("started with 2 predicates"));
-        } else {
-            panic!("Expected TooManyPredicates error, got: {:?}", err);
-        }
+        // pred1, pred2, pred3 + large_pred split into 2 = 5 total predicates
+        // Should span 2 batches
+        assert_eq!(batches.total_predicate_count(), 5);
+        assert_eq!(batches.batch_count(), 2);
     }
 }

src/lang/frontend_ast_split.rs

@@ -34,6 +34,40 @@ pub struct ChainLink {
     pub public_args_out: Vec<String>,
 }
 
+/// Information about a single piece of a split predicate chain
+#[derive(Debug, Clone)]
+pub struct SplitChainPiece {
+    /// Name of this predicate piece (e.g., "foo_1")
+    pub name: String,
+    /// Number of "real" statements in this piece (excludes chain call)
+    pub real_statement_count: usize,
+    /// Whether this piece has a chain call to the next piece
+    pub has_chain_call: bool,
+}
+
+/// Metadata about a split predicate chain
+#[derive(Debug, Clone)]
+pub struct SplitChainInfo {
+    /// Original predicate name (e.g., "foo")
+    pub original_name: String,
+    /// Chain pieces in execution order (innermost continuation first: [foo_2, foo_1, foo])
+    pub chain_pieces: Vec<SplitChainPiece>,
+    /// Total number of "real" user statements (excludes chain calls)
+    pub real_statement_count: usize,
+    /// Maps original statement index → reordered index
+    /// e.g., if original stmt 0 became reordered stmt 3, then `reorder_map[0] = 3`
+    pub reorder_map: Vec<usize>,
+}
+
+/// Result of splitting a predicate
+#[derive(Debug, Clone)]
+pub struct SplitResult {
+    /// The predicates (continuations first, original last if split)
+    pub predicates: Vec<CustomPredicateDef>,
+    /// Split chain info, if splitting occurred (None for non-split)
+    pub chain_info: Option<SplitChainInfo>,
+}
+
 /// Wildcard usage information
 #[derive(Debug, Clone)]
 struct WildcardUsage {
@@ -66,19 +100,25 @@ pub fn validate_predicate_is_splittable(
 pub fn split_predicate_if_needed(
     pred: CustomPredicateDef,
     params: &Params,
-) -> Result<Vec<CustomPredicateDef>, SplittingError> {
+) -> Result<SplitResult, SplittingError> {
     // Early validation
     validate_predicate_is_splittable(&pred, params)?;
 
     // If within limits, no splitting needed
     if pred.statements.len() <= params.max_custom_predicate_arity {
-        return Ok(vec![pred]);
+        return Ok(SplitResult {
+            predicates: vec![pred],
+            chain_info: None,
+        });
     }
 
     // Need to split - execute the splitting algorithm
-    let chain = split_into_chain(pred, params)?;
+    let (predicates, chain_info) = split_into_chain(pred, params)?;
 
-    Ok(chain)
+    Ok(SplitResult {
+        predicates,
+        chain_info: Some(chain_info),
+    })
 }
 
 fn analyze_wildcards(statements: &[StatementTmpl]) -> HashMap<String, WildcardUsage> {
@@ -121,18 +161,33 @@ fn collect_wildcards_from_statement(stmt: &StatementTmpl) -> HashSet<String> {
 }
 
 /// Order constraints optimally to minimize liveness at boundaries
+/// Result of ordering statements optimally for splitting
+struct OrderingResult {
+    /// Reordered statements
+    statements: Vec<StatementTmpl>,
+    /// Maps original statement index → reordered index
+    /// reorder_map[original_idx] = new_idx
+    reorder_map: Vec<usize>,
+}
+
 fn order_constraints_optimally(
     statements: Vec<StatementTmpl>,
     _usage: &HashMap<String, WildcardUsage>,
     params: &Params,
-) -> Vec<StatementTmpl> {
-    // If no splitting needed, preserve original order
-    if statements.len() <= params.max_custom_predicate_arity {
-        return statements;
+) -> OrderingResult {
+    let n = statements.len();
+
+    // If no splitting needed, preserve original order (identity mapping)
+    if n <= params.max_custom_predicate_arity {
+        return OrderingResult {
+            statements,
+            reorder_map: (0..n).collect(),
+        };
     }
 
     let mut ordered = Vec::new();
-    let mut remaining: HashSet<usize> = (0..statements.len()).collect();
+    let mut reorder_map = vec![0; n];
+    let mut remaining: HashSet<usize> = (0..n).collect();
     let mut active_wildcards: HashSet<String> = HashSet::new();
 
     while !remaining.is_empty() {
@@ -146,6 +201,9 @@ fn order_constraints_optimally(
 
         remaining.remove(&best_idx);
         let stmt = &statements[best_idx];
+
+        // Record the mapping: original index best_idx → new index ordered.len()
+        reorder_map[best_idx] = ordered.len();
         ordered.push(stmt.clone());
 
         // Update active wildcards
@@ -160,7 +218,10 @@ fn order_constraints_optimally(
         active_wildcards.retain(|w| needed_later.contains(w));
     }
 
-    ordered
+    OrderingResult {
+        statements: ordered,
+        reorder_map,
+    }
 }
 
 /// Compute tie-breaker metrics for deterministic ordering when scores are equal
@@ -360,16 +421,20 @@ fn generate_refactor_suggestion(
 }
 
 /// Split into chain using bucket-filling approach
+/// Returns the split predicates and metadata about the split
 fn split_into_chain(
     pred: CustomPredicateDef,
     params: &Params,
-) -> Result<Vec<CustomPredicateDef>, SplittingError> {
+) -> Result<(Vec<CustomPredicateDef>, SplitChainInfo), SplittingError> {
     let original_name = pred.name.name.clone();
     let conjunction = pred.conjunction_type;
 
     let usage = analyze_wildcards(&pred.statements);
+    let real_statement_count = pred.statements.len();
 
-    let ordered_statements = order_constraints_optimally(pred.statements, &usage, params);
+    let ordering_result = order_constraints_optimally(pred.statements, &usage, params);
+    let ordered_statements = ordering_result.statements;
+    let reorder_map = ordering_result.reorder_map;
 
     let original_public_args: Vec<String> = pred
         .args
@@ -479,12 +544,43 @@ fn split_into_chain(
         }
     }
 
-    let chain_predicates =
+    // Build SplitChainInfo from chain_links before generating predicates
+    // Pieces are in execution order: innermost continuation first, original last
+    let num_links = chain_links.len();
+    let mut chain_pieces = Vec::new();
+    for i in (0..num_links).rev() {
+        let link = &chain_links[i];
+        let is_last = i == num_links - 1;
+        let name = if i == 0 {
+            original_name.clone()
+        } else {
+            format!("{}_{}", original_name, i)
+        };
+        chain_pieces.push(SplitChainPiece {
+            name,
+            real_statement_count: link.statements.len(),
+            has_chain_call: !is_last,
+        });
+    }
+
+    let chain_info = SplitChainInfo {
+        original_name: original_name.clone(),
+        chain_pieces,
+        real_statement_count,
+        reorder_map,
+    };
+
+    let mut chain_predicates =
         generate_chain_predicates(&original_name, chain_links, conjunction, params)?;
 
-    validate_chain(&chain_predicates, &original_name, params)?;
+    validate_chain(&chain_predicates, params)?;
 
-    Ok(chain_predicates)
+    // Reverse so continuations come before callers in declaration order.
+    // This ensures that when batched, continuations are in earlier batches
+    // and can be referenced by their callers.
+    chain_predicates.reverse();
+
+    Ok((chain_predicates, chain_info))
 }
 
 /// Phase 4: Generate synthetic predicates from chain links
@@ -519,20 +615,19 @@ fn generate_chain_predicates(
                 span: None,
             };
 
-            // Create arguments for chain call: all public args (incoming + promoted)
-            let mut chain_call_args = Vec::new();
-            for arg_name in &link.public_args_in {
-                chain_call_args.push(StatementTmplArg::Wildcard(Identifier {
+            // Create arguments for chain call: use next link's public_args_in
+            // which is the deduplicated union of current public_args_in and public_args_out
+            let next_link = &chain_links[i + 1];
+            let chain_call_args: Vec<StatementTmplArg> = next_link
+                .public_args_in
+                .iter()
+                .map(|arg_name| {
+                    StatementTmplArg::Wildcard(Identifier {
                         name: arg_name.clone(),
                         span: None,
-                }));
-            }
-            for arg_name in &link.public_args_out {
-                chain_call_args.push(StatementTmplArg::Wildcard(Identifier {
-                    name: arg_name.clone(),
-                    span: None,
-                }));
-            }
+                    })
+                })
+                .collect();
 
             let chain_call = StatementTmpl {
                 predicate: next_pred_name,
@@ -587,19 +682,10 @@ fn generate_chain_predicates(
 }
 
 /// Phase 5: Validate the generated chain
-fn validate_chain(
-    chain: &[CustomPredicateDef],
-    original_name: &str,
-    params: &Params,
-) -> Result<(), SplittingError> {
-    if chain.len() > params.max_custom_batch_size {
-        return Err(SplittingError::TooManyPredicatesInChain {
-            predicate: original_name.to_string(),
-            count: chain.len(),
-            max_allowed: params.max_custom_batch_size,
-        });
-    }
-
+///
+/// Note: We no longer check chain length against max_custom_batch_size since
+/// chains can now span multiple batches thanks to multi-batch support.
+fn validate_chain(chain: &[CustomPredicateDef], params: &Params) -> Result<(), SplittingError> {
     for pred in chain {
         // Each predicate should have ≤ max_statements
         assert!(pred.statements.len() <= params.max_custom_predicate_arity);
@@ -681,8 +767,9 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
-        assert_eq!(chain.len(), 1); // No split needed
+        let split_result = result.unwrap();
+        assert_eq!(split_result.predicates.len(), 1); // No split needed
+        assert!(split_result.chain_info.is_none()); // No chain info for non-split
     }
 
     #[test]
@@ -704,14 +791,29 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
+        let split_result = result.unwrap();
+        let chain = &split_result.predicates;
         assert_eq!(chain.len(), 2); // Should split into 2 predicates
 
-        // First predicate: 4 statements + chain call = 5
-        assert_eq!(chain[0].statements.len(), 5);
+        // Chain is reversed: continuation comes first, original comes last
+        // Last predicate (index 1): original name, 4 statements + chain call = 5
+        assert_eq!(chain[1].statements.len(), 5);
+        assert_eq!(chain[1].name.name, "my_pred");
 
-        // Second predicate: 2 remaining statements
-        assert_eq!(chain[1].statements.len(), 2);
+        // First predicate (index 0): continuation, 2 remaining statements
+        assert_eq!(chain[0].statements.len(), 2);
+        assert_eq!(chain[0].name.name, "my_pred_1");
+
+        // Verify chain_info is present
+        let chain_info = split_result.chain_info.as_ref().unwrap();
+        assert_eq!(chain_info.original_name, "my_pred");
+        assert_eq!(chain_info.real_statement_count, 6);
+        assert_eq!(chain_info.chain_pieces.len(), 2);
+        // Pieces are in execution order: innermost first
+        assert_eq!(chain_info.chain_pieces[0].name, "my_pred_1");
+        assert!(!chain_info.chain_pieces[0].has_chain_call);
+        assert_eq!(chain_info.chain_pieces[1].name, "my_pred");
+        assert!(chain_info.chain_pieces[1].has_chain_call);
     }
 
     #[test]
@@ -733,12 +835,15 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
+        let split_result = result.unwrap();
+        let chain = &split_result.predicates;
         assert_eq!(chain.len(), 2); // Should split into 2 predicates
 
-        // First predicate should have wildcards that cross boundary promoted
-        // Check that chain call is present
-        let last_stmt = &chain[0].statements.last().unwrap();
+        // Chain is reversed: continuation first, original last
+        // Original predicate should have chain call as last statement
+        let original = &chain[1];
+        assert_eq!(original.name.name, "complex");
+        let last_stmt = original.statements.last().unwrap();
         assert_eq!(last_stmt.predicate.name, "complex_1");
     }
 
@@ -766,15 +871,29 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
+        let split_result = result.unwrap();
+        let chain = &split_result.predicates;
         assert_eq!(chain.len(), 3); // Should split into 3 predicates
 
-        // First: 4 + chain call = 5
-        assert_eq!(chain[0].statements.len(), 5);
-        // Second: 4 + chain call = 5
+        // Chain is reversed: [_2, _1, original]
+        // chain[0]: large_pred_2 (3 remaining statements)
+        assert_eq!(chain[0].statements.len(), 3);
+        assert_eq!(chain[0].name.name, "large_pred_2");
+        // chain[1]: large_pred_1 (4 + chain call = 5)
         assert_eq!(chain[1].statements.len(), 5);
-        // Third: 3 remaining
-        assert_eq!(chain[2].statements.len(), 3);
+        assert_eq!(chain[1].name.name, "large_pred_1");
+        // chain[2]: large_pred (4 + chain call = 5)
+        assert_eq!(chain[2].statements.len(), 5);
+        assert_eq!(chain[2].name.name, "large_pred");
+
+        // Verify chain_info
+        let chain_info = split_result.chain_info.as_ref().unwrap();
+        assert_eq!(chain_info.real_statement_count, 11);
+        assert_eq!(chain_info.chain_pieces.len(), 3);
+        // Execution order: innermost first
+        assert_eq!(chain_info.chain_pieces[0].name, "large_pred_2");
+        assert_eq!(chain_info.chain_pieces[1].name, "large_pred_1");
+        assert_eq!(chain_info.chain_pieces[2].name, "large_pred");
     }
 
     #[test]
@@ -801,7 +920,8 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
+        let split_result = result.unwrap();
+        let chain = &split_result.predicates;
         // Should split into 2 predicates
         // T is used in first segment and crosses to second, then used again in second
         assert_eq!(chain.len(), 2);
@@ -867,7 +987,8 @@ mod tests {
         let result = split_predicate_if_needed(pred, &params);
         assert!(result.is_ok());
 
-        let chain = result.unwrap();
+        let split_result = result.unwrap();
+        let chain = &split_result.predicates;
         assert_eq!(chain.len(), 2, "Predicate should split into 2 links");
 
         let second_pred = &chain[1];

src/lang/mod.rs

@@ -1,5 +1,34 @@
+//! Podlang front-end: parsing, validation, lowering, and multi-batch output.
+//!
+//! This module is the high-level entrypoint to the Podlang pipeline. It:
+//! - Parses a Podlang document (`parse_podlang`).
+//! - Validates names, imports, and well-formedness (`frontend_ast_validate`).
+//! - Lowers to middleware structures, including automatic predicate splitting and
+//!   dependency-aware packing into one or more custom predicate batches (`frontend_ast_split`,
+//!   `frontend_ast_batch`, `frontend_ast_lower`).
+//!
+//! The result is a [`PodlangOutput`], which contains:
+//! - `custom_batches`: a [`PredicateBatches`] container (possibly empty) with all custom
+//!   predicates defined in the document. Use
+//!   [`PredicateBatches::apply_predicate`](crate::lang::frontend_ast_batch::PredicateBatches::apply_predicate)
+//!   to apply a predicate into a `MainPodBuilder` (recommended primary API), or
+//!   [`apply_predicate_with`](crate::lang::frontend_ast_batch::PredicateBatches::apply_predicate_with)
+//!   for advanced control.
+//! - `request`: a `PodRequest` containing the request templates defined by a `REQUEST(...)` block
+//!   in the document (or empty if none was provided).
+//!
+//! Notes
+//! - Predicate splitting: large predicates are automatically split into a chain of smaller
+//!   predicates while preserving semantics; only the final chain result is public when applying a
+//!   predicate as public.
+//! - Multi-batch packing: predicates are packed dependency-aware; cross-batch references always
+//!   point to earlier batches and forward references cannot occur.
+//! - Backwards compatibility: `PodlangOutput::first_batch()` is provided to ease migration of code
+//!   that expects a single custom predicate batch.
+//!
 pub mod error;
 pub mod frontend_ast;
+pub mod frontend_ast_batch;
 pub mod frontend_ast_lower;
 pub mod frontend_ast_split;
 pub mod frontend_ast_validate;
@@ -9,6 +38,8 @@ pub mod pretty_print;
 use std::sync::Arc;
 
 pub use error::LangError;
+pub use frontend_ast_batch::{MultiOperationError, PredicateBatches};
+pub use frontend_ast_split::{SplitChainInfo, SplitChainPiece, SplitResult};
 pub use parser::{parse_podlang, Pairs, ParseError, Rule};
 pub use pretty_print::PrettyPrint;
 
@@ -17,12 +48,34 @@ use crate::{
     middleware::{CustomPredicateBatch, Params},
 };
 
-#[derive(Debug, Clone, PartialEq)]
+/// Final result of processing a Podlang document.
+///
+/// - `custom_batches`: all custom predicates defined in the document, possibly spanning multiple
+///   batches. Use [`PredicateBatches`] APIs to look up predicates by name and apply them.
+/// - `request`: the request templates defined in the document (empty if not present).
+#[derive(Debug, Clone)]
 pub struct PodlangOutput {
-    pub custom_batch: Arc<CustomPredicateBatch>,
+    pub custom_batches: PredicateBatches,
     pub request: PodRequest,
 }
 
+impl PodlangOutput {
+    /// Get the first batch, if any (for backwards compatibility).
+    ///
+    /// Prefer using `custom_batches` directly if your code expects multiple batches.
+    pub fn first_batch(&self) -> Option<&Arc<CustomPredicateBatch>> {
+        self.custom_batches.first_batch()
+    }
+}
+
+/// Parse, validate, and lower a Podlang document into middleware structures.
+///
+/// - `input`: Podlang source.
+/// - `params`: middleware parameters limiting sizes/arity and controlling lowering behavior.
+/// - `available_batches`: external batches available for `use batch ... from 0x...` imports.
+///
+/// Returns a [`PodlangOutput`] containing custom predicate batches (if any) and a `PodRequest`
+/// (possibly empty).
 pub fn parse(
     input: &str,
     params: &Params,
@@ -37,10 +90,7 @@ pub fn parse(
     let validated = frontend_ast_validate::validate(document, available_batches)?;
     let lowered = frontend_ast_lower::lower(validated, params, "PodlangBatch".to_string())?;
 
-    let custom_batch = lowered.batch.unwrap_or_else(|| {
-        // If no batch, create an empty one
-        CustomPredicateBatch::new(params, "PodlangBatch".to_string(), vec![])
-    });
+    let custom_batches = lowered.batches.unwrap_or_default();
 
     let request = lowered.request.unwrap_or_else(|| {
         // If no request, create an empty one
@@ -48,7 +98,7 @@ pub fn parse(
     });
 
     Ok(PodlangOutput {
-        custom_batch,
+        custom_batches,
         request,
     })
 }
@@ -93,6 +143,11 @@ mod tests {
         PredicateOrWildcard::Predicate(pred)
     }
 
+    // Helper to get the first batch from the output
+    fn first_batch(output: &super::PodlangOutput) -> &Arc<CustomPredicateBatch> {
+        output.first_batch().expect("Expected at least one batch")
+    }
+
     #[test]
     fn test_e2e_simple_predicate() -> Result<(), LangError> {
         let input = r#"
@@ -103,14 +158,12 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
+        let batch_result = first_batch(&processed);
         let request_result = processed.request.templates();
 
         assert_eq!(request_result.len(), 0);
         assert_eq!(batch_result.predicates.len(), 1);
 
-        let batch = batch_result;
-
         // Expected structure
         let expected_statements = vec![StatementTmpl {
             pred_or_wc: pred_lit(Predicate::Native(NativePredicate::Equal)),
@@ -132,7 +185,7 @@ mod tests {
             vec![expected_predicate],
         );
 
-        assert_eq!(batch, expected_batch);
+        assert_eq!(*batch_result, expected_batch);
 
         Ok(())
     }
@@ -148,10 +201,9 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
         let request_templates = processed.request.templates();
 
-        assert_eq!(batch_result.predicates.len(), 0);
+        assert!(processed.custom_batches.is_empty());
         assert!(!request_templates.is_empty());
 
         // Expected structure
@@ -188,14 +240,12 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
+        let batch_result = first_batch(&processed);
         let request_result = processed.request.templates();
 
         assert_eq!(request_result.len(), 0);
         assert_eq!(batch_result.predicates.len(), 1);
 
-        let batch = batch_result;
-
         // Expected structure: Public args: A (index 0). Private args: Temp (index 1)
         let expected_statements = vec![
             StatementTmpl {
@@ -226,7 +276,7 @@ mod tests {
             vec![expected_predicate],
         );
 
-        assert_eq!(batch, expected_batch);
+        assert_eq!(*batch_result, expected_batch);
 
         Ok(())
     }
@@ -245,14 +295,12 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
+        let batch_result = first_batch(&processed);
         let request_templates = processed.request.templates();
 
         assert_eq!(batch_result.predicates.len(), 1);
         assert!(!request_templates.is_empty());
 
-        let batch = batch_result;
-
         // Expected Batch structure
         let expected_pred_statements = vec![StatementTmpl {
             pred_or_wc: pred_lit(Predicate::Native(NativePredicate::Equal)),
@@ -274,7 +322,7 @@ mod tests {
             vec![expected_predicate],
         );
 
-        assert_eq!(batch, expected_batch);
+        assert_eq!(*batch_result, expected_batch);
 
         // Expected Request structure
         // Pod1 -> Wildcard 0, Pod2 -> Wildcard 1
@@ -311,7 +359,7 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
+        let batch_result = first_batch(&processed);
         let request_templates = processed.request.templates();
 
         assert_eq!(batch_result.predicates.len(), 1); // some_pred is defined
@@ -324,7 +372,10 @@ mod tests {
         // Expected structure
         let expected_templates = vec![
             StatementTmpl {
-                pred_or_wc: pred_lit(Predicate::Custom(CustomPredicateRef::new(batch_result, 0))), // Refers to some_pred
+                pred_or_wc: pred_lit(Predicate::Custom(CustomPredicateRef::new(
+                    batch_result.clone(),
+                    0,
+                ))), // Refers to some_pred
                 args: vec![
                     StatementTmplArg::Wildcard(wc("Var1", 0)),        // Var1
                     StatementTmplArg::Literal(Value::from(12345i64)), // 12345
@@ -361,10 +412,9 @@ mod tests {
 
         let params = Params::default();
         let processed = parse(input, &params, &[])?;
-        let batch_result = processed.custom_batch;
         let request_templates = processed.request.templates();
 
-        assert_eq!(batch_result.predicates.len(), 0);
+        assert!(processed.custom_batches.is_empty());
         assert!(!request_templates.is_empty());
 
         let expected_templates = vec![
@@ -509,7 +559,7 @@ mod tests {
         );
 
         assert!(
-            processed.custom_batch.predicates.is_empty(),
+            processed.custom_batches.is_empty(),
             "Expected no custom predicates for a REQUEST only input"
         );
 
@@ -560,7 +610,7 @@ mod tests {
             "Expected no request templates"
         );
         assert_eq!(
-            processed.custom_batch.predicates.len(),
+            first_batch(&processed).predicates.len(),
             4,
             "Expected 4 custom predicates"
         );
@@ -691,7 +741,8 @@ mod tests {
         );
 
         assert_eq!(
-            processed.custom_batch, expected_batch,
+            *first_batch(&processed),
+            expected_batch,
             "Processed ETHDoS predicates do not match expected structure"
         );
 
@@ -739,7 +790,7 @@ mod tests {
         let request_templates = processed.request.templates();
 
         assert!(
-            processed.custom_batch.predicates.is_empty(),
+            processed.custom_batches.is_empty(),
             "No custom predicates should be defined in the main input"
         );
         assert_eq!(request_templates.len(), 1, "Expected one request template");
@@ -860,13 +911,13 @@ mod tests {
             "No request should be defined"
         );
         assert_eq!(
-            processed.custom_batch.predicates.len(),
+            first_batch(&processed).predicates.len(),
             1,
             "Expected one custom predicate to be defined"
         );
 
         // 4. Check the resulting predicate definition
-        let defined_pred = &processed.custom_batch.predicates[0];
+        let defined_pred = &first_batch(&processed).predicates[0];
         assert_eq!(defined_pred.name, "wrapper_pred");
         assert_eq!(defined_pred.statements.len(), 1);
 

src/lang/pretty_print.rs

@@ -395,15 +395,17 @@ mod tests {
             parse(input, &params, available_batches).expect("Initial parsing should succeed");
 
         // Step 2: Pretty-print the parsed batch
-        let pretty_printed = parsed_result.custom_batch.to_podlang_string();
+        let batch = parsed_result.first_batch().expect("Expected batch");
+        let pretty_printed = batch.to_podlang_string();
 
         // Step 3: Parse the pretty-printed result
         let reparsed_result =
             parse(&pretty_printed, &params, available_batches).expect("Reparsing should succeed");
+        let reparsed_batch = reparsed_result.first_batch().expect("Expected batch");
 
         // Step 4: Verify the ASTs are equivalent
         assert_eq!(
-            parsed_result.custom_batch.predicates, reparsed_result.custom_batch.predicates,
+            batch.predicates, reparsed_batch.predicates,
             "Original AST should match reparsed AST.\nOriginal input:\n{}\nPretty-printed:\n{}\n",
             input, pretty_printed
         );
@@ -553,18 +555,17 @@ mod tests {
 
         let params = Params::default();
         let parsed_result = parse(input, &params, &[]).expect("Parsing should succeed");
+        let batch = parsed_result.first_batch().expect("Expected batch");
 
-        let pretty_printed = parsed_result.custom_batch.to_podlang_string();
+        let pretty_printed = batch.to_podlang_string();
 
         println!("Original input:\n{}", input);
         println!("\nPretty-printed output:\n{}", pretty_printed);
 
         let reparsed = parse(&pretty_printed, &params, &[]).expect("Reparsing should succeed");
+        let reparsed_batch = reparsed.first_batch().expect("Expected batch");
 
-        assert_eq!(
-            parsed_result.custom_batch.predicates,
-            reparsed.custom_batch.predicates
-        );
+        assert_eq!(batch.predicates, reparsed_batch.predicates);
     }
 
     #[test]
@@ -627,14 +628,16 @@ mod tests {
 
             let params = Params::default();
             let parsed_result = parse(&input, &params, &[]).expect("Should parse successfully");
+            let batch = parsed_result.first_batch().expect("Expected batch");
 
-            let pretty_printed = parsed_result.custom_batch.to_podlang_string();
+            let pretty_printed = batch.to_podlang_string();
 
             let reparsed_result =
                 parse(&pretty_printed, &params, &[]).expect("Should reparse successfully");
+            let reparsed_batch = reparsed_result.first_batch().expect("Expected batch");
 
             assert_eq!(
-                parsed_result.custom_batch.predicates, reparsed_result.custom_batch.predicates,
+                batch.predicates, reparsed_batch.predicates,
                 "Round-trip failed for string: {:?}\nPretty-printed: {}",
                 test_string, pretty_printed
             );