Skip to main content

egglog_core_relations/action/
mod.rs

1//! Instructions that are executed on the results of a query.
2//!
3//! This allows us to execute the "right-hand-side" of a rule. The
4//! implementation here is optimized to execute on a batch of rows at a time.
5use std::{
6    any::Any,
7    ops::Deref,
8    sync::{
9        Arc,
10        atomic::{AtomicBool, Ordering},
11    },
12};
13
14use crate::{
15    common::HashMap,
16    free_join::{get_column_index_from_tableinfo, invoke_batch, invoke_batch_assign},
17    numeric_id::{DenseIdMap, NumericId},
18};
19use egglog_concurrency::NotificationList;
20use smallvec::SmallVec;
21
22use crate::{
23    BaseValues, ContainerValues, ExternalFunctionId, Offset, WrappedTable,
24    common::Value,
25    free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable},
26    offsets::Subset,
27    pool::{Clear, Pooled, with_pool_set},
28    row_buffer::TaggedRowBuffer,
29    table_spec::{ColumnId, Constraint, MutationBuffer},
30};
31
32use self::mask::{Mask, MaskIter, ValueSource};
33
34#[macro_use]
35pub(crate) mod mask;
36
37#[cfg(test)]
38mod tests;
39
40/// A representation of a value within a query or rule.
41///
42/// A QueryEntry is either a variable bound in a join, or an untyped constant.
43#[derive(Copy, Clone, Debug)]
44pub enum QueryEntry {
45    Var(Variable),
46    Const(Value),
47}
48
49impl From<Variable> for QueryEntry {
50    fn from(var: Variable) -> Self {
51        QueryEntry::Var(var)
52    }
53}
54
55impl From<Value> for QueryEntry {
56    fn from(val: Value) -> Self {
57        QueryEntry::Const(val)
58    }
59}
60
61/// A value that can be written to a table in an action.
62#[derive(Debug, Clone, Copy)]
63pub enum WriteVal {
64    /// A variable or a constant.
65    QueryEntry(QueryEntry),
66    /// A fresh value from the given counter.
67    IncCounter(CounterId),
68    /// The value of the current row index.
69    CurrentVal(usize),
70}
71
72impl<T> From<T> for WriteVal
73where
74    T: Into<QueryEntry>,
75{
76    fn from(val: T) -> Self {
77        WriteVal::QueryEntry(val.into())
78    }
79}
80
81impl From<CounterId> for WriteVal {
82    fn from(ctr: CounterId) -> Self {
83        WriteVal::IncCounter(ctr)
84    }
85}
86
87/// A value that can be written to the database during a merge action.
88#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
89pub enum MergeVal {
90    /// A fresh value from the given counter.
91    Counter(CounterId),
92    /// A standard constant value.
93    Constant(Value),
94}
95
96impl From<CounterId> for MergeVal {
97    fn from(ctr: CounterId) -> Self {
98        MergeVal::Counter(ctr)
99    }
100}
101
102impl From<Value> for MergeVal {
103    fn from(val: Value) -> Self {
104        MergeVal::Constant(val)
105    }
106}
107
108/// Bindings store a sequence of values for a given set of variables.
109///
110/// The intent of bindings is to store a sequence of mappings from [`Variable`] to [`Value`], in a
111/// struct-of-arrays style that is better laid out for processing bindings in batches.
112pub(crate) struct Bindings {
113    matches: usize,
114    /// The maximum number of calls to `push` that we can receive before we clear the
115    /// [`Bindings`].
116    // This is used to preallocate chunks of the flat `data` vector.
117    max_batch_size: usize,
118    data: Pooled<Vec<Value>>,
119    /// Points into `data`. `data[vars[var].. vars[var]+matches]` contains the values for `data`.
120    var_offsets: DenseIdMap<Variable, usize>,
121}
122
123impl std::ops::Index<Variable> for Bindings {
124    type Output = [Value];
125    fn index(&self, var: Variable) -> &[Value] {
126        self.get(var).unwrap()
127    }
128}
129
130impl std::fmt::Debug for Bindings {
131    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        let mut table = f.debug_map();
133        for (var, start) in self.var_offsets.iter() {
134            table.entry(&var, &&self.data[*start..*start + self.matches]);
135        }
136        table.finish()
137    }
138}
139
140impl Bindings {
141    pub(crate) fn new(max_batch_size: usize) -> Self {
142        Bindings {
143            matches: 0,
144            max_batch_size,
145            data: Default::default(),
146            var_offsets: DenseIdMap::new(),
147        }
148    }
149    fn assert_invariant(&self) {
150        #[cfg(debug_assertions)]
151        {
152            assert!(self.matches <= self.max_batch_size);
153            for (var, start) in self.var_offsets.iter() {
154                assert!(
155                    start + self.matches <= self.data.len(),
156                    "Variable {:?} starts at {}, but data only has {} elements",
157                    var,
158                    start,
159                    self.data.len()
160                );
161            }
162        }
163    }
164
165    pub(crate) fn clear(&mut self) {
166        self.matches = 0;
167        self.var_offsets.clear();
168        self.data.clear();
169        self.assert_invariant();
170    }
171
172    fn get(&self, var: Variable) -> Option<&[Value]> {
173        let start = self.var_offsets.get(var)?;
174        Some(&self.data[*start..*start + self.matches])
175    }
176
177    fn add_mapping(&mut self, var: Variable, vals: &[Value]) {
178        let start = self.data.len();
179        self.data.extend_from_slice(vals);
180        // We have a flat representation of the data, meaning that writing more than
181        // `max_batch_size` values to `var` could overwrite values for a different variable, which
182        // would produce some mysterious results that are hard to debug.
183        debug_assert!(vals.len() <= self.max_batch_size);
184        if vals.len() < self.max_batch_size {
185            let target_len = self.data.len() + self.max_batch_size - vals.len();
186            self.data.resize(target_len, Value::stale());
187        }
188        self.var_offsets.insert(var, start);
189    }
190
191    pub(crate) fn insert(&mut self, var: Variable, vals: &[Value]) {
192        if self.var_offsets.n_ids() == 0 {
193            self.matches = vals.len();
194        } else {
195            assert_eq!(self.matches, vals.len());
196        }
197        self.add_mapping(var, vals);
198        self.assert_invariant();
199    }
200
201    /// Push a new set of bindings for the given variables.
202    ///
203    /// # Safety:
204    /// This method assumes that all calls to `push`:
205    /// * Have a mapping for every member of `used_vars`.
206    /// * Are passed the same `used_vars`.
207    ///
208    /// It is unsafe to avoid bounds-checking. This method is called extremely frequently and the
209    /// overhead of boundschecking is noticeable.
210    pub(crate) unsafe fn push(
211        &mut self,
212        map: &DenseIdMap<Variable, Value>,
213        used_vars: &[Variable],
214    ) {
215        if self.matches != 0 {
216            assert!(self.matches < self.max_batch_size);
217            #[cfg(debug_assertions)]
218            {
219                for var in used_vars {
220                    assert!(
221                        self.var_offsets.get(*var).is_some(),
222                        "Variable {:?} not found in bindings {:?}",
223                        var,
224                        self.var_offsets
225                    );
226                }
227            }
228            for var in used_vars {
229                let var = var.index();
230                // Safe version: this degrades some benchmarks by ~6%
231                // let start = self.var_offsets.raw()[var].unwrap();
232                // self.data[start + self.matches] = map.raw()[var].unwrap();
233                unsafe {
234                    let start = self.var_offsets.raw().get_unchecked(var).unwrap_unchecked();
235                    *self.data.get_unchecked_mut(start + self.matches) =
236                        map.raw().get_unchecked(var).unwrap_unchecked();
237                }
238            }
239        } else {
240            for (var, val) in map.iter() {
241                self.add_mapping(var, &[*val]);
242            }
243        }
244
245        self.matches += 1;
246        self.assert_invariant();
247    }
248
249    /// A method that removes the bindings for the given variable and allows for its values to be
250    /// used independently from the [`Bindings`] struct. This is helpful when an operation needs to
251    /// mutably borrow the values for one value while reading the values for another.
252    ///
253    /// To add the values back, use [`Bindings::replace`].
254    pub(crate) fn take(&mut self, var: Variable) -> Option<ExtractedBinding> {
255        let mut vals: Pooled<Vec<Value>> = with_pool_set(|ps| ps.get());
256        vals.extend_from_slice(self.get(var)?);
257        let start = self.var_offsets.take(var)?;
258        Some(ExtractedBinding {
259            var,
260            offset: start,
261            vals,
262        })
263    }
264
265    /// Replace a binding extracted with [`Bindings::take`].
266    ///
267    /// # Panics
268    /// This method will panic if the length of the values in `bdg` does not match the current
269    /// number of matches in `Bindings`. It may panic if `bdg` was extracted from a different
270    /// [`Bindings`] than the one it is being replaced in.
271    pub(crate) fn replace(&mut self, bdg: ExtractedBinding) {
272        // Replace the binding with the new values.
273        let ExtractedBinding {
274            var,
275            offset,
276            mut vals,
277        } = bdg;
278        assert_eq!(vals.len(), self.matches);
279        self.data
280            .splice(offset..offset + self.matches, vals.drain(..));
281        self.var_offsets.insert(var, offset);
282    }
283}
284
285/// A binding that has been extracted from a [`Bindings`] struct via the [`Bindings::take`] method.
286///
287/// This allows for a variable's contents to be read while the [`Bindings`] struct has been
288/// borrowed mutably. The contents will not be readable until [`Bindings::replace`] is called.
289pub(crate) struct ExtractedBinding {
290    var: Variable,
291    offset: usize,
292    pub(crate) vals: Pooled<Vec<Value>>,
293}
294
295#[derive(Default)]
296pub(crate) struct PredictedVals {
297    #[allow(clippy::type_complexity)]
298    data: HashMap<(TableId, SmallVec<[Value; 3]>), Pooled<Vec<Value>>>,
299}
300
301impl Clear for PredictedVals {
302    fn reuse(&self) -> bool {
303        self.data.capacity() > 0
304    }
305    fn clear(&mut self) {
306        self.data.clear()
307    }
308    fn bytes(&self) -> usize {
309        self.data.capacity()
310            * (std::mem::size_of::<(TableId, SmallVec<[Value; 3]>)>()
311                + std::mem::size_of::<Pooled<Vec<Value>>>())
312    }
313}
314
315impl PredictedVals {
316    pub(crate) fn get_val(
317        &mut self,
318        table: TableId,
319        key: &[Value],
320        default: impl FnOnce() -> Pooled<Vec<Value>>,
321    ) -> impl Deref<Target = Pooled<Vec<Value>>> + '_ {
322        self.data
323            .entry((table, SmallVec::from_slice(key)))
324            .or_insert_with(default)
325    }
326}
327
328#[derive(Copy, Clone)]
329pub(crate) struct DbView<'a> {
330    pub(crate) external_context: ExternalContext<'a>,
331    pub(crate) table_info: &'a DenseIdMap<TableId, TableInfo>,
332    pub(crate) counters: &'a Counters,
333    pub(crate) external_funcs: &'a ExternalFunctions,
334    pub(crate) bases: &'a BaseValues,
335    pub(crate) containers: &'a ContainerValues,
336    pub(crate) notification_list: &'a NotificationList<TableId>,
337}
338
339/// A borrowed value an embedder can make visible to every [`ExecutionState`]
340/// created for one operation, for its [`ExternalFunction`]s to read back with
341/// [`ExecutionState::external_context`].
342///
343/// It is borrowed for exactly the operation that supplied it, so an embedder
344/// cannot mutate the state it shared while that operation runs.
345///
346/// [`ExternalFunction`]: crate::ExternalFunction
347pub type ExternalContext<'a> = Option<&'a (dyn Any + Send + Sync)>;
348
349/// A handle on a database that may be in the process of running a rule.
350///
351/// An ExecutionState grants immutable access to the (much of) the database, and also provides a
352/// limited API to mutate database contents.
353///
354/// A few important notes:
355///
356/// ## Some tables may be missing
357/// Callers external to this crate cannot construct an `ExecutionState` directly. Depending on the
358/// context, some tables may not be available. In particular, when running [`crate::Table::merge`]
359/// operations, only a table's read-side dependencies are available for reading (sim. for writing).
360/// This allows tables that do not need access to one another to be merged in parallel.
361///
362/// When executing a rule, ExecutionState has access to all tables.
363///
364/// ## Limited Mutability
365/// Callers can only stage insertsions and deletions to tables. These changes are not applied until
366/// the next call to `merge` on the underlying table.
367///
368/// ## Predicted Values
369/// ExecutionStates provide a means of synchronizing the results of a pending write across
370/// different executions of a rule. This is particularly important in the case where the result of
371/// an operation (such as "lookup or insert new id" operatiosn) is a fresh id. A common
372/// ExecutionState ensures that future lookups will see the same id (even across calls to
373/// [`ExecutionState::clone`]).
374pub struct ExecutionState<'a> {
375    pub(crate) predicted: PredictedVals,
376    pub(crate) db: DbView<'a>,
377    buffers: MutationBuffers<'a>,
378    /// Whether any mutations have been staged via this ExecutionState.
379    pub(crate) changed: bool,
380    /// Atomic flag for early stopping of rule execution.
381    /// This flag is shared across all handles (clones) of this ExecutionState.
382    stop_match: Arc<AtomicBool>,
383}
384
385/// A basic wrapper around an map from table id to a mutation buffer for that table that also
386/// tracks if a table has been modified.
387struct MutationBuffers<'a> {
388    notify_list: &'a NotificationList<TableId>,
389    buffers: DenseIdMap<TableId, Box<dyn MutationBuffer>>,
390}
391
392impl Clone for MutationBuffers<'_> {
393    fn clone(&self) -> Self {
394        let mut res = MutationBuffers::new(self.notify_list, Default::default());
395        for (id, buf) in self.buffers.iter() {
396            res.buffers.insert(id, buf.fresh_handle());
397        }
398        res
399    }
400}
401
402impl<'a> MutationBuffers<'a> {
403    fn new(
404        notify_list: &'a NotificationList<TableId>,
405        buffers: DenseIdMap<TableId, Box<dyn MutationBuffer>>,
406    ) -> MutationBuffers<'a> {
407        MutationBuffers {
408            notify_list,
409            buffers,
410        }
411    }
412    fn lazy_init(&mut self, table_id: TableId, f: impl FnOnce() -> Box<dyn MutationBuffer>) {
413        self.buffers.get_or_insert(table_id, f);
414    }
415    fn stage_insert(&mut self, table_id: TableId, row: &[Value]) {
416        self.buffers[table_id].stage_insert(row);
417        self.notify_list.notify(table_id);
418    }
419
420    fn stage_remove(&mut self, table_id: TableId, key: &[Value]) {
421        self.buffers[table_id].stage_remove(key);
422        self.notify_list.notify(table_id);
423    }
424}
425
426impl Clone for ExecutionState<'_> {
427    fn clone(&self) -> Self {
428        ExecutionState {
429            predicted: Default::default(),
430            db: self.db,
431            buffers: self.buffers.clone(),
432            changed: false,
433            stop_match: Arc::clone(&self.stop_match),
434        }
435    }
436}
437
438impl<'a> ExecutionState<'a> {
439    pub(crate) fn new(
440        db: DbView<'a>,
441        buffers: DenseIdMap<TableId, Box<dyn MutationBuffer>>,
442    ) -> Self {
443        ExecutionState {
444            predicted: Default::default(),
445            db,
446            buffers: MutationBuffers::new(db.notification_list, buffers),
447            changed: false,
448            stop_match: Arc::new(AtomicBool::new(false)),
449        }
450    }
451
452    /// Stage an insertion of the given row into `table`.
453    ///
454    /// If you are using `egglog`, consider using `egglog_bridge::TableAction`.
455    pub fn stage_insert(&mut self, table: TableId, row: &[Value]) {
456        self.buffers
457            .lazy_init(table, || self.db.table_info[table].table.new_buffer());
458        self.buffers.stage_insert(table, row);
459        self.changed = true;
460    }
461
462    /// Stage a removal of the given row from `table` if it is present.
463    ///
464    /// If you are using `egglog`, consider using `egglog_bridge::TableAction`.
465    pub fn stage_remove(&mut self, table: TableId, key: &[Value]) {
466        self.buffers
467            .lazy_init(table, || self.db.table_info[table].table.new_buffer());
468        self.buffers.stage_remove(table, key);
469        self.changed = true;
470    }
471
472    /// Call an external function.
473    pub fn call_external_func(
474        &mut self,
475        func: ExternalFunctionId,
476        args: &[Value],
477    ) -> Option<Value> {
478        self.db.external_funcs[func].invoke(self, args)
479    }
480
481    /// The value the caller of this operation supplied as its
482    /// [`ExternalContext`], if any.
483    pub fn external_context(&self) -> ExternalContext<'a> {
484        self.db.external_context
485    }
486
487    pub fn inc_counter(&self, ctr: CounterId) -> usize {
488        self.db.counters.inc(ctr)
489    }
490
491    pub fn read_counter(&self, ctr: CounterId) -> usize {
492        self.db.counters.read(ctr)
493    }
494
495    /// Iterate over the identifiers of all tables visible to this execution state.
496    pub fn table_ids(&self) -> impl Iterator<Item = TableId> + '_ {
497        self.db.table_info.iter().map(|(id, _)| id)
498    }
499
500    /// Get an immutable reference to the table with id `table`.
501    /// Dangerous: Reading from a table during action execution may break the semi-naive evaluation
502    pub fn get_table(&self, table: TableId) -> &'a WrappedTable {
503        &self.db.table_info[table].table
504    }
505
506    /// Call `f` on each visible row in `table` whose `col` equals `value`,
507    /// with the whole row as a value slice.
508    pub fn for_each_matching_col(
509        &self,
510        table: TableId,
511        col: ColumnId,
512        value: Value,
513        mut f: impl FnMut(&[Value]),
514    ) {
515        let table_info = &self.db.table_info[table];
516        let constraint = Constraint::EqConst { col, val: value };
517
518        // Same order of preference as `Database::process_constraints`, for the
519        // one equality this takes: a sort the table already has, else an index
520        // on the column if it can be cached, else the constraint applied during
521        // the scan.
522        let cacheable = !*table_info
523            .spec
524            .uncacheable_columns
525            .get(col)
526            .unwrap_or(&false);
527        let (subset, slow) = if let Some(subset) = table_info.table.fast_subset(&constraint) {
528            (subset, Vec::new())
529        } else if cacheable {
530            let index = get_column_index_from_tableinfo(table_info, col);
531            let subset = match index.get().unwrap().get_subset(&value) {
532                Some(subset) => with_pool_set(|ps| subset.to_owned(&ps.get_pool())),
533                // No rows hold this key.
534                None => Subset::empty(),
535            };
536            (subset, Vec::new())
537        } else {
538            (table_info.table.all(), vec![constraint])
539        };
540
541        let imp = &table_info.table;
542        let cols: SmallVec<[_; 8]> = (0..imp.spec().arity()).map(ColumnId::from_usize).collect();
543        let mut cur = Offset::new(0);
544        let mut buf = TaggedRowBuffer::new_inline(imp.spec().arity());
545
546        macro_rules! drain_buf {
547            ($buf:expr) => {
548                for (_, row) in $buf.iter() {
549                    f(row);
550                }
551                $buf.clear();
552            };
553        }
554
555        while let Some(next) = imp.scan_project(subset.as_ref(), &cols, cur, 1024, &slow, &mut buf)
556        {
557            drain_buf!(buf);
558            cur = next;
559        }
560        drain_buf!(buf);
561    }
562
563    /// Get the human-readable name for a table, if one exists.
564    pub fn table_name(&self, table: TableId) -> Option<&'a str> {
565        self.db.table_info[table].name()
566    }
567
568    pub fn base_values(&self) -> &'a BaseValues {
569        self.db.bases
570    }
571
572    pub fn container_values(&self) -> &'a ContainerValues {
573        self.db.containers
574    }
575
576    /// Get the _current_ value for a given key in `table`, or otherwise insert
577    /// the unique _next_ value.
578    ///
579    /// Insertions into tables are not performed immediately, but rules and
580    /// merge functions sometimes need to get the result of an insertion. For
581    /// such cases, executions keep a cache of "predicted" values for a given
582    /// mapping that manage the insertions, etc.
583    ///
584    /// If you are using `egglog`, consider using `egglog_bridge::TableAction`.
585    pub fn predict_val(
586        &mut self,
587        table: TableId,
588        key: &[Value],
589        vals: impl ExactSizeIterator<Item = MergeVal>,
590    ) -> Pooled<Vec<Value>> {
591        if let Some(row) = self.db.table_info[table].table.get_row(key) {
592            return row.vals;
593        }
594        Pooled::cloned(
595            self.predicted
596                .get_val(table, key, || {
597                    Self::construct_new_row(
598                        &self.db,
599                        &mut self.buffers,
600                        &mut self.changed,
601                        table,
602                        key,
603                        vals,
604                    )
605                })
606                .deref(),
607        )
608    }
609
610    fn construct_new_row(
611        db: &DbView,
612        buffers: &mut MutationBuffers,
613        changed: &mut bool,
614        table: TableId,
615        key: &[Value],
616        vals: impl ExactSizeIterator<Item = MergeVal>,
617    ) -> Pooled<Vec<Value>> {
618        with_pool_set(|ps| {
619            let mut new = ps.get::<Vec<Value>>();
620            new.reserve(key.len() + vals.len());
621            new.extend_from_slice(key);
622            for val in vals {
623                new.push(match val {
624                    MergeVal::Counter(ctr) => Value::from_usize(db.counters.inc(ctr)),
625                    MergeVal::Constant(c) => c,
626                })
627            }
628            buffers.lazy_init(table, || db.table_info[table].table.new_buffer());
629            buffers.stage_insert(table, &new);
630            *changed = true;
631            new
632        })
633    }
634
635    /// A variant of [`ExecutionState::predict_val`] that avoids materializing the full row, and
636    /// instead only returns the value in the given column.
637    pub fn predict_col(
638        &mut self,
639        table: TableId,
640        key: &[Value],
641        vals: impl ExactSizeIterator<Item = MergeVal>,
642        col: ColumnId,
643    ) -> Value {
644        if let Some(val) = self.db.table_info[table].table.get_row_column(key, col) {
645            return val;
646        }
647        self.predicted.get_val(table, key, || {
648            Self::construct_new_row(
649                &self.db,
650                &mut self.buffers,
651                &mut self.changed,
652                table,
653                key,
654                vals,
655            )
656        })[col.index()]
657    }
658
659    /// Trigger early stopping by setting the stop_match flag.
660    /// This causes rule execution to halt at the next opportunity.
661    ///
662    /// Uses Release ordering to ensure all prior writes are visible to threads that observe this flag.
663    pub fn trigger_early_stop(&self) {
664        self.stop_match.store(true, Ordering::Release);
665    }
666
667    /// Check if early stopping has been requested.
668    ///
669    /// Uses Acquire ordering to ensure we see all writes that happened before the flag was set.
670    pub fn should_stop(&self) -> bool {
671        self.stop_match.load(Ordering::Acquire)
672    }
673}
674
675impl ExecutionState<'_> {
676    /// Returns the number of matches that make it to the end of the instructions
677    pub(crate) fn run_instrs(&mut self, instrs: &[Instr], bindings: &mut Bindings) -> usize {
678        if bindings.var_offsets.next_id().rep() == 0 {
679            // If we have no variables, we want to run the rules once.
680            bindings.matches = 1;
681        }
682
683        // Vectorized execution for larger batch sizes
684        let mut mask = with_pool_set(|ps| Mask::new(0..bindings.matches, ps));
685        for instr in instrs {
686            if mask.is_empty() {
687                return 0;
688            }
689            self.run_instr(&mut mask, instr, bindings);
690        }
691        mask.count_ones()
692    }
693    fn run_instr(&mut self, mask: &mut Mask, inst: &Instr, bindings: &mut Bindings) {
694        fn assert_impl(
695            bindings: &mut Bindings,
696            mask: &mut Mask,
697            l: &QueryEntry,
698            r: &QueryEntry,
699            op: impl Fn(Value, Value) -> bool,
700        ) {
701            match (l, r) {
702                (QueryEntry::Var(v1), QueryEntry::Var(v2)) => {
703                    mask.iter(&bindings[*v1])
704                        .zip(&bindings[*v2])
705                        .retain(|(v1, v2)| op(*v1, *v2));
706                }
707                (QueryEntry::Var(v), QueryEntry::Const(c))
708                | (QueryEntry::Const(c), QueryEntry::Var(v)) => {
709                    mask.iter(&bindings[*v]).retain(|v| op(*v, *c));
710                }
711                (QueryEntry::Const(c1), QueryEntry::Const(c2)) => {
712                    if !op(*c1, *c2) {
713                        mask.clear();
714                    }
715                }
716            }
717        }
718
719        match inst {
720            Instr::LookupOrInsertDefault {
721                table: table_id,
722                args,
723                default,
724                dst_col,
725                dst_var,
726            } => {
727                let pool = with_pool_set(|ps| ps.get_pool::<Vec<Value>>().clone());
728                self.buffers.lazy_init(*table_id, || {
729                    self.db.table_info[*table_id].table.new_buffer()
730                });
731                let table = &self.db.table_info[*table_id].table;
732                // Do two passes over the current vector. First, do a round of lookups. Then, for
733                // any offsets where the lookup failed, insert the default value.
734                let mut mask_copy = mask.clone();
735                table.lookup_row_vectorized(&mut mask_copy, bindings, args, *dst_col, *dst_var);
736                mask_copy.symmetric_difference(mask);
737                if mask_copy.is_empty() {
738                    return;
739                }
740                let mut out = bindings.take(*dst_var).unwrap();
741                for_each_binding_with_mask!(mask_copy, args.as_slice(), bindings, |iter| {
742                    iter.assign_vec(&mut out.vals, |offset, key| {
743                        // First, check if the entry is already in the table:
744                        // if let Some(row) = table.get_row_column(&key, *dst_col) {
745                        //     return row;
746                        // }
747                        // If not, insert the default value.
748                        //
749                        // We avoid doing this more than once by using the
750                        // `predicted` map.
751                        let prediction_key = (
752                            *table_id,
753                            SmallVec::<[Value; 3]>::from_slice(key.as_slice()),
754                        );
755                        let buffers = &mut self.buffers;
756                        // Bind some mutable references because the closure passed
757                        // to or_insert_with is `move`.
758                        let ctrs = &self.db.counters;
759                        let bindings = &bindings;
760                        let pool = pool.clone();
761                        let row =
762                            self.predicted
763                                .data
764                                .entry(prediction_key)
765                                .or_insert_with(move || {
766                                    let mut row = pool.get();
767                                    row.extend_from_slice(key.as_slice());
768                                    // Extend the key with the default values.
769                                    row.reserve(default.len());
770                                    for val in default {
771                                        let val = match val {
772                                            WriteVal::QueryEntry(QueryEntry::Const(c)) => *c,
773                                            WriteVal::QueryEntry(QueryEntry::Var(v)) => {
774                                                bindings[*v][offset]
775                                            }
776                                            WriteVal::IncCounter(ctr) => {
777                                                Value::from_usize(ctrs.inc(*ctr))
778                                            }
779                                            WriteVal::CurrentVal(ix) => row[*ix],
780                                        };
781                                        row.push(val)
782                                    }
783                                    // Insert it into the table.
784                                    buffers.stage_insert(*table_id, &row);
785                                    row
786                                });
787                        row[dst_col.index()]
788                    });
789                });
790                bindings.replace(out);
791            }
792            Instr::LookupWithDefault {
793                table,
794                args,
795                dst_col,
796                dst_var,
797                default,
798            } => {
799                let table = &self.db.table_info[*table].table;
800                table.lookup_with_default_vectorized(
801                    mask, bindings, args, *dst_col, *default, *dst_var,
802                );
803            }
804            Instr::Lookup {
805                table,
806                args,
807                dst_col,
808                dst_var,
809            } => {
810                let table = &self.db.table_info[*table].table;
811                table.lookup_row_vectorized(mask, bindings, args, *dst_col, *dst_var);
812            }
813
814            Instr::LookupWithFallback {
815                table: table_id,
816                table_key,
817                func,
818                func_args,
819                dst_col,
820                dst_var,
821            } => {
822                let table = &self.db.table_info[*table_id].table;
823                let mut lookup_result = mask.clone();
824                table.lookup_row_vectorized(
825                    &mut lookup_result,
826                    bindings,
827                    table_key,
828                    *dst_col,
829                    *dst_var,
830                );
831                let mut to_call_func = lookup_result.clone();
832                to_call_func.symmetric_difference(mask);
833                if to_call_func.is_empty() {
834                    return;
835                }
836
837                // Call the given external function on all entries where the lookup failed.
838                invoke_batch_assign(
839                    self.db.external_funcs[*func].as_ref(),
840                    self,
841                    &mut to_call_func,
842                    bindings,
843                    func_args,
844                    *dst_var,
845                );
846                // The new mask should be the lanes where the lookup succeeded or where `func`
847                // succeeded.
848                lookup_result.union(&to_call_func);
849                *mask = lookup_result;
850            }
851            Instr::Insert { table, vals } => {
852                for_each_binding_with_mask!(mask, vals.as_slice(), bindings, |iter| {
853                    iter.for_each(|vals| {
854                        self.stage_insert(*table, vals.as_slice());
855                    })
856                });
857            }
858            Instr::InsertIfEq { table, l, r, vals } => match (l, r) {
859                (QueryEntry::Var(v1), QueryEntry::Var(v2)) => {
860                    for_each_binding_with_mask!(mask, vals.as_slice(), bindings, |iter| {
861                        iter.zip(&bindings[*v1])
862                            .zip(&bindings[*v2])
863                            .for_each(|((vals, v1), v2)| {
864                                if v1 == v2 {
865                                    self.stage_insert(*table, &vals);
866                                }
867                            })
868                    })
869                }
870                (QueryEntry::Var(v), QueryEntry::Const(c))
871                | (QueryEntry::Const(c), QueryEntry::Var(v)) => {
872                    for_each_binding_with_mask!(mask, vals.as_slice(), bindings, |iter| {
873                        iter.zip(&bindings[*v]).for_each(|(vals, cond)| {
874                            if cond == c {
875                                self.stage_insert(*table, &vals);
876                            }
877                        })
878                    })
879                }
880                (QueryEntry::Const(c1), QueryEntry::Const(c2)) => {
881                    if c1 == c2 {
882                        for_each_binding_with_mask!(mask, vals.as_slice(), bindings, |iter| iter
883                            .for_each(|vals| {
884                                self.stage_insert(*table, &vals);
885                            }))
886                    }
887                }
888            },
889            Instr::Remove { table, args } => {
890                for_each_binding_with_mask!(mask, args.as_slice(), bindings, |iter| {
891                    iter.for_each(|args| {
892                        self.stage_remove(*table, args.as_slice());
893                    })
894                });
895            }
896            Instr::External { func, args, dst } => {
897                invoke_batch(
898                    self.db.external_funcs[*func].as_ref(),
899                    self,
900                    mask,
901                    bindings,
902                    args,
903                    *dst,
904                );
905            }
906            Instr::ExternalWithFallback {
907                f1,
908                args1,
909                f2,
910                args2,
911                dst,
912            } => {
913                let mut f1_result = mask.clone();
914                invoke_batch(
915                    self.db.external_funcs[*f1].as_ref(),
916                    self,
917                    &mut f1_result,
918                    bindings,
919                    args1,
920                    *dst,
921                );
922                let mut to_call_f2 = f1_result.clone();
923                to_call_f2.symmetric_difference(mask);
924                if to_call_f2.is_empty() {
925                    return;
926                }
927                // Call the given external function on all entries where the first call failed.
928                invoke_batch_assign(
929                    self.db.external_funcs[*f2].as_ref(),
930                    self,
931                    &mut to_call_f2,
932                    bindings,
933                    args2,
934                    *dst,
935                );
936                // The new mask should be the lanes where either `f1` or `f2` succeeded.
937                f1_result.union(&to_call_f2);
938                *mask = f1_result;
939            }
940            Instr::AssertAnyNe { ops, divider } => {
941                for_each_binding_with_mask!(mask, ops.as_slice(), bindings, |iter| {
942                    iter.retain(|vals| {
943                        vals[0..*divider]
944                            .iter()
945                            .zip(&vals[*divider..])
946                            .any(|(l, r)| l != r)
947                    })
948                })
949            }
950            Instr::AssertEq(l, r) => assert_impl(bindings, mask, l, r, |l, r| l == r),
951            Instr::AssertNe(l, r) => assert_impl(bindings, mask, l, r, |l, r| l != r),
952            Instr::ReadCounter { counter, dst } => {
953                let mut vals = with_pool_set(|ps| ps.get::<Vec<Value>>());
954                let ctr_val = Value::from_usize(self.read_counter(*counter));
955                vals.resize(bindings.matches, ctr_val);
956                bindings.insert(*dst, &vals);
957            }
958        }
959    }
960}
961
962#[derive(Debug, Clone)]
963pub(crate) enum Instr {
964    /// Look up the value of the given table, inserting a new entry with a
965    /// default value if it is not there.
966    LookupOrInsertDefault {
967        table: TableId,
968        args: Vec<QueryEntry>,
969        default: Vec<WriteVal>,
970        dst_col: ColumnId,
971        dst_var: Variable,
972    },
973
974    /// Look up the value of the given table; if the value is not there, use the
975    /// given default.
976    LookupWithDefault {
977        table: TableId,
978        args: Vec<QueryEntry>,
979        dst_col: ColumnId,
980        dst_var: Variable,
981        default: QueryEntry,
982    },
983
984    /// Look up a value of the given table, halting execution if it is not
985    /// there.
986    Lookup {
987        table: TableId,
988        args: Vec<QueryEntry>,
989        dst_col: ColumnId,
990        dst_var: Variable,
991    },
992
993    /// Look up the given key in the table: if the value is not present in the given table, then
994    /// call the given external function with the given arguments. If the external function returns
995    /// a value, that value is returned in the given `dst_var`. If the lookup fails and the
996    /// external function does not return a value, then execution is halted.
997    LookupWithFallback {
998        table: TableId,
999        table_key: Vec<QueryEntry>,
1000        func: ExternalFunctionId,
1001        func_args: Vec<QueryEntry>,
1002        dst_col: ColumnId,
1003        dst_var: Variable,
1004    },
1005
1006    /// Insert the given return value value with the provided arguments into the
1007    /// table.
1008    Insert {
1009        table: TableId,
1010        vals: Vec<QueryEntry>,
1011    },
1012
1013    /// Insert `vals` into `table` if `l` and `r` are equal.
1014    InsertIfEq {
1015        table: TableId,
1016        l: QueryEntry,
1017        r: QueryEntry,
1018        vals: Vec<QueryEntry>,
1019    },
1020
1021    /// Remove the entry corresponding to `args` in `func`.
1022    Remove {
1023        table: TableId,
1024        args: Vec<QueryEntry>,
1025    },
1026
1027    /// Bind the result of the external function to a variable.
1028    External {
1029        func: ExternalFunctionId,
1030        args: Vec<QueryEntry>,
1031        dst: Variable,
1032    },
1033
1034    /// Bind the result of the external function to a variable. If the first external function
1035    /// fails, then use the second external function. If both fail, execution is haulted, (as in a
1036    /// single failure of `External`).
1037    ExternalWithFallback {
1038        f1: ExternalFunctionId,
1039        args1: Vec<QueryEntry>,
1040        f2: ExternalFunctionId,
1041        args2: Vec<QueryEntry>,
1042        dst: Variable,
1043    },
1044
1045    /// Continue execution iff the two variables are equal.
1046    AssertEq(QueryEntry, QueryEntry),
1047
1048    /// Continue execution iff the two variables are not equal.
1049    AssertNe(QueryEntry, QueryEntry),
1050
1051    /// For the two slices: ops[0..divider] and ops[divider..], continue
1052    /// execution iff there is one pair of values at the same offset that are
1053    /// not equal.
1054    AssertAnyNe {
1055        ops: Vec<QueryEntry>,
1056        divider: usize,
1057    },
1058
1059    /// Read the value of a counter and write it to the given variable.
1060    ReadCounter {
1061        /// The counter to broadcast.
1062        counter: CounterId,
1063        /// The variable to write the value to.
1064        dst: Variable,
1065    },
1066}