egglog_bridge/

lib.rs

//! An implementation of egglog-style queries on top of core-relations.
//!
//! This module translates a well-typed egglog-esque query into the abstractions
//! from the `core-relations` crate. The main higher-level functionality that it
//! implements are seminaive evaluation, default values, and merge functions.
//!
//! This crate is essentially involved in desugaring: it elaborates the encoding
//! of core egglog functionality, but it does not implement algorithms for
//! joins, union-finds, etc.

use std::{
    fmt::Debug,
    hash::Hash,
    iter, mem,
    ops::{Index, IndexMut},
    sync::{Arc, Mutex},
};

use crate::core_relations::{
    BaseValue, BaseValueId, BaseValues, ColumnId, Constraint, ContainerValue, ContainerValues,
    CounterId, Database, DisplacedTable, ExecutionState, ExternalFunction, ExternalFunctionId,
    MergeVal, Offset, PlanStrategy, SortedWritesTable, TableId, TaggedRowBuffer, Value,
    WrappedTable,
};
use crate::numeric_id::{DenseIdMap, DenseIdMapWithReuse, NumericId, define_id};
use egglog_concurrency::ThreadPool;
use egglog_core_relations as core_relations;
use egglog_numeric_id as numeric_id;
use egglog_reports::{IterationReport, ReportLevel, RuleSetReport};
use hashbrown::HashMap;
use indexmap::IndexSet;
use log::info;
use once_cell::sync::Lazy;
use smallvec::SmallVec;
use web_time::{Duration, Instant};

pub mod macros;
pub(crate) mod rule;
#[cfg(test)]
mod tests;

pub use rule::{Function, QueryEntry, RuleBuilder};
use thiserror::Error;

/// A live registry of action handles for use by typed primitives.
///
/// Maps table name to [`TableAction`] (plus the shared [`UnionAction`]
/// and the default-panic external-function id) and is owned by the
/// bridge `EGraph`. The state wrappers (`PureState`/`ReadState`/
/// `WriteState`/`FullState`) live in the `egglog` crate; they read
/// from this registry at invoke time to back name-indexed action
/// methods. Held by the bridge `EGraph` inside an `Arc<RwLock<_>>`.
#[derive(Clone)]
pub struct ActionRegistry {
    table_actions: hashbrown::HashMap<String, TableAction>,
    union_action: UnionAction,
    default_panic_id: ExternalFunctionId,
}

impl ActionRegistry {
    pub(crate) fn new(union_action: UnionAction, default_panic_id: ExternalFunctionId) -> Self {
        Self {
            table_actions: hashbrown::HashMap::new(),
            union_action,
            default_panic_id,
        }
    }

    pub(crate) fn register_table(&mut self, name: String, action: TableAction) {
        self.table_actions.insert(name, action);
    }

    /// Look up the [`TableAction`] for a table by name, or `None` if
    /// no table with that name has been registered.
    pub fn lookup_table(&self, name: &str) -> Option<&TableAction> {
        self.table_actions.get(name)
    }

    /// Snapshot the registered table names and their current row counts.
    pub fn table_sizes(&self, state: &ExecutionState) -> Vec<(&str, usize)> {
        self.table_actions
            .iter()
            .map(|(name, action)| (name.as_str(), action.row_count(state)))
            .collect()
    }

    /// The shared [`UnionAction`] for this EGraph's union-find.
    pub fn union_action(&self) -> &UnionAction {
        &self.union_action
    }

    /// The default panic external function id, used by the egglog
    /// crate's `ActionView::panic`.
    pub fn default_panic_id(&self) -> ExternalFunctionId {
        self.default_panic_id
    }
}

#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub enum ColumnTy {
    Id,
    Base(BaseValueId),
}

define_id!(pub RuleId, u32, "An egglog-style rule");
define_id!(pub FunctionId, u32, "An id representing an egglog function");
define_id!(pub(crate) Timestamp, u32, "An abstract timestamp used to track execution of egglog rules");
impl Timestamp {
    fn to_value(self) -> Value {
        Value::new(self.rep())
    }
}

/// The state associated with an egglog program.
#[derive(Clone)]
pub struct EGraph {
    db: Database,
    uf_table: TableId,
    id_counter: CounterId,
    timestamp_counter: CounterId,
    rules: DenseIdMapWithReuse<RuleId, RuleInfo>,
    funcs: DenseIdMap<FunctionId, FunctionInfo>,
    panic_message: SideChannel<String>,
    /// This is a cache of all the different panic messages that we may use while executing rules
    /// against the EGraph. Oftentimes, these messages are generated dynamically: keeping this map
    /// around allows us to cache external function ids with repeat panic messages and they can
    /// also serve as a debugging tool in the case that the number of panic messages grows without
    /// bound.
    panic_funcs: HashMap<String, ExternalFunctionId>,
    report_level: ReportLevel,
    /// Live registry of name-indexed action handles. Shared (via
    /// `Arc<RwLock<_>>`) with state wrappers and primitive callbacks
    /// in the egglog crate so name-indexed action methods on
    /// [`WriteState`] / [`FullState`] can resolve table actions at
    /// invoke time. Mutated in place from [`add_table`](EGraph::add_table).
    action_registry: Arc<std::sync::RwLock<ActionRegistry>>,
    threads: usize,
    thread_pool: Option<Arc<ThreadPool>>,
}

pub type Result<T> = std::result::Result<T, anyhow::Error>;

fn normalize_thread_count(threads: usize) -> usize {
    #[cfg(target_family = "wasm")]
    {
        if threads > 1 {
            panic!("cannot use more than 1 thread on wasm");
        }
        1
    }
    #[cfg(not(target_family = "wasm"))]
    {
        if threads == 0 {
            std::thread::available_parallelism()
                .map(usize::from)
                .unwrap_or(1)
        } else {
            threads
        }
    }
}

fn install_thread_pool<R>(thread_pool: Option<Arc<ThreadPool>>, f: impl FnOnce() -> R) -> R {
    match thread_pool {
        Some(thread_pool) => thread_pool.install(f),
        None => egglog_concurrency::without_current_pool(f),
    }
}

impl Default for EGraph {
    fn default() -> Self {
        Self::new(1)
    }
}

/// Properties of a function added to an [`EGraph`].
pub struct FunctionConfig {
    /// The function's schema. The last column in the schema is the return type.
    pub schema: Vec<ColumnTy>,
    /// The behavior of the function when lookups are made on keys not currently present.
    pub default: DefaultVal,
    /// How to resolve FD conflicts for the function.
    pub merge: MergeFn,
    /// The function's name
    pub name: String,
    /// Whether or not subsumption is enabled for this function.
    pub can_subsume: bool,
}

impl EGraph {
    /// Create an e-graph configured to use `threads` worker threads.
    ///
    /// Passing `1` keeps execution serial and does not allocate a thread pool.
    /// Passing `0` uses the host's available parallelism.
    pub fn new(threads: usize) -> Self {
        let threads = normalize_thread_count(threads);
        let thread_pool = (threads > 1).then(|| Arc::new(ThreadPool::new(threads)));
        let (mut db, uf_table, id_counter, ts_counter) =
            install_thread_pool(thread_pool.clone(), || {
                let mut db = Database::new();
                let uf_table = db.add_table_named(
                    DisplacedTable::default(),
                    "$uf".into(),
                    iter::empty(),
                    iter::empty(),
                );
                let id_counter = db.add_counter();
                let ts_counter = db.add_counter();
                // Start the timestamp counter at 1.
                db.inc_counter(ts_counter);
                (db, uf_table, id_counter, ts_counter)
            });

        // Register a default panic external function so the typed
        // state wrappers' `panic()` method has an id to call. This
        // also seeds `panic_funcs` so a later `new_panic` with the
        // same message reuses the id.
        let panic_message: SideChannel<String> = Default::default();
        let mut panic_funcs: HashMap<String, ExternalFunctionId> = Default::default();
        let default_panic_msg = "primitive panicked".to_string();
        let default_panic_id = db.add_external_function(Box::new(Panic(
            default_panic_msg.clone(),
            panic_message.clone(),
        )));
        panic_funcs.insert(default_panic_msg, default_panic_id);

        let union_action = UnionAction {
            table: uf_table,
            timestamp: ts_counter,
        };
        let action_registry = Arc::new(std::sync::RwLock::new(ActionRegistry::new(
            union_action,
            default_panic_id,
        )));

        Self {
            db,
            uf_table,
            id_counter,
            timestamp_counter: ts_counter,
            rules: Default::default(),
            funcs: Default::default(),
            panic_message,
            panic_funcs,
            report_level: Default::default(),
            action_registry,
            threads,
            thread_pool,
        }
    }

