Expose extraction support APIs

core-relations/src/action/mod.rs

@@ -12,18 +12,20 @@ use std::{
 
 use crate::{
     common::HashMap,
-    free_join::{invoke_batch, invoke_batch_assign},
+    free_join::{get_column_index_from_tableinfo, invoke_batch, invoke_batch_assign},
     numeric_id::{DenseIdMap, NumericId},
 };
 use egglog_concurrency::NotificationList;
 use smallvec::SmallVec;
 
 use crate::{
-    BaseValues, ContainerValues, ExternalFunctionId, WrappedTable,
+    BaseValues, ContainerValues, ExternalFunctionId, Offset, WrappedTable,
     common::Value,
     free_join::{CounterId, Counters, ExternalFunctions, TableId, TableInfo, Variable},
+    offsets::Subset,
     pool::{Clear, Pooled, with_pool_set},
-    table_spec::{ColumnId, MutationBuffer},
+    row_buffer::TaggedRowBuffer,
+    table_spec::{ColumnId, Constraint, MutationBuffer},
 };
 
 use self::mask::{Mask, MaskIter, ValueSource};
@@ -483,6 +485,66 @@ impl<'a> ExecutionState<'a> {
         &self.db.table_info[table].table
     }
 
+    /// Iterate over visible rows in `table` whose `col` equals `value`.
+    ///
+    /// Cacheable columns use the table's lazy column index; uncacheable columns
+    /// fall back to scanning with the equality constraint. The callback
+    /// receives each matching row as a full value slice.
+    pub fn for_each_matching_col(
+        &self,
+        table: TableId,
+        col: ColumnId,
+        value: Value,
+        mut f: impl FnMut(&[Value]),
+    ) {
+        let table_info = &self.db.table_info[table];
+        let constraint = Constraint::EqConst { col, val: value };
+        let (mut subset, _fast, mut slow) = table_info
+            .table
+            .split_fast_slow(std::slice::from_ref(&constraint));
+
+        debug_assert!(slow.iter().all(|c| matches!(c, Constraint::EqConst { .. })));
+
+        if !*table_info
+            .spec
+            .uncacheable_columns
+            .get(col)
+            .unwrap_or(&false)
+        {
+            let index = get_column_index_from_tableinfo(table_info, col);
+            match index.get().unwrap().get_subset(&value) {
+                Some(s) => {
+                    with_pool_set(|ps| subset.intersect(s, &ps.get_pool()));
+                }
+                None => {
+                    subset = Subset::empty();
+                }
+            }
+            slow.clear();
+        }
+
+        let imp = &table_info.table;
+        let cols: SmallVec<[_; 8]> = (0..imp.spec().arity()).map(ColumnId::from_usize).collect();
+        let mut cur = Offset::new(0);
+        let mut buf = TaggedRowBuffer::new_inline(imp.spec().arity());
+
+        macro_rules! drain_buf {
+            ($buf:expr) => {
+                for (_, row) in $buf.non_stale() {
+                    f(row);
+                }
+                $buf.clear();
+            };
+        }
+
+        while let Some(next) = imp.scan_project(subset.as_ref(), &cols, cur, 1024, &slow, &mut buf)
+        {
+            drain_buf!(buf);
+            cur = next;
+        }
+        drain_buf!(buf);
+    }
+
     /// Get the human-readable name for a table, if one exists.
     pub fn table_name(&self, table: TableId) -> Option<&'a str> {
         self.db.table_info[table].name()

core-relations/src/free_join/mod.rs

@@ -868,7 +868,10 @@ fn get_index_from_tableinfo(table_info: &TableInfo, cols: &[ColumnId]) -> HashIn
 /// The core logic behind getting and updating a column index.
 ///
 /// This is the single-column analog to [`get_index_from_tableinfo`].
-fn get_column_index_from_tableinfo(table_info: &TableInfo, col: ColumnId) -> HashColumnIndex {
+pub(crate) fn get_column_index_from_tableinfo(
+    table_info: &TableInfo,
+    col: ColumnId,
+) -> HashColumnIndex {
     let index: Arc<_> = table_info.column_indexes.get_or_insert(col, || {
         Arc::new(ResettableOnceLock::new(Index::new(
             vec![col],

egglog-bridge/src/lib.rs

@@ -1375,6 +1375,28 @@ impl TableAction {
         drain_buf!(buf);
     }
 
+    /// Iterate over rows whose output column is `value`.
+    ///
+    /// This reaches the table through an [`ExecutionState`], so read-capable
+    /// state wrappers can use the same lazy column indexes as direct backend
+    /// callers without exposing lower-level scan buffers or projection details.
+    pub fn for_each_output_value(
+        &self,
+        state: &ExecutionState,
+        value: Value,
+        mut f: impl FnMut(ScanEntry<'_>),
+    ) {
+        let schema_math = self.table_math;
+        let output_col = ColumnId::from_usize(self.input_arity());
+        state.for_each_matching_col(self.table, output_col, value, |row| {
+            let subsumed = schema_math.subsume && row[schema_math.subsume_col()] == SUBSUMED;
+            f(ScanEntry {
+                vals: &row[0..schema_math.func_cols],
+                subsumed,
+            });
+        });
+    }
+
     /// 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

src/exec_state.rs

@@ -395,6 +395,34 @@ pub trait Read<'a, 'db: 'a>: Core<'a, 'db> + RegistrySealed<'a, 'db> {
         Ok(())
     }
 
+    /// Call `f` on each [`Enode`] of a constructor / relation table whose
+    /// output eclass is `eclass`.
+    ///
+    /// This uses the backend's indexed output-column lookup instead of scanning
+    /// the whole constructor table. Errors with `WrongSubtype` if `name` is a
+    /// function.
+    fn enodes_for_eclass(
+        &self,
+        name: &str,
+        eclass: Value,
+        mut f: impl FnMut(Enode<'_>),
+    ) -> Result<(), Error> {
+        let action = lookup_action(self.registry(), name)?;
+        check_subtype(name, &action, TableKind::Constructor, "constructor")?;
+        action.for_each_output_value(self.es(), eclass, |row| {
+            let (eclass, children) = row
+                .vals
+                .split_last()
+                .expect("constructor row has at least an eclass column");
+            f(Enode {
+                children,
+                eclass: *eclass,
+                subsumed: row.subsumed,
+            });
+        });
+        Ok(())
+    }
+
     /// Call `f` on each [`FunctionEntry`] of a function table. Errors
     /// with `WrongSubtype` if `name` is a constructor. To stop early,
     /// use [`Read::function_entries_while`].