perf: normalize single-external query var order

MilkBlock · MilkBlock · commit a52475df33ff · 2026-04-12T12:00:05.000+08:00
diff --git a/core-relations/src/query.rs b/core-relations/src/query.rs
@@ -1,6 +1,6 @@
 //! APIs for building a query of a database.
 
-use std::{iter::once, sync::Arc};
+use std::{iter::once, mem, sync::Arc};
 
 use crate::numeric_id::{DenseIdMap, IdVec, NumericId, define_id};
 use smallvec::SmallVec;
@@ -444,7 +444,60 @@ impl RuleBuilder<'_, '_> {
         }
     }
 
+    fn normalize_single_external_arg_order(&mut self) {
+        let Some(arg_vars) = single_external_arg_vars(&self.qb.instrs) else {
+            return;
+        };
+
+        let mut preferred = SmallVec::<[Variable; 4]>::new();
+        preferred.extend(arg_vars.iter().copied());
+        preferred.extend(
+            self.qb
+                .query
+                .var_info
+                .iter()
+                .map(|(var, _)| var)
+                .filter(|var| !arg_vars.contains(var)),
+        );
+
+        let current = SmallVec::<[Variable; 4]>::from_iter(
+            self.qb.query.var_info.iter().map(|(var, _)| var),
+        );
+        if preferred == current {
+            return;
+        }
+
+        let mut remap = DenseIdMap::with_capacity(self.qb.query.var_info.n_ids());
+        for (idx, old_var) in preferred.iter().enumerate() {
+            remap.insert(*old_var, Variable::from_usize(idx));
+        }
+
+        let mut old_var_info = mem::take(&mut self.qb.query.var_info);
+        for old_var in preferred {
+            self.qb
+                .query
+                .var_info
+                .push(old_var_info.take(old_var).expect("all vars must be present"));
+        }
+
+        for (_, atom) in self.qb.query.atoms.iter_mut() {
+            let old_var_to_column = mem::take(&mut atom.var_to_column);
+            atom.var_to_column = old_var_to_column
+                .into_iter()
+                .map(|(var, col)| (remap[var], col))
+                .collect();
+            for (_, var) in atom.column_to_var.iter_mut() {
+                *var = remap[*var];
+            }
+        }
+
+        for instr in self.qb.instrs.iter_mut() {
+            remap_instr(instr, &remap);
+        }
+    }
+
     pub fn build_with_description(mut self, desc: impl Into<String>) -> RuleId {
+        self.normalize_single_external_arg_order();
         let var_info = &self.qb.query.var_info;
         let symbol_map = self.build_symbol_map();
         // Generate an id for our actions and slot them in.
@@ -769,6 +822,127 @@ impl RuleBuilder<'_, '_> {
     }
 }
 