    /// Return a copy of this e-graph configured with a different thread count.
    pub fn with_num_threads(mut self, threads: usize) -> Self {
        self.set_num_threads(threads);
        self
    }

    /// Set the number of worker threads used by this e-graph.
    ///
    /// Passing `1` disables the pool. Passing `0` uses available parallelism.
    pub fn set_num_threads(&mut self, threads: usize) {
        let threads = normalize_thread_count(threads);
        self.threads = threads;
        self.thread_pool = (threads > 1).then(|| Arc::new(ThreadPool::new(threads)));
    }

    /// Return the configured thread count.
    pub fn num_threads(&self) -> usize {
        self.threads
    }

    fn thread_pool(&self) -> Option<Arc<ThreadPool>> {
        self.thread_pool.clone()
    }

    fn next_ts(&self) -> Timestamp {
        Timestamp::from_usize(self.db.read_counter(self.timestamp_counter))
    }

    fn inc_ts(&mut self) {
        self.db.inc_counter(self.timestamp_counter);
    }

    /// Get a mutable reference to the underlying table of base values for this
    /// `EGraph`.
    pub fn base_values_mut(&mut self) -> &mut BaseValues {
        self.db.base_values_mut()
    }

    /// Get a mutable reference to the underlying table of containers for this
    /// `EGraph`.
    pub fn container_values_mut(&mut self) -> &mut ContainerValues {
        self.db.container_values_mut()
    }

    /// Get a reference to the underlying table of containers for this `EGraph`.
    pub fn container_values(&self) -> &ContainerValues {
        self.db.container_values()
    }

    /// Intern the given container value into the EGraph.
    pub fn get_container_value<C: ContainerValue>(&mut self, val: C) -> Value {
        self.register_container_ty::<C>();
        self.db
            .with_execution_state(|state| state.clone().container_values().register_val(val, state))
    }

    /// Register the given [`ContainerValue`] type with this EGraph.
    ///
    /// The given container will use the EGraph's union-find to manage rebuilding and the merging
    /// of containers with a common id.
    pub fn register_container_ty<C: ContainerValue>(&mut self) {
        let uf_table = self.uf_table;
        let ts_counter = self.timestamp_counter;
        self.db.container_values_mut().register_type::<C>(
            self.id_counter,
            move |state, old, new| {
                if old != new {
                    let next_ts = Value::from_usize(state.read_counter(ts_counter));
                    state.stage_insert(uf_table, &[old, new, next_ts]);
                    std::cmp::min(old, new)
                } else {
                    old
                }
            },
        );
    }

    /// Get a reference to the underlying table of base values for this `EGraph`.
    pub fn base_values(&self) -> &BaseValues {
        self.db.base_values()
    }

    /// Create a [`QueryEntry`] for a base value.
    pub fn base_value_constant<T>(&self, x: T) -> QueryEntry
    where
        T: BaseValue,
    {
        QueryEntry::Const {
            val: self.base_values().get(x),
            ty: ColumnTy::Base(self.base_values().get_ty::<T>()),
        }
    }

    /// Register a low-level external function. The callback receives a
    /// raw `&mut ExecutionState`.
    ///
    /// # Seminaive-safety trust boundary
    ///
    /// Like [`EGraph::with_execution_state`], this is a raw escape —
    /// the registered function has unrestricted access and is not
    /// tracked by the per-context validity system. Prefer building
    /// primitives via the higher-level `egglog::Primitive` /
    /// `egglog::EGraph::add_primitive` API, which enforces #772's
    /// seminaive-safety contract.
    pub fn register_external_func(
        &mut self,
        func: Box<dyn ExternalFunction + 'static>,
    ) -> ExternalFunctionId {
        self.db.add_external_function(func)
    }

    pub fn free_external_func(&mut self, func: ExternalFunctionId) {
        self.db.free_external_function(func)
    }

    /// Generate a fresh id.
    pub fn fresh_id(&mut self) -> Value {
        Value::from_usize(self.db.inc_counter(self.id_counter))
    }

    /// Look up the canonical value for `val` in the union-find.
    ///
    /// If the value has never been inserted into the union-find, `val` is returned.
    fn get_canon_in_uf(&self, val: Value) -> Value {
        let table = self.db.get_table(self.uf_table);
        let row = table.get_row(&[val]);
        row.map(|row| row.vals[1]).unwrap_or(val)
    }

    /// Get the canonical representation for `val` based on type.
    ///
    /// For [`ColumnTy::Id`], it looks up the union find; otherwise,
    /// it returns the value itself.
    pub fn get_canon_repr(&self, val: Value, ty: ColumnTy) -> Value {
        match ty {
            ColumnTy::Id => self.get_canon_in_uf(val),
            ColumnTy::Base(_) => val,
        }
    }

    /// Load the given values into the database.
    ///
    /// # Panics
    /// This method panics if the values do not match the arity of the function.
    ///
    /// NB: this is not an efficient interface for bulk loading. We should add
    /// one that allows us to pass through a series of RowBuffers before
    /// incrementing the timestamp.
    pub fn add_values(&mut self, values: impl IntoIterator<Item = (FunctionId, Vec<Value>)>) {
        let mut extended_row = Vec::<Value>::new();
        let mut bufs = DenseIdMap::default();
        for (func, row) in values.into_iter() {
            let table_info = &self.funcs[func];
            let schema_math = SchemaMath {
                subsume: table_info.can_subsume,
                func_cols: table_info.schema.len(),
            };
            let table_id = table_info.table;
            extended_row.extend_from_slice(&row);
            schema_math.write_table_row(
                &mut extended_row,
                RowVals {
                    timestamp: self.next_ts().to_value(),
                    subsume: schema_math.subsume.then_some(NOT_SUBSUMED),
                    ret_val: None, // already filled in.
                },
            );
            let buf = bufs.get_or_insert(table_id, || self.db.new_buffer(table_id));
            buf.stage_insert(&extended_row);
            extended_row.clear();
        }
        // Flush the buffers.
        mem::drop(bufs);
        self.flush_updates();
    }

    /// A term-oriented means of adding data to the database: hand back a "term
    /// id" for the given function and keys for the function.
    ///
    /// # Panics
    /// This method panics if the values do not match the arity of the function.
    pub fn add_term(&mut self, func: FunctionId, inputs: &[Value]) -> Value {
        let info = &self.funcs[func];
        let schema_math = SchemaMath {
            subsume: info.can_subsume,
            func_cols: info.schema.len(),
        };
        let mut extended_row = Vec::new();
        extended_row.extend_from_slice(inputs);
        let res = self.fresh_id();
        schema_math.write_table_row(
            &mut extended_row,
            RowVals {
                timestamp: self.next_ts().to_value(),
                ret_val: Some(res),
                subsume: schema_math.subsume.then_some(NOT_SUBSUMED),
            },
        );
        extended_row[schema_math.ret_val_col()] = res;
        let table_id = self.funcs[func].table;
        self.db.new_buffer(table_id).stage_insert(&extended_row);
        self.flush_updates();
        self.get_canon_in_uf(res)
    }

