Optimize primitive arg layout fast paths

MilkBlock · MilkBlock · commit 5bb2cb40d5cc · 2026-04-12T10:03:46.000+08:00
diff --git a/core-relations/src/action/mod.rs b/core-relations/src/action/mod.rs
@@ -22,7 +22,7 @@ use crate::{
     BaseValues, ContainerValues, ExternalFunctionId, WrappedTable,
     common::Value,
     free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable},
-    pool::{Clear, Pooled, with_pool_set},
+    pool::{Clear, Pool, Pooled, with_pool_set},
     table_spec::{ColumnId, MutationBuffer},
 };
 
@@ -55,6 +55,12 @@ impl From<Value> for QueryEntry {
     }
 }
 
+#[derive(Debug, Clone)]
+pub(crate) enum ArgBindings {
+    Owned(Vec<QueryEntry>),
+    BorrowedWindow { start: Variable, len: usize },
+}
+
 /// A value that can be written to a table in an action.
 #[derive(Debug, Clone, Copy)]
 pub enum WriteVal {
@@ -279,6 +285,75 @@ impl Bindings {
     }
 }
 
+fn invoke_batch_from_args(
+    this: &dyn crate::free_join::ExternalFunction,
+    state: &mut ExecutionState,
+    mask: &mut Mask,
+    bindings: &mut Bindings,
+    args: &ArgBindings,
+    out_var: Variable,
+) {
+    match args {
+        ArgBindings::Owned(args) => invoke_batch(this, state, mask, bindings, args, out_var),
+        ArgBindings::BorrowedWindow { start, len } => {
+            let pool: Pool<Vec<Value>> = with_pool_set(|ps| ps.get_pool().clone());
+            let mut out = pool.get();
+            out.reserve(mask.len());
+            let start_index = start.index();
+            let arity = *len;
+            let iter = mask.iter_dynamic(
+                with_pool_set(crate::pool::PoolSet::get_pool),
+                (0..arity)
+                    .map(|i| ValueSource::Slice(&bindings[Variable::from_usize(start_index + i)])),
+            );
+            iter.fill_vec(&mut out, Value::stale, |_, args| {
+                this.invoke(state, args.as_slice())
+            });
+            bindings.insert(out_var, &out);
+        }
+    }
+}
+
+fn invoke_batch_assign_from_args(
+    this: &dyn crate::free_join::ExternalFunction,
+    state: &mut ExecutionState,
+    mask: &mut Mask,
+    bindings: &mut Bindings,
+    args: &ArgBindings,
+    out_var: Variable,
+) {
+    match args {
+        ArgBindings::Owned(args) => invoke_batch_assign(this, state, mask, bindings, args, out_var),
+        ArgBindings::BorrowedWindow { start, len } => {
+            let mut out = bindings.take(out_var).expect("out_var must be bound");
+            let start_index = start.index();
+            let arity = *len;
+            let iter = mask.iter_dynamic(
+                with_pool_set(crate::pool::PoolSet::get_pool),
+                (0..arity)
+                    .map(|i| ValueSource::Slice(&bindings[Variable::from_usize(start_index + i)])),
+            );
+            iter.assign_vec_and_retain(&mut out.vals, |_, args| this.invoke(state, &args));
+            bindings.replace(out);
+        }
+    }
+}
+
+fn single_external_borrowed_window(
+    instrs: &[Instr],
+) -> Option<(ExternalFunctionId, Variable, usize, Variable)> {
+    match instrs {
+        [
+            Instr::External {
+                func,
+                args: ArgBindings::BorrowedWindow { start, len },
+                dst,
+            },
+        ] => Some((*func, *start, *len, *dst)),
+        _ => None,
+    }
+}
+
 /// A binding that has been extracted from a [`Bindings`] struct via the [`Bindings::take`] method.
 ///
 /// This allows for a variable's contents to be read while the [`Bindings`] struct has been
@@ -603,6 +678,33 @@ impl ExecutionState<'_> {
             bindings.matches = 1;
         }
 