+fn single_external_arg_vars(instrs: &[Instr]) -> Option<SmallVec<[Variable; 4]>> {
+    let [Instr::External { args, .. }] = instrs else {
+        return None;
+    };
+
+    let mut vars = SmallVec::<[Variable; 4]>::new();
+    for arg in args {
+        let QueryEntry::Var(var) = arg else {
+            return None;
+        };
+        if vars.contains(var) {
+            return None;
+        }
+        vars.push(*var);
+    }
+    Some(vars)
+}
+
+fn remap_query_entry(entry: &mut QueryEntry, remap: &DenseIdMap<Variable, Variable>) {
+    if let QueryEntry::Var(var) = entry {
+        *var = remap[*var];
+    }
+}
+
+fn remap_write_val(val: &mut WriteVal, remap: &DenseIdMap<Variable, Variable>) {
+    if let WriteVal::QueryEntry(entry) = val {
+        remap_query_entry(entry, remap);
+    }
+}
+
+fn remap_instr(instr: &mut Instr, remap: &DenseIdMap<Variable, Variable>) {
+    match instr {
+        Instr::LookupOrInsertDefault {
+            args,
+            default,
+            dst_var,
+            ..
+        } => {
+            args.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            default
+                .iter_mut()
+                .for_each(|val| remap_write_val(val, remap));
+            *dst_var = remap[*dst_var];
+        }
+        Instr::LookupWithDefault {
+            args,
+            default,
+            dst_var,
+            ..
+        } => {
+            args.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            remap_query_entry(default, remap);
+            *dst_var = remap[*dst_var];
+        }
+        Instr::Lookup {
+            args, dst_var, ..
+        } => {
+            args.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            *dst_var = remap[*dst_var];
+        }
+        Instr::LookupWithFallback {
+            table_key,
+            func_args,
+            dst_var,
+            ..
+        } => {
+            table_key
+                .iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            func_args
+                .iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            *dst_var = remap[*dst_var];
+        }
+        Instr::Insert { vals, .. } => {
+            vals.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+        }
+        Instr::InsertIfEq { l, r, vals, .. } => {
+            remap_query_entry(l, remap);
+            remap_query_entry(r, remap);
+            vals.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+        }
+        Instr::Remove { args, .. } => {
+            args.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+        }
+        Instr::External { args, dst, .. } => {
+            args.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            *dst = remap[*dst];
+        }
+        Instr::ExternalWithFallback {
+            args1, args2, dst, ..
+        } => {
+            args1
+                .iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            args2
+                .iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+            *dst = remap[*dst];
+        }
+        Instr::AssertEq(l, r) | Instr::AssertNe(l, r) => {
+            remap_query_entry(l, remap);
+            remap_query_entry(r, remap);
+        }
+        Instr::AssertAnyNe { ops, .. } => {
+            ops.iter_mut()
+                .for_each(|entry| remap_query_entry(entry, remap));
+        }
+        Instr::ReadCounter { dst, .. } => {
+            *dst = remap[*dst];
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct Atom {
     pub(crate) table: TableId,
diff --git a/core-relations/src/tests.rs b/core-relations/src/tests.rs
@@ -122,6 +122,51 @@ fn basic_query() {
     assert_eq!(res, Vec::from_iter((0..10).chain([13, 17].into_iter())));
 }
 
+#[test]
+fn single_external_rule_normalizes_gap_vars_into_consecutive_window() {
+    let mut db = Database::default();
+    let tuple4 = db.add_table(
+        SortedWritesTable::new(4, 4, None, vec![], Box::new(|_, _, _, _| false)),
+        iter::empty(),
+        iter::empty(),
+    );
+    let echo = db.add_external_function(Box::new(make_external_func(|_, args| args.first().copied())));
+
+    let mut rsb = RuleSetBuilder::new(&mut db);
+    let mut query = rsb.new_rule();
+    let a = query.new_var_named("a");
+    let b = query.new_var_named("b");
+    let c = query.new_var_named("c");
+    let d = query.new_var_named("d");
+    query
+        .add_atom(tuple4, &[a.into(), b.into(), c.into(), d.into()], &[])
+        .unwrap();
+    let mut rule = query.build();
+    let _ = rule.call_external(echo, &[a.into(), b.into(), d.into()]).unwrap();
+    let rule_id = rule.build_with_description("single_external_gap");
+    let rule_set = rsb.build();
+
+    let (_, _, _, action_id) = rule_set.plans.get(rule_id).unwrap();
+    let action = &rule_set.actions[*action_id];
+    let [crate::action::Instr::External { args, .. }] = action.instrs.as_ref().as_slice() else {
+        panic!("expected single external instruction");
+    };
+
+    let vars: Vec<_> = args
+        .iter()
+        .map(|arg| match arg {
+            crate::action::QueryEntry::Var(var) => *var,
+            crate::action::QueryEntry::Const(_) => panic!("expected var-only external args"),
+        })
+        .collect();
+
+    assert_eq!(
+        vars,
+        vec![crate::free_join::Variable::from_usize(0), crate::free_join::Variable::from_usize(1), crate::free_join::Variable::from_usize(2)],
+        "eligible single-external rules should normalize arg vars into a consecutive window",
+    );
+}
+
 #[test]
 fn line_graph_1_fj_puresize() {
     line_graph_1_test(PlanStrategy::PureSize);