    /// Lookup the id associated with a function `func` and the given arguments
    /// (`key`).
    pub fn lookup_id(&self, func: FunctionId, key: &[Value]) -> Option<Value> {
        let info = &self.funcs[func];
        let schema_math = SchemaMath {
            subsume: info.can_subsume,
            func_cols: info.schema.len(),
        };
        let table_id = info.table;
        let table = self.db.get_table(table_id);
        let row = table.get_row(key)?;
        Some(row.vals[schema_math.ret_val_col()])
    }

    pub fn approx_table_size(&self, table: FunctionId) -> usize {
        self.db.estimate_size(self.funcs[table].table, None)
    }

    pub fn table_size(&self, table: FunctionId) -> usize {
        self.db.get_table(self.funcs[table].table).len()
    }

    /// Remove every row from the given function's backing table.
    ///
    /// This is the bulk counterpart to staging a `remove` for every key in the
    /// table: the underlying `Database::clear_table` drops the row buffer in
    /// O(1)-in-row-count time and bumps the table's major generation, which
    /// lazily invalidates any cached subsets or indexes a later reader might
    /// consult. Any rows staged for this table by an in-flight
    /// `MutationBuffer` are dropped along with the table contents.
    ///
    /// Callers that have staged inserts/removes for *other* tables that they
    /// want flushed first should call [`EGraph::flush_updates`] before
    /// clearing.
    pub fn clear_table(&mut self, func: FunctionId) {
        let table_id = self.funcs[func].table;
        self.db.clear_table(table_id);
    }

