More Podlang parser performance improvements

Author: robknight

Cargo.toml

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

src/frontend/custom.rs

@@ -1,5 +1,10 @@
 #![allow(unused)]
-use std::{collections::HashMap, fmt, hash as h, iter, iter::zip, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    fmt, hash as h, iter,
+    iter::zip,
+    sync::Arc,
+};
 
 use schemars::JsonSchema;
 
@@ -132,6 +137,9 @@ pub struct CustomPredicateBatchBuilder {
     params: Params,
     pub name: String,
     pub predicates: Vec<CustomPredicate>,
+    /// Names added via `predicate()`, kept alongside `predicates` so the
+    /// duplicate-name check doesn't rescan the whole batch per insertion.
+    predicate_names: HashSet<String>,
     /// Forward references to resolve in finish(): (predicate_idx, statement_idx, arg_idx, name)
     pending_self_pred_hashes: Vec<(usize, usize, usize, String)>,
 }
@@ -142,6 +150,7 @@ impl CustomPredicateBatchBuilder {
             params,
             name,
             predicates: Vec::new(),
+            predicate_names: HashSet::new(),
             pending_self_pred_hashes: Vec::new(),
         }
     }
@@ -176,7 +185,7 @@ impl CustomPredicateBatchBuilder {
         priv_args: &[&str],
         sts: &[StatementTmplBuilder],
     ) -> Result<Predicate> {
-        if self.predicates.iter().any(|p| p.name == name) {
+        if self.predicate_names.contains(name) {
             return Err(Error::custom(format!(
                 "Duplicate predicate name '{}' in batch",
                 name
@@ -271,6 +280,7 @@ impl CustomPredicateBatchBuilder {
                 .collect(),
         )?;
         self.predicates.push(custom_predicate);
+        self.predicate_names.insert(name.to_string());
         self.pending_self_pred_hashes.extend(pending);
         Ok(Predicate::BatchSelf(self.predicates.len() - 1))
     }

src/lang/frontend_ast_lower.rs

@@ -539,7 +539,7 @@ impl<'a> Lowerer<'a> {
         &self,
     ) -> Result<Vec<frontend_ast_split::SplitResult>, LoweringError> {
         let doc = self.validated.document();
-        let predicates: Vec<CustomPredicateDef> = doc
+        let mut predicates: Vec<CustomPredicateDef> = doc
             .items
             .iter()
             .filter_map(|item| match item {
@@ -548,23 +548,15 @@ impl<'a> Lowerer<'a> {
             })
             .collect();
 
-        // Apply splitting to each predicate as needed. The typed-key rewrite
-        // happens before splitting so split chain pieces inherit `Index` keys
-        // unchanged. The search cache is shared across the module: modules
-        // routinely contain families of same-shape predicates, which then
-        // pay for the ordering search once.
-        let split_started = std::time::Instant::now();
-        let mut search_cache = frontend_ast_split::SplitSearchCache::default();
-        let mut split_results = Vec::new();
-        for mut pred in predicates {
-            self.rewrite_typed_dot_access(&mut pred);
-            let result = frontend_ast_split::split_predicate_if_needed(
-                pred,
-                self.params,
-                &mut search_cache,
-            )?;
-            split_results.push(result);
+        // The typed-key rewrite happens before splitting so split chain
+        // pieces inherit `Index` keys unchanged.
+        for pred in &mut predicates {
+            self.rewrite_typed_dot_access(pred);
         }
+
+        let split_started = std::time::Instant::now();
+        let split_results =
+            frontend_ast_split::split_predicates_if_needed(predicates, self.params)?;
         log::debug!(
             "predicate splitting: {:?} ({} predicates)",
             split_started.elapsed(),

src/lang/frontend_ast_split.rs

@@ -97,8 +97,9 @@ pub(super) fn validate_predicate_is_splittable(
 }
 
 /// Split a predicate into a chain if it exceeds statement limit. Callers
-/// splitting many predicates under one `Params` (module lowering) share one
-/// cache so same-shape predicates search only once.
+/// splitting many predicates under one `Params` share one cache so same-shape
+/// predicates search only once; whole-module callers should prefer
+/// [`split_predicates_if_needed`], which also parallelizes the searches.
 pub fn split_predicate_if_needed(
     pred: CustomPredicateDef,
     params: &Params,
@@ -124,6 +125,73 @@ pub fn split_predicate_if_needed(
     })
 }
 
+/// Split a whole module's predicates. The link assignment is a pure function
+/// of a predicate's statement/wildcard shape, so each distinct shape is
+/// searched once and the searches run in parallel across shapes.
+pub fn split_predicates_if_needed(
+    predicates: Vec<CustomPredicateDef>,
+    params: &Params,
+) -> Result<Vec<SplitResult>, SplittingError> {
+    use rayon::prelude::*;
+
+    for pred in &predicates {
+        validate_predicate_is_splittable(pred)?;
+    }
+
+    // Predicates over the statement cap get a prepared split input; the rest
+    // pass through whole.
+    let inputs: Vec<Option<SplitInput>> = predicates
+        .iter()
+        .map(|pred| {
+            (pred.statements.len() > Params::max_custom_predicate_arity())
+                .then(|| prepare_split_input(pred))
+        })
+        .collect();
+
+    // One representative per distinct shape, in first-encounter order so a
+    // search failure surfaces the same predicate the sequential loop would
+    // have blamed.
+    let mut seen_shapes: HashSet<&SplitShape> = HashSet::new();
+    let mut representatives: Vec<(&str, &SplitInput)> = Vec::new();
+    for (pred, input) in predicates.iter().zip(&inputs) {
+        if let Some(input) = input {
+            if seen_shapes.insert(&input.shape) {
+                representatives.push((pred.name.name.as_str(), input));
+            }
+        }
+    }
+
+    let searched: Vec<Result<LinkAssignment, SplittingError>> = representatives
+        .par_iter()
+        .map(|(name, input)| search_link_assignment(name, input, params))
+        .collect();
+    let mut assignment_of_shape: HashMap<&SplitShape, LinkAssignment> = HashMap::new();
+    for ((_, input), result) in representatives.iter().zip(searched) {
+        assignment_of_shape.insert(&input.shape, result?);
+    }
+
+    let results: Vec<Result<SplitResult, SplittingError>> = predicates
+        .into_par_iter()
+        .zip(&inputs)
+        .map(|(pred, input)| match input {
+            None => Ok(SplitResult {
+                predicates: vec![pred],
+                chain_info: None,
+            }),
+            Some(input) => {
+                let assignment = assignment_of_shape[&input.shape].clone();
+                let (chain_predicates, chain_info) =
+                    build_chain_from_assignment(&pred, input, assignment, params)?;
+                Ok(SplitResult {
+                    predicates: chain_predicates,
+                    chain_info: Some(chain_info),
+                })
+            }
+        })
+        .collect();
+    results.into_iter().collect()
+}
+
 fn collect_wildcards_from_statement(stmt: &StatementTmpl) -> HashSet<String> {
     stmt.wildcard_names().map(str::to_string).collect()
 }
@@ -852,6 +920,7 @@ enum CandidateSearch {
 fn candidate_orderings(input: &SplitInput, params: &Params) -> CandidateSearch {
     use rand::{seq::SliceRandom, SeedableRng};
     use rand_chacha::ChaCha20Rng;
+    use rayon::prelude::*;
 
     let num_statements = input.shape.num_statements;
     let max_args = Params::max_statement_args();
@@ -882,26 +951,29 @@ fn candidate_orderings(input: &SplitInput, params: &Params) -> CandidateSearch {
         refinement_seeds.push(shuffled);
     }
 
-    let mut refined: Vec<(Vec<usize>, usize)> = Vec::new();
-    let mut shortcut_attempted = false;
-    for seed in refinement_seeds {
-        let result = refine_ordering(
-            seed,
-            &input.shape.statements_using,
-            &input.shape.is_original_public,
-            max_args,
-            REFINE_ITERATIONS,
-        );
-        let result_cost = input.excess_cost(&result);
-        if result_cost == 0 && !shortcut_attempted {
-            shortcut_attempted = true;
-            if let Some(assignment) = input.partition_at_lower_bound(&result, params) {
-                return CandidateSearch::Found(assignment);
-            }
-            // A cost-0 ordering that still needs extra links: refine the
-            // remaining seeds and let the full partition loop decide.
+    // Each refinement depends only on its seed, so the seeds run in parallel.
+    let mut refined: Vec<(Vec<usize>, usize)> = refinement_seeds
+        .into_par_iter()
+        .map(|seed| {
+            let result = refine_ordering(
+                seed,
+                &input.shape.statements_using,
+                &input.shape.is_original_public,
+                max_args,
+                REFINE_ITERATIONS,
+            );
+            let result_cost = input.excess_cost(&result);
+            (result, result_cost)
+        })
+        .collect();
+    // The first cost-0 refinement in seed order gets one shortcut attempt: the
+    // stable sort below keeps seed order among ties, so it is the ordering the
+    // partition loop would try first anyway. A cost-0 ordering that still
+    // needs extra links falls through to the full partition loop.
+    if let Some((zero_cost_ordering, _)) = refined.iter().find(|(_, cost)| *cost == 0) {
+        if let Some(assignment) = input.partition_at_lower_bound(zero_cost_ordering, params) {
+            return CandidateSearch::Found(assignment);
         }
-        refined.push((result, result_cost));
     }
     refined.sort_by_key(|(_, refined_cost)| *refined_cost);
 
@@ -931,12 +1003,7 @@ fn split_into_chain(
     params: &Params,
     search_cache: &mut SplitSearchCache,
 ) -> Result<(Vec<CustomPredicateDef>, SplitChainInfo), SplittingError> {
-    let original_name = pred.name.name.clone();
-    let conjunction = pred.conjunction_type;
-    let real_statement_count = pred.statements.len();
-
     let input = prepare_split_input(pred);
-    let num_statements = input.shape.num_statements;
 
     // The link assignment depends only on the statement/wildcard usage shape
     // (plus `params`, fixed for the cache's lifetime), and modules routinely
@@ -945,13 +1012,30 @@ fn split_into_chain(
     let assignment = if let Some(assignment) = search_cache.assignments.get(&input.shape) {
         assignment.clone()
     } else {
-        let assignment = search_link_assignment(&original_name, &input, params)?;
+        let assignment = search_link_assignment(&pred.name.name, &input, params)?;
         search_cache
             .assignments
             .insert(input.shape.clone(), assignment.clone());
         assignment
     };
 
+    build_chain_from_assignment(pred, &input, assignment, params)
+}
+
+/// Everything downstream of the ordering search: reorder the statements per
+/// the assignment, derive each link's argument lists, and emit the chain's
+/// predicate definitions plus [`SplitChainInfo`].
+fn build_chain_from_assignment(
+    pred: &CustomPredicateDef,
+    input: &SplitInput,
+    assignment: LinkAssignment,
+    params: &Params,
+) -> Result<(Vec<CustomPredicateDef>, SplitChainInfo), SplittingError> {
+    let original_name = pred.name.name.clone();
+    let conjunction = pred.conjunction_type;
+    let real_statement_count = pred.statements.len();
+    let num_statements = input.shape.num_statements;
+
     // Reorder map: original index -> position in flattened chain.
     let mut reorder_map = vec![0usize; num_statements];
     {

src/lang/frontend_ast_validate.rs

@@ -540,13 +540,17 @@ impl Validator {
         pred_def: &CustomPredicateDef,
     ) -> Result<(), ValidationError> {
         let pred_name = pred_def.name.name.clone();
+        let wildcard_scope = self
+            .symbols
+            .wildcard_scopes
+            .get(&pred_name)
+            .expect("Wildcard scope should exist after pass 1");
+
+        // The collision check depends only on the scope, so it runs once per
+        // predicate rather than once per statement.
+        self.validate_wildcard_names(&wildcard_scope.wildcards.keys().collect())?;
 
         for stmt in &pred_def.statements {
-            let wildcard_scope = self
-                .symbols
-                .wildcard_scopes
-                .get(&pred_name)
-                .expect("Wildcard scope should exist after pass 1");
             self.validate_statement(stmt, Some((&pred_name, wildcard_scope)))?;
         }
 
@@ -592,12 +596,6 @@ impl Validator {
         let pred_name = stmt.predicate.predicate_name();
         let pred_span = stmt.predicate.span();
 
-        let wc_names = match wildcard_context {
-            Some((_, wc_scope)) => wc_scope.wildcards.keys().collect(),
-            None => HashSet::new(),
-        };
-        self.validate_wildcard_names(&wc_names)?;
-
         // Check if predicate exists
         let pred_info = match &stmt.predicate {
             PredicateRef::Qualified { module, predicate } => {
@@ -642,7 +640,9 @@ impl Validator {
                 } else if let Some(info) = self.symbols.predicates.get(pred_name) {
                     // Custom or imported predicate
                     Some(info.clone())
-                } else if wc_names.contains(&pred_name.to_string()) {
+                } else if wildcard_context
+                    .is_some_and(|(_, scope)| scope.wildcards.contains_key(pred_name))
+                {
                     None
                 } else {
                     return Err(ValidationError::UndefinedPredicate {

src/lang/module.rs

@@ -333,9 +333,31 @@ fn build_single_batch(
     params: &Params,
     batch_name: &str,
 ) -> Result<Arc<CustomPredicateBatch>, BatchingError> {
+    use rayon::prelude::*;
+
     let mut builder = CustomPredicateBatchBuilder::new(params.clone(), batch_name.to_string());
 
-    for pred in predicates {
+    // Statement-template construction is a pure function of the symbol table
+    // and reference map, so it runs in parallel; only the batch insertion
+    // below stays sequential.
+    let prepared: Vec<Vec<StatementTmplBuilder>> = predicates
+        .par_iter()
+        .map(|pred| {
+            pred.statements
+                .iter()
+                .map(|stmt| {
+                    build_statement_with_resolved_refs(
+                        stmt,
+                        reference_map,
+                        &pred.name.name,
+                        symbols,
+                    )
+                })
+                .collect::<Result<_, _>>()
+        })
+        .collect::<Result<_, _>>()?;
+
+    for (pred, statement_builders) in predicates.iter().zip(prepared) {
         let name = &pred.name.name;
 
         // Collect argument names
@@ -353,13 +375,6 @@ fn build_single_batch(
             .map(|args| args.iter().map(|a| a.name.as_str()).collect())
             .unwrap_or_default();
 
-        // Build statement templates with resolved predicates
-        let statement_builders: Vec<StatementTmplBuilder> = pred
-            .statements
-            .iter()
-            .map(|stmt| build_statement_with_resolved_refs(stmt, reference_map, name, symbols))
-            .collect::<Result<_, _>>()?;
-
         let conjunction = pred.conjunction_type == ConjunctionType::And;
 
         builder

src/middleware/custom.rs

@@ -466,8 +466,9 @@ impl fmt::Debug for CustomPredicateBatchData {
 // TODO: Rename Batch for Module everywhere in the code base
 impl CustomPredicateBatchData {
     fn new_full(predicates: Vec<CustomPredicate>) -> Self {
+        use rayon::prelude::*;
         let kvs: HashMap<RawValue, RawValue> = predicates
-            .iter()
+            .par_iter()
             .enumerate()
             .map(|(index, pred)| {
                 let cp_hash = hash_fields(&pred.to_fields());