+        if let Some((func, start, len, dst)) = single_external_borrowed_window(instrs) {
+            let matches = bindings.matches;
+            let mut out: Pooled<Vec<Value>> = with_pool_set(|ps| ps.get());
+            out.resize(matches, Value::stale());
+            let mut transposed: Pooled<Vec<Value>> = with_pool_set(|ps| ps.get());
+            transposed.resize(matches * len, Value::stale());
+            let start_index = start.index();
+            let mut succeeded = 0usize;
+            for row in 0..matches {
+                for col in 0..len {
+                    transposed[row * len + col] =
+                        bindings[Variable::from_usize(start_index + col)][row];
+                }
+                if self.should_stop() {
+                    break;
+                }
+                if let Some(value) =
+                    self.call_external_func(func, &transposed[row * len..(row + 1) * len])
+                {
+                    out[row] = value;
+                    succeeded += 1;
+                }
+            }
+            bindings.insert(dst, &out);
+            return succeeded;
+        }
+
         // Vectorized execution for larger batch sizes
         let mut mask = with_pool_set(|ps| Mask::new(0..bindings.matches, ps));
         for instr in instrs {
@@ -758,7 +860,7 @@ impl ExecutionState<'_> {
                 }
 
                 // Call the given external function on all entries where the lookup failed.