    /// Read the contents of the given function.
    ///
    /// The callback `f` is called with each row and its subsumption status.
    pub fn for_each(&self, table: FunctionId, mut f: impl FnMut(ScanEntry<'_>)) {
        self.for_each_while(table, |row| {
            f(row);
            true
        });
    }

    /// Iterate over the rows of a function table, calling `f` on each row. If `f` returns `false`
    /// the function returns early and stops reading rows from the table.
    pub fn for_each_while(&self, table: FunctionId, mut f: impl FnMut(ScanEntry<'_>) -> bool) {
        let info = &self.funcs[table];
        let table = self.funcs[table].table;
        let schema_math = SchemaMath {
            subsume: info.can_subsume,
            func_cols: info.schema.len(),
        };
        let imp = self.db.get_table(table);
        let all = imp.all();
        let mut cur = Offset::new(0);
        let mut buf = TaggedRowBuffer::new(imp.spec().arity());
        // This somewhat awkward iteration strategy is forced on us by the `scan_bounded` API. We
        // should look into ways to avoid this cludge where the loop body effectively must be
        // repeated at the end. The obvious and idiomatic ways to do this all require
        // `dyn`-compatibility on `Table` or dynamic dispatch per row.
        macro_rules! drain_buf {
            ($buf:expr) => {
                for (_, row) in $buf.non_stale() {
                    let subsumed =
                        schema_math.subsume && row[schema_math.subsume_col()] == SUBSUMED;
                    if !f(ScanEntry {
                        vals: &row[0..schema_math.func_cols],
                        subsumed,
                    }) {
                        return;
                    }
                }
                $buf.clear();
            };
        }
        while let Some(next) = imp.scan_bounded(all.as_ref(), cur, 32, &mut buf) {
            drain_buf!(buf);
            cur = next;
        }
        drain_buf!(buf);
    }

    /// A basic method for dumping the state of the database to `log::info!`.
    ///
    /// For large tables, this is unlikely to give particularly useful output.
    pub fn dump_debug_info(&self) {
        info!("=== View Tables ===");
        for (id, info) in self.funcs.iter() {
            let table = self.db.get_table(info.table);
            self.scan_table(table, |row| {
                info!(
                    "View Table {name} / {id:?} / {table:?}: {row:?}",
                    name = info.name,
                    table = info.table
                )
            });
        }
    }

    /// A helper for scanning the entries in a table.
    fn scan_table(&self, table: &WrappedTable, mut f: impl FnMut(&[Value])) {
        const BATCH_SIZE: usize = 128;
        let all = table.all();
        let mut cur = Offset::new(0);
        let mut out = TaggedRowBuffer::new(table.spec().arity());
        while let Some(next) = table.scan_bounded(all.as_ref(), cur, BATCH_SIZE, &mut out) {
            out.non_stale().for_each(|(_, row)| f(row));
            out.clear();
            cur = next;
        }
        out.non_stale().for_each(|(_, row)| f(row));
    }

    /// Register a function in this EGraph.
    pub fn add_table(&mut self, config: FunctionConfig) -> FunctionId {
        let FunctionConfig {
            schema,
            default,
            merge,
            name,
            can_subsume,
        } = config;
        assert!(
            !schema.is_empty(),
            "must have at least one column in schema"
        );
        let to_rebuild: Vec<ColumnId> = schema
            .iter()
            .enumerate()
            .filter(|(_, ty)| matches!(ty, ColumnTy::Id))
            .map(|(i, _)| ColumnId::from_usize(i))
            .collect();
        let schema_math = SchemaMath {
            subsume: can_subsume,
            func_cols: schema.len(),
        };
        let n_args = schema_math.num_keys();
        let n_cols = schema_math.table_columns();
        let next_func_id = self.funcs.next_id();
        let mut read_deps = IndexSet::<TableId>::new();
        let mut write_deps = IndexSet::<TableId>::new();
        merge.fill_deps(self, &mut read_deps, &mut write_deps);
        let merge_fn = merge.to_callback(schema_math, &name, self);
        let table = install_thread_pool(self.thread_pool(), || {
            SortedWritesTable::new(
                n_args,
                n_cols,
                Some(ColumnId::from_usize(schema.len())),
                to_rebuild,
                merge_fn,
            )
        });
        let name: Arc<str> = name.into();
        let table_id = self.db.add_table_named(
            table,
            name.clone(),
            read_deps.iter().copied(),
            write_deps.iter().copied(),
        );

        let res = self.funcs.push(FunctionInfo {
            table: table_id,
            schema: schema.clone(),
            incremental_rebuild_rules: Default::default(),
            nonincremental_rebuild_rule: RuleId::new(!0),
            default_val: default,
            can_subsume,
            name,
        });
        debug_assert_eq!(res, next_func_id);
        let incremental_rebuild_rules = self.incremental_rebuild_rules(res, &schema);
        let nonincremental_rebuild_rule = self.nonincremental_rebuild(res, &schema);
        let info = &mut self.funcs[res];
        info.incremental_rebuild_rules = incremental_rebuild_rules;
        info.nonincremental_rebuild_rule = nonincremental_rebuild_rule;
        let action = TableAction::new(self, res);
        let table_name = self.funcs[res].name.to_string();
        self.action_registry
            .write()
            .unwrap()
            .register_table(table_name, action);
        res
    }

    /// A handle to the live [`ActionRegistry`] for this EGraph.
    /// The handle is shared (`Arc<RwLock<_>>`); cloning the outer
    /// `Arc` does not duplicate the underlying registry. Used by the
    /// egglog crate's primitive machinery to thread the registry into
    /// state wrappers at invoke time.
    pub fn action_registry(&self) -> &Arc<std::sync::RwLock<ActionRegistry>> {
        &self.action_registry
    }

    /// Run the given rules, returning whether the database changed.
    ///
    /// If the given rules are malformed, this method can return an error.
    pub fn run_rules(&mut self, rules: &[RuleId]) -> Result<IterationReport> {
        let thread_pool = self.thread_pool();
        install_thread_pool(thread_pool, || self.run_rules_inner(rules))
    }

    fn run_rules_inner(&mut self, rules: &[RuleId]) -> Result<IterationReport> {
        let ts = self.next_ts();

        let uf_size_before = self.db.get_table(self.uf_table).len();
        let rule_set_report =
            run_rules_impl(&mut self.db, &mut self.rules, rules, ts, self.report_level)?;
        if let Some(message) = self.panic_message.lock().unwrap().take() {
            return Err(PanicError(message).into());
        }

        let mut iteration_report = IterationReport {
            rule_set_report,
            rebuild_time: Duration::ZERO,
        };
        let uf_size_after = self.db.get_table(self.uf_table).len();
        if uf_size_before == uf_size_after {
            // No new unions: skip the full rebuild but still advance the
            // timestamp so that seminaive evaluation sees a fresh epoch.
            // Rebuilding is only necessary when new unions have been made because ids may need to be updated.
            // Adding terms doesn't necessarily touch the union-find, only doing a union between existing ids does.
            self.inc_ts();
            return Ok(iteration_report);
        }

        let rebuild_timer = Instant::now();
        self.rebuild()?;
        iteration_report.rebuild_time = rebuild_timer.elapsed();

        if let Some(message) = self.panic_message.lock().unwrap().take() {
            return Err(PanicError(message).into());
        }

        Ok(iteration_report)
    }

    fn rebuild(&mut self) -> Result<()> {
        let do_parallel = egglog_concurrency::current_num_threads() > 1;
        if self.db.get_table(self.uf_table).rebuilder(&[]).is_some() {
            // The UF implementation supports "native"  rebuilding.
            let mut tables = Vec::with_capacity(self.funcs.next_id().index());
            for (_, func) in self.funcs.iter() {
                tables.push(func.table);
            }
            loop {
                // Order matters here: we need to rebuild containers first and then rebuild the
                // tables. Why?
                //
                // Say we have a sort that can map to and from a vector containing only itself:
                // (sort X)
                // (function to-vec (X) (Vec X) :no-merge)
                // (constructor from-vec (Vec X) X)
                // (constructor Num (i64) X)
                // (constructor Add (X X) X)
                //
                // Along with rules:
                // (rule ((= x (Num i))) ((set (to-vec x) (vec-of x))))
                // (rule ((= x (Add i j))) ((set (to-vec x) (vec-of x))))
                // (rule ((= x (from-vec v))) ((set (to-vec x) v))
                // (rewrite (Add (Num i) (Num j)) (Num (+ i j)))
                //
                // These rules, while redundant, should be safe. However, if we rebuild tables
                // before containers some schedules can cause us to violate the `:no-merge`
                // directive, which asserts that all values written for a key are equal.
                //
                // Suppose we start off with x1=(Num 1), x2=(Num 3), and x3=(Add (Num 1) (Num 2)) as
                // expressions, with `to-vec` and `from-vec` entries for all three expressions.
                // We'll call (to-vec xi) vi for all i.
                //
                // Now suppose we run the `rewrite` above: now, x3 = x2. But v3 will only equal v2
                // _after_ we rebuild the `Vec` container. That means that if we rebuild `to-vec`
                // we will collapse the the rows for x3 and x2, but then fail to merge v3 and v2
                // because they are not (yet) equal.
                //
                // Rebuilding containers first will find that v3 and v2 are equal, and the rest of
                // the rules can proceed.
                let container_rebuild = self.db.rebuild_containers(self.uf_table);
                let next_ts = self.next_ts().to_value();
                let table_rebuild = self.db.apply_rebuild(self.uf_table, &tables, next_ts);
                // Container rebuild can make a parent row newly matchable without
                // changing the row's stored id. Re-timestamp those parents so
                // seminaive sees the newly enabled match on the next pass.
                let dirty_ids: Vec<Value> = container_rebuild.dirty_ids().iter().copied().collect();
                let refreshed_rows = self
                    .db
                    .refresh_rows_for_values(&tables, &dirty_ids, next_ts);
                self.inc_ts();
                if !table_rebuild && !refreshed_rows && !container_rebuild.changed() {
                    break;
                }
            }
            return Ok(());
        }
        if do_parallel {
            return self.rebuild_parallel();
        }
        let start = Instant::now();

        // The database changed. Rebuild. New entries should land after the given rules.
        let mut changed = true;
        while changed {
            changed = false;
            // We need to iterate rebuilding to a fixed point. Future scans
            // should look only at the latest updates.
            self.inc_ts();
            let ts = self.next_ts();
            for (_, info) in self.funcs.iter_mut() {
                let last_rebuilt_at = self.rules[info.nonincremental_rebuild_rule].last_run_at;
                let table_size = self.db.estimate_size(info.table, None);
                let uf_size = self.db.estimate_size(
                    self.uf_table,
                    Some(Constraint::GeConst {
                        col: ColumnId::new(2),
                        val: last_rebuilt_at.to_value(),
                    }),
                );
                if incremental_rebuild(uf_size, table_size, false) {
                    marker_incremental_rebuild(|| -> Result<()> {
                        // Run each of the incremental rules serially.
                        //
                        // This is to avoid recanonicalizing the same row multiple
                        // times.
                        for rule in &info.incremental_rebuild_rules {
                            changed |= run_rules_impl(
                                &mut self.db,
                                &mut self.rules,
                                &[*rule],
                                ts,
                                ReportLevel::TimeOnly,
                            )?
                            .changed;
                        }
                        // Reset the rule we did not run. These two should be equivalent.
                        self.rules[info.nonincremental_rebuild_rule].last_run_at = ts;
                        Ok(())
                    })?;
                } else {
                    marker_nonincremental_rebuild(|| -> Result<()> {
                        changed |= run_rules_impl(
                            &mut self.db,
                            &mut self.rules,
                            &[info.nonincremental_rebuild_rule],
                            ts,
                            ReportLevel::TimeOnly,
                        )?
                        .changed;
                        for rule in &info.incremental_rebuild_rules {
                            self.rules[*rule].last_run_at = ts;
                        }
                        Ok(())
                    })?;
                }
            }
        }
        log::info!("rebuild took {:?}", start.elapsed());
        Ok(())
    }

    /// A variant of `rebuild` that attempts to combine rebuild rules into
    /// larger rulesets to increase parallelism. This kind of preprocessing can
    /// slow processing down in a single-threaded setting, so it is only used
    /// when the number of active threads is greater than 1.
    fn rebuild_parallel(&mut self) -> Result<()> {
        let start = Instant::now();
        #[derive(Default)]
        struct RebuildState {
            nonincremental: Vec<FunctionId>,
            incremental: DenseIdMap<usize, SmallVec<[FunctionId; 2]>>,
        }

        impl RebuildState {
            fn clear(&mut self) {
                self.nonincremental.clear();
                self.incremental.iter_mut().for_each(|(_, v)| v.clear());
            }
        }

        let mut changed = true;
        let mut state = RebuildState::default();
        let mut scratch = Vec::new();
        while changed {
            changed = false;
            state.clear();
            self.inc_ts();
            // First, figure out which functions will be rebuilt nonincrementally,
            // vs. incrementally. Group them together.
            for (func, info) in self.funcs.iter_mut() {
                let last_rebuilt_at = self.rules[info.nonincremental_rebuild_rule].last_run_at;
                let table_size = self.db.estimate_size(info.table, None);
                let uf_size = self.db.estimate_size(
                    self.uf_table,
                    Some(Constraint::GeConst {
                        col: ColumnId::new(2),
                        val: last_rebuilt_at.to_value(),
                    }),
                );
                if incremental_rebuild(uf_size, table_size, true) {
                    for (i, _) in info.incremental_rebuild_rules.iter().enumerate() {
                        state.incremental.get_or_default(i).push(func);
                    }
                } else {
                    state.nonincremental.push(func);
                }
            }
            let ts = self.next_ts();
            for func in state.nonincremental.iter().copied() {
                scratch.push(self.funcs[func].nonincremental_rebuild_rule);
                for rule in &self.funcs[func].incremental_rebuild_rules {
                    self.rules[*rule].last_run_at = ts;
                }
            }
            changed |= run_rules_impl(
                &mut self.db,
                &mut self.rules,
                &scratch,
                ts,
                ReportLevel::TimeOnly,
            )?
            .changed;
            scratch.clear();
            let ts = self.next_ts();
            for (i, funcs) in state.incremental.iter() {
                for func in funcs.iter().copied() {
                    let info = &mut self.funcs[func];
                    scratch.push(info.incremental_rebuild_rules[i]);
                    self.rules[info.nonincremental_rebuild_rule].last_run_at = ts;
                }
                changed |= run_rules_impl(
                    &mut self.db,
                    &mut self.rules,
                    &scratch,
                    ts,
                    ReportLevel::TimeOnly,
                )?
                .changed;
                scratch.clear();
            }
        }
        log::info!("rebuild took {:?}", start.elapsed());
        Ok(())
    }

    fn incremental_rebuild_rules(&mut self, table: FunctionId, schema: &[ColumnTy]) -> Vec<RuleId> {
        schema
            .iter()
            .enumerate()
            .filter_map(|(i, ty)| match ty {
                ColumnTy::Id => {
                    Some(self.incremental_rebuild_rule(table, schema, ColumnId::from_usize(i)))
                }
                ColumnTy::Base(_) => None,
            })
            .collect()
    }

    fn incremental_rebuild_rule(
        &mut self,
        table: FunctionId,
        schema: &[ColumnTy],
        col: ColumnId,
    ) -> RuleId {
        let subsume = self.funcs[table].can_subsume;
        let table_id = self.funcs[table].table;
        let uf_table = self.uf_table;
        // Two atoms, one binding a whole tuple, one binding a displaced column
        let mut rb = self.new_rule(&format!("incremental rebuild {table:?}, {col:?}"), true);
        rb.set_plan_strategy(PlanStrategy::MinCover);
        let mut vars = Vec::<QueryEntry>::with_capacity(schema.len());
        for ty in schema {
            vars.push(rb.new_var(*ty).into());
        }
        let canon_val: QueryEntry = rb.new_var(ColumnTy::Id).into();
        let subsume_var = subsume.then(|| rb.new_var(ColumnTy::Id));
        rb.add_atom_with_timestamp_and_func(
            table_id,
            Some(table),
            subsume_var.clone().map(QueryEntry::from),
            &vars,
        );
        rb.add_atom_with_timestamp_and_func(
            uf_table,
            None,
            None,
            &[vars[col.index()].clone(), canon_val.clone()],
        );
        rb.set_focus(1); // Set the uf atom as the sole focus.

        // Now canonicalize the entire row.
        let mut canon = Vec::<QueryEntry>::with_capacity(schema.len());
        for (i, (var, ty)) in vars.iter().zip(schema.iter()).enumerate() {
            canon.push(if i == col.index() {
                canon_val.clone()
            } else if let ColumnTy::Id = ty {
                rb.lookup_uf(var.clone()).unwrap().into()
            } else {
                var.clone()
            })
        }

        // Remove the old row and insert the new one.
        rb.rebuild_row(table, &vars, &canon, subsume_var);
        rb.build()
    }

    fn nonincremental_rebuild(&mut self, table: FunctionId, schema: &[ColumnTy]) -> RuleId {
        let can_subsume = self.funcs[table].can_subsume;
        let table_id = self.funcs[table].table;
        let mut rb = self.new_rule(&format!("nonincremental rebuild {table:?}"), false);
        rb.set_plan_strategy(PlanStrategy::MinCover);
        let mut vars = Vec::<QueryEntry>::with_capacity(schema.len());
        for ty in schema {
            vars.push(rb.new_var(*ty).into());
        }
        let subsume_var = can_subsume.then(|| rb.new_var(ColumnTy::Id));
        rb.add_atom_with_timestamp_and_func(
            table_id,
            Some(table),
            subsume_var.clone().map(QueryEntry::from),
            &vars,
        );
        let mut lhs = SmallVec::<[QueryEntry; 4]>::new();
        let mut rhs = SmallVec::<[QueryEntry; 4]>::new();
        let mut canon = Vec::<QueryEntry>::with_capacity(schema.len());
        for (var, ty) in vars.iter().zip(schema.iter()) {
            canon.push(if let ColumnTy::Id = ty {
                lhs.push(var.clone());
                let canon_var = QueryEntry::from(rb.lookup_uf(var.clone()).unwrap());
                rhs.push(canon_var.clone());
                canon_var
            } else {
                var.clone()
            })
        }
        rb.check_for_update(&lhs, &rhs).unwrap();
        rb.rebuild_row(table, &vars, &canon, subsume_var);
        rb.build()
    }

    /// Gives the user a handle to the underlying ExecutionState. Useful for staging updates
    /// to the database.
    ///
    /// The staged updates are not immediately reflected in the EGraph, so you may want to
    /// manually flush the updates using [`EGraph::flush_updates`].
    ///
    /// # Seminaive-safety trust boundary
    ///
    /// This method hands out a raw `&mut ExecutionState`, which bypasses
    /// the egglog crate's `Read` / `Write` capability wrappers
    /// (`PureState`, `WriteState`, `ReadState`, `FullState`) that
    /// enforce #772's seminaive-safety model. Treat it as top-level
    /// / global-action context: appropriate for one-shot database
    /// manipulation from outside any rule, not for use inside
    /// primitive implementations.
    pub fn with_execution_state<R>(&self, f: impl FnOnce(&mut ExecutionState<'_>) -> R) -> R {
        let thread_pool = self.thread_pool();
        install_thread_pool(thread_pool, || self.db.with_execution_state(f))
    }

    /// Like [`EGraph::with_execution_state`], but also reports whether `f`
    /// staged any mutation. A read-only closure leaves the flag `false`, so
    /// callers can skip a [`EGraph::flush_updates`] that would otherwise be a
    /// no-op merge plus a spurious timestamp bump.
    pub fn with_execution_state_tracked<R>(
        &self,
        f: impl FnOnce(&mut ExecutionState<'_>) -> R,
    ) -> (R, bool) {
        self.db.with_execution_state_tracked(f)
    }

    /// Flush the pending update buffers to the EGraph.
    /// Returns `true` if the database is updated.
    pub fn flush_updates(&mut self) -> bool {
        let thread_pool = self.thread_pool();
        install_thread_pool(thread_pool, || self.flush_updates_inner())
    }

    fn flush_updates_inner(&mut self) -> bool {
        let uf_size_before = self.db.get_table(self.uf_table).len();
        let updated = self.db.merge_all();
        self.inc_ts();
        let uf_size_after = self.db.get_table(self.uf_table).len();
        if uf_size_before != uf_size_after {
            // Rebuilding is only necessary when new unions have been made because ids may need to be updated.
            // Adding terms doesn't necessarily touch the union-find, only doing a union between existing ids does.
            self.rebuild().unwrap();
        }
        updated
    }

    pub fn set_report_level(&mut self, level: ReportLevel) {
        self.report_level = level;
    }
}

#[derive(Clone)]
struct RuleInfo {
    last_run_at: Timestamp,
    query: rule::Query,
    cached_plan: Option<CachedPlanInfo>,
    desc: Arc<str>,
}

#[derive(Clone)]
struct CachedPlanInfo {
    plan: Arc<core_relations::CachedPlan>,
    /// A mapping from index into a [`rule::Query`]'s atoms to the atoms in the underlying cached
    /// plan.
    atom_mapping: Vec<core_relations::AtomId>,
}

#[derive(Clone)]
struct FunctionInfo {
    table: TableId,
    schema: Vec<ColumnTy>,
    incremental_rebuild_rules: Vec<RuleId>,
    nonincremental_rebuild_rule: RuleId,
    default_val: DefaultVal,
    can_subsume: bool,
    name: Arc<str>,
}

impl FunctionInfo {
    fn ret_ty(&self) -> ColumnTy {
        self.schema.last().copied().unwrap()
    }
}

/// How defaults are computed for the given function.
#[derive(Copy, Clone)]
pub enum DefaultVal {
    /// Generate a fresh UF id.
    FreshId,
    /// Cause an egglog-level panic if a lookup fails.
    Fail,
    /// Insert a constant of some kind.
    Const(Value),
}

/// How to resolve FD conflicts for a table.
pub enum MergeFn {
    /// Panic if the old and new values don't match.
    AssertEq,
    /// Use congruence to resolve FD conflicts.
    UnionId,
    /// The output of a merge is determined by applying the given ExternalFunction to the result
    /// of the argument merge functions.
    Primitive(ExternalFunctionId, Vec<MergeFn>),
    /// The output of a merge is determined by looking up the value for the given function and the
    /// given arguments in the egraph.
    Function(FunctionId, Vec<MergeFn>),
    /// Always return the old value for the given function.
    Old,
    /// Always return the new value for the given function.
    New,
    /// Always overwrite the new value for the given function with a constant. This is more useful
    /// as a "base case" in a more complicated merge function (e.g. one that clamps a value between
    /// 1 and 100) than it is as a standalone merge function.
    Const(Value),
}

impl MergeFn {
    fn fill_deps(
        &self,
        egraph: &EGraph,
        read_deps: &mut IndexSet<TableId>,
        write_deps: &mut IndexSet<TableId>,
    ) {
        use MergeFn::*;
        match self {
            Primitive(_, args) => {
                args.iter()
                    .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps));
                write_deps.insert(egraph.uf_table);
            }
            Function(func, args) => {
                read_deps.insert(egraph.funcs[*func].table);
                write_deps.insert(egraph.funcs[*func].table);
                args.iter()
                    .for_each(|arg| arg.fill_deps(egraph, read_deps, write_deps));
            }
            UnionId => {
                write_deps.insert(egraph.uf_table);
            }
            AssertEq | Old | New | Const(..) => {}
        }
    }

    fn to_callback(
        &self,
        schema_math: SchemaMath,
        function_name: &str,
        egraph: &mut EGraph,
    ) -> Box<core_relations::MergeFn> {
        let resolved = self.resolve(function_name, egraph);

        Box::new(move |state, cur, new, out| {
            let timestamp = new[schema_math.ts_col()];

            let mut changed = false;

            let ret_val = {
                let cur = cur[schema_math.ret_val_col()];
                let new = new[schema_math.ret_val_col()];
                let out = resolved.run(state, cur, new, timestamp);
                changed |= cur != out;
                out
            };

            let subsume = schema_math.subsume.then(|| {
                let cur = cur[schema_math.subsume_col()];
                let new = new[schema_math.subsume_col()];
                let out = combine_subsumed(cur, new);
                changed |= cur != out;
                out
            });
            if changed {
                out.extend_from_slice(new);
                schema_math.write_table_row(
                    out,
                    RowVals {
                        timestamp,
                        subsume,
                        ret_val: Some(ret_val),
                    },
                );
            }

            changed
        })
    }

    fn resolve(&self, function_name: &str, egraph: &mut EGraph) -> ResolvedMergeFn {
        match self {
            MergeFn::Const(v) => ResolvedMergeFn::Const(*v),
            MergeFn::Old => ResolvedMergeFn::Old,
            MergeFn::New => ResolvedMergeFn::New,
            MergeFn::AssertEq => ResolvedMergeFn::AssertEq {
                panic: egraph.new_panic(format!(
                    "Illegal merge attempted for function {function_name}"
                )),
            },
            MergeFn::UnionId => ResolvedMergeFn::UnionId {
                uf_table: egraph.uf_table,
            },
            // NB: The primitive and function-based merge functions heap allocate a single callback
            // for each layer of nesting. This introduces a bit of overhead, particularly for cases
            // that look like `(f old new)` or `(f new old)`. We could special-case common cases in
            // this function if that overhead shows up.
            MergeFn::Primitive(prim, args) => ResolvedMergeFn::Primitive {
                prim: *prim,
                args: args
                    .iter()
                    .map(|arg| arg.resolve(function_name, egraph))
                    .collect::<Vec<_>>(),
                panic: egraph.new_panic(format!(
                    "Merge function for {function_name} primitive call failed"
                )),
            },
            MergeFn::Function(func, args) => {
                let func_info = &egraph.funcs[*func];
                assert_eq!(
                    func_info.schema.len(),
                    args.len() + 1,
                    "Merge function for {function_name} must match function arity for {}",
                    func_info.name
                );
                ResolvedMergeFn::Function {
                    func: TableAction::new(egraph, *func),
                    panic: egraph.new_panic(format!(
                        "Lookup on {} failed in the merge function for {function_name}",
                        func_info.name
                    )),
                    args: args
                        .iter()
                        .map(|arg| arg.resolve(function_name, egraph))
                        .collect::<Vec<_>>(),
                }
            }
        }
    }
}

/// This enum is taking the place of a
/// `Box<dyn Fn(&mut ExecutionState, Value, Value, Value) -> Value + Send + Sync>`
/// to avoid extra boxes. It stores the data needed to run a `MergeFn` without
/// holding onto any references, so it can be `move`d inside the `core_relations::MergeFn`.
enum ResolvedMergeFn {
    Const(Value),
    Old,
    New,
    AssertEq {
        panic: ExternalFunctionId,
    },
    UnionId {
        uf_table: TableId,
    },
    Primitive {
        prim: ExternalFunctionId,
        args: Vec<ResolvedMergeFn>,
        panic: ExternalFunctionId,
    },
    Function {
        func: TableAction,
        args: Vec<ResolvedMergeFn>,
        panic: ExternalFunctionId,
    },
}

impl ResolvedMergeFn {
    fn run(&self, state: &mut ExecutionState, cur: Value, new: Value, ts: Value) -> Value {
        match self {
            ResolvedMergeFn::Const(v) => *v,
            ResolvedMergeFn::Old => cur,
            ResolvedMergeFn::New => new,
            ResolvedMergeFn::AssertEq { panic } => {
                if cur != new {
                    let res = state.call_external_func(*panic, &[]);
                    assert_eq!(res, None);
                }
                cur
            }
            ResolvedMergeFn::UnionId { uf_table } => {
                if cur != new {
                    state.stage_insert(*uf_table, &[cur, new, ts]);
                    // We pick the minimum when unioning. This matches the original egglog
                    // behavior. THIS MUST MATCH THE UNION-FIND IMPLEMENTATION!
                    std::cmp::min(cur, new)
                } else {
                    cur
                }
            }
            // NB: The primitive and function-based merge functions heap allocate a single callback
            // for each layer of nesting. This introduces a bit of overhead, particularly for cases
            // that look like `(f old new)` or `(f new old)`. We could special-case common cases in
            // this function if that overhead shows up.
            ResolvedMergeFn::Primitive { prim, args, panic } => {
                let args = args
                    .iter()
                    .map(|arg| arg.run(state, cur, new, ts))
                    .collect::<Vec<_>>();

                match state.call_external_func(*prim, &args) {
                    Some(result) => result,
                    None => {
                        let res = state.call_external_func(*panic, &[]);
                        assert_eq!(res, None);
                        cur
                    }
                }
            }
            ResolvedMergeFn::Function { func, args, panic } => {
                // see github.com/egraphs-good/egglog/pull/287
                if cur == new {
                    return cur;
                }

                let args = args
                    .iter()
                    .map(|arg| arg.run(state, cur, new, ts))
                    .collect::<Vec<_>>();

                // Merge functions dispatch to another function that may be
                // a constructor (mint fresh id on miss) or a custom function
                // (return `None` → panic). `lookup_or_insert` preserves
                // both behaviors; the pure-read `lookup` would skip
                // constructor minting.
                func.lookup_or_insert(state, &args).unwrap_or_else(|| {
                    let res = state.call_external_func(*panic, &[]);
                    assert_eq!(res, None);
                    cur
                })
            }
        }
    }
}

/// Coarse classification of a table — `Constructor` mints a fresh
/// eclass id when a row is missed; `Function` does not. Mirrors the
/// `FunctionSubtype` split on the egglog side without dragging that
/// type into the bridge crate.
#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub enum TableKind {
    Function,
    Constructor,
}

/// This is an intern-able struct that holds all the data needed
/// to do table operations with an [`ExecutionState`], assuming
/// that the [`FunctionId`] for the table is known ahead of time.
#[derive(Debug, PartialEq, Eq, Hash, Clone)]
pub struct TableAction {
    table: TableId,
    table_math: SchemaMath,
    default: Option<MergeVal>,
    timestamp: CounterId,
    kind: TableKind,
}

impl TableAction {
    /// Create a new `TableAction` to be used later.
    /// This requires access to the `egglog_bridge::EGraph`.
    pub fn new(egraph: &EGraph, func: FunctionId) -> TableAction {
        let func_info = &egraph.funcs[func];
        let kind = match &func_info.default_val {
            DefaultVal::FreshId => TableKind::Constructor,
            DefaultVal::Fail | DefaultVal::Const(_) => TableKind::Function,
        };
        TableAction {
            table: func_info.table,
            table_math: SchemaMath {
                func_cols: func_info.schema.len(),
                subsume: func_info.can_subsume,
            },
            default: match &func_info.default_val {
                DefaultVal::FreshId => Some(MergeVal::Counter(egraph.id_counter)),
                DefaultVal::Fail => None,
                DefaultVal::Const(val) => Some(MergeVal::Constant(*val)),
            },
            timestamp: egraph.timestamp_counter,
            kind,
        }
    }

