Optimize primitive arg layout fast paths

Author: MilkBlock

core-relations/src/action/mod.rs

@@ -279,6 +279,34 @@ impl Bindings {
     }
 }
 
+fn single_external_borrowed_window(
+    instrs: &[Instr],
+) -> Option<(ExternalFunctionId, Variable, usize, Variable)> {
+    match instrs {
+        [
+            Instr::External {
+                func,
+                args,
+                dst,
+            },
+        ] => {
+            let mut start = None;
+            for (idx, arg) in args.iter().enumerate() {
+                let QueryEntry::Var(v) = arg else {
+                    return None;
+                };
+                if idx == 0 {
+                    start = Some(*v);
+                } else if v.index() != start?.index() + idx {
+                    return None;
+                }
+            }
+            start.map(|start| (*func, start, args.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 +631,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 {

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,37 @@ 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: vec![Variable::from_usize(0).into(), Variable::from_usize(1).into()],
+        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)]
+    );
+}