-                invoke_batch_assign(
+                invoke_batch_assign_from_args(
                     self.db.external_funcs[*func].as_ref(),
                     self,
                     &mut to_call_func,
@@ -817,7 +919,7 @@ impl ExecutionState<'_> {
                 });
             }
             Instr::External { func, args, dst } => {
-                invoke_batch(
+                invoke_batch_from_args(
                     self.db.external_funcs[*func].as_ref(),
                     self,
                     mask,
@@ -834,7 +936,7 @@ impl ExecutionState<'_> {
                 dst,
             } => {
                 let mut f1_result = mask.clone();
-                invoke_batch(
+                invoke_batch_from_args(
                     self.db.external_funcs[*f1].as_ref(),
                     self,
                     &mut f1_result,
@@ -848,7 +950,7 @@ impl ExecutionState<'_> {
                     return;
                 }
                 // Call the given external function on all entries where the first call failed.
-                invoke_batch_assign(
+                invoke_batch_assign_from_args(
                     self.db.external_funcs[*f2].as_ref(),
                     self,
                     &mut to_call_f2,
@@ -921,7 +1023,7 @@ pub(crate) enum Instr {
         table: TableId,
         table_key: Vec<QueryEntry>,
         func: ExternalFunctionId,
-        func_args: Vec<QueryEntry>,
+        func_args: ArgBindings,
         dst_col: ColumnId,
         dst_var: Variable,
     },
@@ -950,7 +1052,7 @@ pub(crate) enum Instr {
     /// Bind the result of the external function to a variable.
     External {
         func: ExternalFunctionId,
-        args: Vec<QueryEntry>,
+        args: ArgBindings,
         dst: Variable,
     },
 
@@ -959,9 +1061,9 @@ pub(crate) enum Instr {
     /// single failure of `External`).
     ExternalWithFallback {
         f1: ExternalFunctionId,
-        args1: Vec<QueryEntry>,
+        args1: ArgBindings,
         f2: ExternalFunctionId,
-        args2: Vec<QueryEntry>,
+        args2: ArgBindings,
         dst: Variable,
     },
 
diff --git a/core-relations/src/action/tests.rs b/core-relations/src/action/tests.rs
@@ -1,5 +1,7 @@
 use crate::{
     action::mask::{IterResult, ValueSource},
+    free_join::Variable,
+    numeric_id::NumericId,
     pool::{PoolSet, with_pool_set},
 };
 
@@ -148,3 +150,40 @@ fn test_early_stop_multiple_clones() {
     assert!(state2.should_stop());
     assert!(state3.should_stop());
 }
+
+#[test]
+fn run_instrs_single_external_borrowed_window_executes() {
+    let mut db = crate::free_join::Database::default();
+    let sum = db.add_external_function(Box::new(crate::make_external_func(|_, args| {
+        let [x, y] = args else { panic!() };
+        Some(crate::common::Value::from_usize(
+            (x.rep() + y.rep()) as usize,
+        ))
+    })));
+    let mut state = crate::action::ExecutionState::new(db.read_only_view(), Default::default());
+    let mut bindings = super::Bindings::new(8);
+    bindings.insert(
+        Variable::from_usize(0),
+        &[crate::table_shortcuts::v(1), crate::table_shortcuts::v(2)],
+    );
+    bindings.insert(
+        Variable::from_usize(1),
+        &[crate::table_shortcuts::v(10), crate::table_shortcuts::v(20)],
+    );
+
+    let instrs = vec![super::Instr::External {
+        func: sum,
+        args: super::ArgBindings::BorrowedWindow {
+            start: Variable::from_usize(0),
+            len: 2,
+        },
+        dst: Variable::from_usize(2),
+    }];
+
+    let succeeded = state.run_instrs(&instrs, &mut bindings);
+    assert_eq!(succeeded, 2);
+    assert_eq!(
+        &bindings[Variable::from_usize(2)],
+        &[crate::table_shortcuts::v(11), crate::table_shortcuts::v(22)]
+    );
+}
diff --git a/core-relations/src/query.rs b/core-relations/src/query.rs
@@ -8,7 +8,7 @@ use thiserror::Error;
 
 use crate::{
     BaseValueId, CounterId, ExternalFunctionId, PoolSet,
-    action::{Instr, QueryEntry, WriteVal},
+    action::{ArgBindings, Instr, QueryEntry, WriteVal},
     common::HashMap,
     free_join::{
         ActionId, AtomId, Database, ProcessedConstraints, SubAtom, TableId, TableInfo, VarInfo,
@@ -243,6 +243,18 @@ impl<'outer, 'a> QueryBuilder<'outer, 'a> {
         }
     }
 
+    fn mark_used_arg_bindings(&mut self, entries: &ArgBindings) {
+        match entries {
+            ArgBindings::Owned(entries) => self.mark_used(entries),
+            ArgBindings::BorrowedWindow { start, len } => {
+                for offset in 0..*len {
+                    self.query.var_info[Variable::from_usize(start.index() + offset)].used_in_rhs =
+                        true;
+                }
+            }
+        }
+    }
+
     fn mark_defined(&mut self, entry: &QueryEntry) {
         // TODO: use some of this information in query planning, e.g. dedup at match time.
         if let QueryEntry::Var(v) = entry {
@@ -622,14 +634,32 @@ impl RuleBuilder<'_, '_> {
         let res = self.qb.new_var();
         self.qb.instrs.push(Instr::External {
             func,
-            args: args.to_vec(),
+            args: ArgBindings::Owned(args.to_vec()),
             dst: res,
         });
         self.qb.mark_used(args);
         self.qb.mark_defined(&res.into());
         Ok(res)
     }
 
+    pub fn call_external_window(
+        &mut self,
+        func: ExternalFunctionId,
+        start: Variable,
+        len: usize,
+    ) -> Result<Variable, QueryError> {
+        let res = self.qb.new_var();
+        let args = ArgBindings::BorrowedWindow { start, len };
+        self.qb.instrs.push(Instr::External {
+            func,
+            args: args.clone(),
+            dst: res,
+        });
+        self.qb.mark_used_arg_bindings(&args);
+        self.qb.mark_defined(&res.into());
+        Ok(res)
+    }
+
     /// Look up the given key in the given table. If the lookup fails, then call the given external
     /// function with the given arguments. Bind the result to the returned variable. If the
     /// external function returns None (and the lookup fails) then the execution of the rule halts.
@@ -648,7 +678,7 @@ impl RuleBuilder<'_, '_> {
             table,
             table_key: key.to_vec(),
             func,
-            func_args: func_args.to_vec(),
+            func_args: ArgBindings::Owned(func_args.to_vec()),
             dst_var: res,
             dst_col,
         });
@@ -668,9 +698,9 @@ impl RuleBuilder<'_, '_> {
         let res = self.qb.new_var();
         self.qb.instrs.push(Instr::ExternalWithFallback {
             f1,
-            args1: args1.to_vec(),
+            args1: ArgBindings::Owned(args1.to_vec()),
             f2,
-            args2: args2.to_vec(),
+            args2: ArgBindings::Owned(args2.to_vec()),
             dst: res,
         });
         self.qb.mark_used(args1);
@@ -679,6 +709,32 @@ impl RuleBuilder<'_, '_> {
         Ok(res)
     }
 
+    pub fn call_external_with_fallback_window(
+        &mut self,
+        f1: ExternalFunctionId,
+        start1: Variable,
+        len1: usize,
+        f2: ExternalFunctionId,
+        args2: &[QueryEntry],
+    ) -> Result<Variable, QueryError> {
+        let res = self.qb.new_var();
+        let args1 = ArgBindings::BorrowedWindow {
+            start: start1,
+            len: len1,
+        };
+        self.qb.instrs.push(Instr::ExternalWithFallback {
+            f1,
+            args1: args1.clone(),
+            f2,
+            args2: ArgBindings::Owned(args2.to_vec()),
+            dst: res,
+        });
+        self.qb.mark_used_arg_bindings(&args1);
+        self.qb.mark_used(args2);
+        self.qb.mark_defined(&res.into());
+        Ok(res)
+    }
+
     /// Continue execution iff the two arguments are equal.
     pub fn assert_eq(&mut self, l: QueryEntry, r: QueryEntry) {
         self.qb.instrs.push(Instr::AssertEq(l, r));
diff --git a/egglog-bridge/src/rule.rs b/egglog-bridge/src/rule.rs