    /// Whether this table is a `Function` (no auto-insert) or a
    /// `Constructor` (mints a fresh eclass id on miss).
    pub fn kind(&self) -> TableKind {
        self.kind
    }

    /// Number of input columns (schema minus the trailing output
    /// column).
    pub fn input_arity(&self) -> usize {
        self.table_math.func_cols - 1
    }

    /// Look up a row and return its return-value column, or `None` if the
    /// key is not present. **This is a pure read**: it never inserts a row,
    /// regardless of the table's configured [`DefaultVal`].
    ///
    /// For the lookup-or-insert behavior that mints fresh eclass IDs for
    /// constructors, use [`TableAction::lookup_or_insert`].
    pub fn lookup(&self, state: &ExecutionState, key: &[Value]) -> Option<Value> {
        state
            .get_table(self.table)
            .get_row(key)
            .map(|row| row.vals[self.table_math.ret_val_col()])
    }

    /// Return the current number of rows in this table.
    pub fn row_count(&self, state: &ExecutionState) -> usize {
        state.get_table(self.table).len()
    }

    /// Iterate this table's rows, calling `f` on each function row.
    /// Mirrors [`EGraph::for_each`] but reaches the table through an
    /// [`ExecutionState`] — so it's callable from primitive bodies via
    /// the typed `Read`-style API.
    pub fn for_each(&self, state: &ExecutionState, mut f: impl FnMut(ScanEntry<'_>)) {
        self.for_each_while(state, |row| {
            f(row);
            true
        });
    }

    /// Like [`TableAction::for_each`], but stops as soon as `f`
    /// returns `false`.
    pub fn for_each_while(&self, state: &ExecutionState, mut f: impl FnMut(ScanEntry<'_>) -> bool) {
        let schema_math = self.table_math;
        let imp = state.get_table(self.table);
        let all = imp.all();
        let mut cur = Offset::new(0);
        let mut buf = TaggedRowBuffer::new(imp.spec().arity());
        macro_rules! drain_buf {
            ($buf:expr) => {
                for (_, row) in $buf.non_stale() {
                    let subsumed =
                        schema_math.subsume && row[schema_math.subsume_col()] == SUBSUMED;
                    if !f(ScanEntry {
                        vals: &row[0..schema_math.func_cols],
                        subsumed,
                    }) {
                        return;
                    }
                }
                $buf.clear();
            };
        }
        while let Some(next) = imp.scan_bounded(all.as_ref(), cur, 32, &mut buf) {
            drain_buf!(buf);
            cur = next;
        }
        drain_buf!(buf);
    }

    /// Look up a row, inserting the configured default value if absent.
    /// For constructor tables this mints a fresh eclass ID; for custom
    /// functions (no default) this behaves identically to
    /// [`TableAction::lookup`].
    ///
    /// This is a write operation — only safe in action contexts. See
    /// issue #772.
    pub fn lookup_or_insert(&self, state: &mut ExecutionState, key: &[Value]) -> Option<Value> {
        match self.default {
            Some(default) => {
                let timestamp =
                    MergeVal::Constant(Value::from_usize(state.read_counter(self.timestamp)));
                let mut merge_vals = SmallVec::<[MergeVal; 3]>::new();
                SchemaMath {
                    func_cols: 1,
                    ..self.table_math
                }
                .write_table_row(
                    &mut merge_vals,
                    RowVals {
                        timestamp,
                        subsume: self
                            .table_math
                            .subsume
                            .then_some(MergeVal::Constant(NOT_SUBSUMED)),
                        ret_val: Some(default),
                    },
                );
                Some(
                    state.predict_val(self.table, key, merge_vals.iter().copied())
                        [self.table_math.ret_val_col()],
                )
            }
            None => self.lookup(state, key),
        }
    }

    /// Insert a row into this table.
    pub fn insert(&self, state: &mut ExecutionState, row: impl Iterator<Item = Value>) {
        let ts = Value::from_usize(state.read_counter(self.timestamp));
        let mut scratch = row.collect::<SmallVec<[_; 8]>>();
        self.table_math.write_table_row(
            &mut scratch,
            RowVals {
                timestamp: ts,
                subsume: self.table_math.subsume.then_some(NOT_SUBSUMED),
                ret_val: None,
            },
        );
        state.stage_insert(self.table, &scratch);
    }

    /// Delete a row from this table.
    pub fn remove(&self, state: &mut ExecutionState, key: &[Value]) {
        state.stage_remove(self.table, key);
    }

    /// Subsume a row in this table.
    pub fn subsume(&self, state: &mut ExecutionState, key: impl Iterator<Item = Value>) {
        let ts = Value::from_usize(state.read_counter(self.timestamp));
        let mut scratch = key.collect::<SmallVec<[_; 8]>>();

        let ret_val = self.lookup(state, &scratch).expect("subsume lookup failed");

        self.table_math.write_table_row(
            &mut scratch,
            RowVals {
                timestamp: ts,
                subsume: Some(SUBSUMED),
                ret_val: Some(ret_val),
            },
        );
        state.stage_insert(self.table, &scratch);
    }
}

/// A variant of `TableAction` for the union-find.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct UnionAction {
    table: TableId,
    timestamp: CounterId,
}

impl UnionAction {
    /// Create a new `UnionAction` to be used later.
    /// This requires access to the `egglog_bridge::EGraph`.
    pub fn new(egraph: &EGraph) -> UnionAction {
        UnionAction {
            table: egraph.uf_table,
            timestamp: egraph.timestamp_counter,
        }
    }

    /// Union two values.
    pub fn union(&self, state: &mut ExecutionState, x: Value, y: Value) {
        let ts = Value::from_usize(state.read_counter(self.timestamp));
        state.stage_insert(self.table, &[x, y, ts]);
    }
}

fn run_rules_impl(
    db: &mut Database,
    rule_info: &mut DenseIdMapWithReuse<RuleId, RuleInfo>,
    rules: &[RuleId],
    next_ts: Timestamp,
    report_level: ReportLevel,
) -> Result<RuleSetReport> {
    for rule in rules {
        let info = &mut rule_info[*rule];
        if info.cached_plan.is_none() {
            info.cached_plan = Some(info.query.build_cached_plan(db, &info.desc)?);
        }
    }
    let mut rsb = db.new_rule_set();
    for rule in rules {
        let info = &mut rule_info[*rule];
        let cached_plan = info.cached_plan.as_ref().unwrap();
        info.query
            .add_rules_from_cached(&mut rsb, info.last_run_at, cached_plan);
        info.last_run_at = next_ts;
    }
    let ruleset = rsb.build();
    Ok(db.run_rule_set(&ruleset, report_level))
}

// These markers are just used to make it easy to distinguish time spent in
// incremental vs. nonincremental rebuilds in time-based profiles.

#[inline(never)]
fn marker_incremental_rebuild<R>(f: impl FnOnce() -> R) -> R {
    f()
}

#[inline(never)]
fn marker_nonincremental_rebuild<R>(f: impl FnOnce() -> R) -> R {
    f()
}

/// A useful type definition for external functions that need to pass data
/// to outside code, such as `Panic`.
pub type SideChannel<T> = Arc<Mutex<Option<T>>>;

/// An external function used to grab a value out of the database matching a
/// particular query.
//
// TODO: once we have parallelism wired in, we'll want to replace this with a
// more efficient solution (e.g. one based on crossbeam or arcswap).
/// This is a variant on [`Panic`] that avoids eager construction of the panic message.
///
/// The main thing this is used for is to avoid constructing the panic message ahead of time during
/// a call to [`RuleBuilder::call_external_func`]; these panic messages are often quite rare and
/// may never need to be constructed at all. Furthermore, a closure to produce the panic message in
/// most cases need only close over a few cheap-to-clone values.
///
/// The downside of this, and why we do not use it everywhere, is that there's no natural "key"
/// that we can use to cache duplicate panic messages. We would need a more complex API to support
/// both and fully replace our use of `Panic`.
struct LazyPanic<F>(Arc<Lazy<String, F>>, SideChannel<String>);

impl<F: FnOnce() -> String + Send> ExternalFunction for LazyPanic<F> {
    fn invoke(&self, state: &mut core_relations::ExecutionState, args: &[Value]) -> Option<Value> {
        assert!(args.is_empty());
        state.trigger_early_stop();
        let mut guard = self.1.lock().unwrap();
        if guard.is_none() {
            *guard = Some(Lazy::force(&self.0).clone());
        }
        None
    }
}

impl<F> Clone for LazyPanic<F> {
    fn clone(&self) -> Self {
        LazyPanic(self.0.clone(), self.1.clone())
    }
}

/// An external function used to store a message when a panic occurs.
//
// TODO: once we have parallelism wired in, we'll want to replace this with a
// more efficient solution (e.g. one based on crossbeam or arcswap).
#[derive(Clone)]
struct Panic(String, SideChannel<String>);

impl EGraph {
    /// Create a new `ExternalFunction` that panics with the given message.
    pub fn new_panic(&mut self, message: String) -> ExternalFunctionId {
        *self
            .panic_funcs
            .entry(message.to_string())
            .or_insert_with(|| {
                let panic = Panic(message, self.panic_message.clone());
                self.db.add_external_function(Box::new(panic))
            })
    }

    pub fn new_panic_lazy(
        &mut self,
        message: impl FnOnce() -> String + Send + 'static,
    ) -> ExternalFunctionId {
        let lazy = Lazy::new(message);
        let panic = LazyPanic(Arc::new(lazy), self.panic_message.clone());
        self.db.add_external_function(Box::new(panic))
    }
}

impl ExternalFunction for Panic {
    fn invoke(&self, state: &mut core_relations::ExecutionState, args: &[Value]) -> Option<Value> {
        // TODO (egglog feature): change this to support interpolating panic messages
        assert!(args.is_empty());

        state.trigger_early_stop();
        let mut guard = self.1.lock().unwrap();
        if guard.is_none() {
            *guard = Some(self.0.clone());
        }
        None
    }
}

/// Heuristic for deciding whether to do an incremental or nonincremental
/// rebuild for a given table.
fn incremental_rebuild(uf_size: usize, table_size: usize, parallel: bool) -> bool {
    if parallel {
        uf_size <= (table_size / 16)
    } else {
        uf_size <= (table_size / 8)
    }
}

pub(crate) const SUBSUMED: Value = Value::new_const(1);
pub(crate) const NOT_SUBSUMED: Value = Value::new_const(0);
fn combine_subsumed(v1: Value, v2: Value) -> Value {
    std::cmp::max(v1, v2)
}

/// A struct helping with some calculations of where some information is stored at the
/// core-relations Table level for a given function.
///
/// Functions can have multiple "output columns" in the underlying core-relations layer depending
/// on whether different features are enabled. Roughly, tables are laid out as:
///
/// > `[key0, ..., keyn, return value, timestamp, subsume?]`
///
/// Where there are `n+1` key columns and columns marked with a question mark are optional,
/// depending on the egraph and table-level configuration.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
struct SchemaMath {
    /// Whether or not the table is enabled for subsumption.
    subsume: bool,
    /// The number of columns in the function (including the return value).
    func_cols: usize,
}

/// A struct containing possible non-key portions of a table row. To be used with
/// [`SchemaMath::write_table_row`].
///
/// This is the write side (building a row); [`ScanEntry`] is the
/// read side (a row yielded by a table scan).
struct RowVals<T> {
    /// The timestamp for the row.
    timestamp: T,
    /// The subsumption tag for the row. Only relevant if the table has subsumption enabled.
    subsume: Option<T>,
    /// The return value of the row. Return values are mandatory but callers may have already
    /// filled it in.
    ret_val: Option<T>,
}

/// A raw row yielded by a table scan; `vals` includes the trailing
/// output/eclass column.
///
/// Public so the `egglog` crate can consume it, but **do not re-export
/// it from `egglog`'s public API** — the user-facing row types are
/// `egglog::FunctionEntry` and `egglog::Enode`.
#[derive(Clone, Debug)]
pub struct ScanEntry<'a> {
    pub vals: &'a [Value],
    pub subsumed: bool,
}

impl SchemaMath {
    fn write_table_row<T: Clone>(
        &self,
        row: &mut impl HasResizeWith<T>,
        RowVals {
            timestamp,
            subsume,
            ret_val,
        }: RowVals<T>,
    ) {
        row.resize_with(self.table_columns(), || timestamp.clone());
        row[self.ts_col()] = timestamp;
        if let Some(ret_val) = ret_val {
            row[self.ret_val_col()] = ret_val;
        }
        if let Some(subsume) = subsume {
            row[self.subsume_col()] = subsume;
        } else {
            assert!(
                !self.subsume,
                "subsume flag must be provided if subsumption is enabled"
            );
        }
    }

    fn num_keys(&self) -> usize {
        self.func_cols - 1
    }

    fn table_columns(&self) -> usize {
        self.func_cols + 1 /* timestamp */ + if self.subsume { 1 } else { 0 }
    }

    fn ret_val_col(&self) -> usize {
        self.func_cols - 1
    }

    fn ts_col(&self) -> usize {
        self.func_cols
    }

    #[track_caller]
    fn subsume_col(&self) -> usize {
        assert!(self.subsume);
        self.func_cols + 1
    }
}

#[derive(Error, Debug)]
#[error("Panic: {0}")]
struct PanicError(String);

/// Basic ad-hoc polymorphism around `resize_with` in order to get [`SchemaMath::write_table_row`]
/// to work with both `Vec` and `SmallVec`.
trait HasResizeWith<T>:
    AsMut<[T]> + AsRef<[T]> + Index<usize, Output = T> + IndexMut<usize, Output = T>
{
    fn resize_with<F>(&mut self, new_size: usize, f: F)
    where
        F: FnMut() -> T;
}

impl<T> HasResizeWith<T> for Vec<T> {
    fn resize_with<F>(&mut self, new_size: usize, f: F)
    where
        F: FnMut() -> T,
    {
        self.resize_with(new_size, f);
    }
}

impl<T, A: smallvec::Array<Item = T>> HasResizeWith<T> for SmallVec<A> {
    fn resize_with<F>(&mut self, new_size: usize, f: F)
    where
        F: FnMut() -> T,
    {
        self.resize_with(new_size, f);
    }
}