diff --git a/core-relations/src/containers/mod.rs b/core-relations/src/containers/mod.rs index d56658c14..40ecda60e 100644 --- a/core-relations/src/containers/mod.rs +++ b/core-relations/src/containers/mod.rs @@ -490,6 +490,7 @@ impl ContainerValues { id: ContainerValueId, id_counter: CounterId, merge_fn: impl MergeFn + 'static, + base_values: crate::BaseValues, ) -> ContainerValueId { let id = self.container_ids.insert(TypeId::of::(), id); self.data.get_or_insert(id, || { @@ -497,6 +498,7 @@ impl ContainerValues { id, Box::new(merge_fn), id_counter, + base_values, )) }); assert!( diff --git a/core-relations/src/containers/sequence.rs b/core-relations/src/containers/sequence.rs index 904387656..950a06efc 100644 --- a/core-relations/src/containers/sequence.rs +++ b/core-relations/src/containers/sequence.rs @@ -14,7 +14,7 @@ use crossbeam_queue::SegQueue; use rustc_hash::FxHasher; use crate::{ - ExecutionState, Offset, RowId, SequenceTable, Subset, TableChange, TableVersion, + BaseValues, ExecutionState, Offset, RowId, SequenceTable, Subset, TableChange, TableVersion, TaggedRowBuffer, Value, common::{HashMap, IndexSet, ShardData, ShardId}, numeric_id::{DenseIdMap, NumericId}, @@ -37,11 +37,11 @@ use super::{ /// directly. pub trait SequenceContainerValue: ContainerValue { /// Append the canonical serialized key to `out`. - fn encode_sequence(&self, out: &mut Vec); + fn encode_sequence(&self, base_values: &BaseValues, out: &mut Vec); /// Reconstruct the Rust container used by slow primitives and external /// APIs. - fn decode_sequence(sequence: &[Value]) -> Self; + fn decode_sequence(sequence: &[Value], base_values: &BaseValues) -> Self; /// Return the fast primitive view of the serialized key. /// @@ -66,12 +66,13 @@ pub trait SequenceContainerValue: ContainerValue { /// the slow deserialize/rebuild/serialize round trip. fn rebuild_sequence( sequence: &[Value], + base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { - let mut container = Self::decode_sequence(sequence); + let mut container = Self::decode_sequence(sequence, base_values); if container.rebuild_contents(rebuilder) { - container.encode_sequence(out); + container.encode_sequence(base_values, out); true } else { false @@ -79,18 +80,20 @@ pub trait SequenceContainerValue: ContainerValue { } } -#[derive(Clone, Copy)] +#[derive(Clone)] struct SequenceCodec { - encode: fn(&C, &mut Vec), - decode: fn(&[Value]) -> C, + base_values: BaseValues, + encode: fn(&C, &BaseValues, &mut Vec), + decode: fn(&[Value], &BaseValues) -> C, values: for<'a> fn(&'a [Value]) -> &'a [Value], - rebuild: fn(&[Value], &dyn ValueRebuilder, &mut Vec) -> bool, + rebuild: fn(&[Value], &BaseValues, &dyn ValueRebuilder, &mut Vec) -> bool, marker: PhantomData C>, } impl SequenceCodec { - fn new() -> Self { + fn new(base_values: BaseValues) -> Self { Self { + base_values, encode: C::encode_sequence, decode: C::decode_sequence, values: C::sequence_values, @@ -156,6 +159,7 @@ impl SequenceContainerEnv { id: ContainerValueId, merge: Box, counter: crate::CounterId, + base_values: BaseValues, ) -> Self { let table_merge = move |state: &mut ExecutionState, current: &[Value], @@ -185,7 +189,7 @@ impl SequenceContainerEnv { counter, table, reverse, - codec: SequenceCodec::new(), + codec: SequenceCodec::new(base_values), } } } @@ -206,12 +210,15 @@ impl SequenceContainerEnv { } pub(super) fn get_container(&self, value: Value) -> Option { - Some((self.codec.decode)(self.get_key(value)?)) + Some((self.codec.decode)( + self.get_key(value)?, + &self.codec.base_values, + )) } pub(super) fn get_or_insert(&self, container: &C, exec_state: &mut ExecutionState) -> Value { let mut key = Vec::new(); - (self.codec.encode)(container, &mut key); + (self.codec.encode)(container, &self.codec.base_values, &mut key); self.get_or_insert_key(&key, exec_state) } @@ -254,13 +261,14 @@ impl SequenceContainerEnv { ) -> Option { Some((self.codec.decode)( self.get_key_with_predictions(exec_state, value)?, + &self.codec.base_values, )) } pub(super) fn for_each(&self, f: &mut impl FnMut(&C, Value)) { self.table .scan_key_values(self.table.all().as_ref(), |_, key, values| { - let container = (self.codec.decode)(key); + let container = (self.codec.decode)(key, &self.codec.base_values); f(&container, values[0]); }); } @@ -383,6 +391,7 @@ impl DynamicContainerEnv for SequenceContainerEnv { exec_state: &mut ExecutionState, ) -> ContainerRebuildSummary { let rebuild_sequence = self.codec.rebuild; + let base_values = &self.codec.base_values; let stable_changed_ids = SegQueue::new(); let previous = self.table.version(); let rebuild_row = |row: &[Value], rebuilt: &mut Vec| { @@ -393,7 +402,7 @@ impl DynamicContainerEnv for SequenceContainerEnv { let key = &row[..key_end]; let old_id = row[key_end]; let new_id = rebuilder.rebuild_val(old_id); - let key_changed = rebuild_sequence(key, rebuilder, rebuilt); + let key_changed = rebuild_sequence(key, base_values, rebuilder, rebuilt); if !key_changed && new_id == old_id { debug_assert!(rebuilt.is_empty()); return false; @@ -454,7 +463,12 @@ impl DynamicContainerEnv for SequenceContainerEnv { ) -> Option { let key = self.get_key_with_predictions(exec_state, value)?.to_vec(); let mut rebuilt = Vec::new(); - if !(self.codec.rebuild)(&key, &ClosureRebuilder { remap }, &mut rebuilt) { + if !(self.codec.rebuild)( + &key, + &self.codec.base_values, + &ClosureRebuilder { remap }, + &mut rebuilt, + ) { return Some(value); } Some(self.get_or_insert_key(&rebuilt, exec_state)) @@ -488,11 +502,11 @@ mod tests { } impl SequenceContainerValue for TestSequence { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &crate::BaseValues, out: &mut Vec) { out.extend_from_slice(&self.0); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &crate::BaseValues) -> Self { Self(sequence.to_vec()) } @@ -502,6 +516,7 @@ mod tests { fn rebuild_sequence( sequence: &[Value], + _base_values: &crate::BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -525,6 +540,7 @@ mod tests { crate::ContainerValueId::new(0), Box::new(|_state: &mut ExecutionState, left, right| left.min(right)), counter, + db.base_values().clone(), ) } diff --git a/core-relations/src/containers/tests.rs b/core-relations/src/containers/tests.rs index 94263d5f6..0ec1c8515 100644 --- a/core-relations/src/containers/tests.rs +++ b/core-relations/src/containers/tests.rs @@ -54,11 +54,11 @@ impl ContainerValue for LegacyVecContainer { } impl SequenceContainerValue for VecContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &crate::BaseValues, out: &mut Vec) { out.extend_from_slice(&self.0); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &crate::BaseValues) -> Self { Self(sequence.to_vec()) } @@ -68,6 +68,7 @@ impl SequenceContainerValue for VecContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &crate::BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { diff --git a/core-relations/src/free_join/mod.rs b/core-relations/src/free_join/mod.rs index 7b9f8ac2b..83a47b16a 100644 --- a/core-relations/src/free_join/mod.rs +++ b/core-relations/src/free_join/mod.rs @@ -602,11 +602,15 @@ impl Database { self.validate_dependencies(&read_deps, "read"); self.validate_dependencies(&write_deps, "write"); + let base_values = self.base_values.clone(); let participant = self.allocate_storage_id(); let container = crate::ContainerValueId::from_table_id(participant); - let registered = self - .container_values - .register_sequence_type::(container, id_counter, merge_fn); + let registered = self.container_values.register_sequence_type::( + container, + id_counter, + merge_fn, + base_values, + ); assert_eq!(registered, container); self.deps .add_participant(participant, read_deps, write_deps); diff --git a/src/exec_state.rs b/src/exec_state.rs index 060ce1a4a..3583446a8 100644 --- a/src/exec_state.rs +++ b/src/exec_state.rs @@ -41,7 +41,7 @@ use crate::core_relations::{ use crate::{ ast::{FunctionSubtype, Literal, ResolvedExpr}, core::ResolvedCall, - sort::{F, S}, + sort::{F, FunctionContainer, FunctionSequence, PreparedFunction, ResolvedFunction, S}, typechecking::FuncType, }; use egglog_bridge::{ActionRegistry, TableAction, TableKind}; @@ -238,14 +238,56 @@ pub trait Core<'a, 'db: 'a>: Internal<'a, 'db> { /// context — there is no `ctx` parameter to lie about. /// /// [`EGraph::add_pure_primitive`]: crate::EGraph::add_pure_primitive - fn apply_function( + fn apply_function(&mut self, fc: &FunctionContainer, args: &[Value]) -> Option { + let ctx = self.ctx(); + let mut pure = PureState::wrap(self.raw_exec_state(), ctx); + fc.apply(&mut pure, args) + } + + /// Dispatch an `UnstableFn` directly from its serialized sequence. + /// + /// This convenience path prepares the function for one application. The + /// partial arguments are copied before invoking the target because the + /// callback may intern another container and grow the prediction storage + /// backing the borrowed sequence. Higher-order loops should call + /// [`Core::prepare_function`] once and reuse the result; rebuild and + /// occurrence indexing never resolve the descriptor. + fn apply_function_value(&mut self, function: Value, args: &[Value]) -> Option { + let prepared = self.prepare_function(function)?; + self.apply_prepared_function(&prepared, args) + } + + /// Resolve and copy the immutable portion of an `UnstableFn` once so a + /// higher-order primitive can reuse it for many callbacks. + fn prepare_function(&self, function: Value) -> Option { + let (descriptor, partial_args) = + self.with_container_sequence::(function, |sequence| { + let sequence = FunctionSequence::parse(sequence); + (sequence.descriptor(), sequence.args().to_vec()) + })?; + let resolved: ResolvedFunction = self.base_values().unwrap(descriptor); + Some(PreparedFunction::new(resolved, partial_args)) + } + + /// Invoke a function prepared by [`Core::prepare_function`]. + fn apply_prepared_function( &mut self, - fc: &crate::sort::FunctionContainer, + function: &PreparedFunction, args: &[Value], ) -> Option { let ctx = self.ctx(); let mut pure = PureState::wrap(self.raw_exec_state(), ctx); - fc.apply(&mut pure, args) + function.apply(&mut pure, args) + } + + /// Resolve the descriptor of a serialized `UnstableFn` without + /// reconstructing its Rust container or partial-argument sorts. + fn resolve_function_value(&self, function: Value) -> Option { + let descriptor = self + .with_container_sequence::(function, |sequence| { + FunctionSequence::parse(sequence).descriptor() + })?; + Some(self.base_values().unwrap(descriptor)) } /// Dispatch an already type-specialized primitive in the current diff --git a/src/sort/fn.rs b/src/sort/fn.rs index fb45019ed..1465ef4dd 100644 --- a/src/sort/fn.rs +++ b/src/sort/fn.rs @@ -6,13 +6,15 @@ //! To create a function value, use the `(unstable-fn "name" [])` primitive and to apply it use the `(unstable-app function arg1 arg2 ...)` primitive. //! The number of args must match the number of arguments in the function sort. //! -//! The value is stored similar to the `vec` sort, as an index into a set, where each item in -//! the set is a `(Symbol, Vec<(Sort, Value)>)` pairs. The Symbol is the function name, and the `Vec<(Sort, Value)>` is -//! the list of partially applied arguments. +//! Each value is interned as a sequence containing an opaque resolved-function +//! descriptor, a self-describing rebuild mask, and the partially applied +//! arguments. The mask lets congruence closure rebuild function keys without +//! resolving their descriptors or reconstructing Rust container values. use std::any::TypeId; use std::sync::Mutex; use crate::exec_state::Internal; +use crate::numeric_id::NumericId; use enum_map::EnumMap; use super::*; @@ -29,6 +31,148 @@ pub struct FunctionContainer( pub ExternalFunctionId, ); +const FUNCTION_REBUILD_MASK_BITS: usize = 31; +const FUNCTION_REBUILD_MASK: u32 = (1 << FUNCTION_REBUILD_MASK_BITS) - 1; + +/// A validated view over an `UnstableFn` sequence key. +/// +/// The key is `[descriptor, argc, mask..., args...]`. `argc` determines both +/// the number of 31-bit rebuild-mask words and the number of trailing args, so +/// parsing never needs to resolve the descriptor through `BaseValues`. +#[derive(Clone, Copy, Debug)] +pub(crate) struct FunctionSequence<'a> { + descriptor: Value, + masks: &'a [Value], + args: &'a [Value], +} + +/// A resolved `UnstableFn` plus an owned copy of its partial arguments. +/// +/// Higher-order container primitives prepare this once before their element +/// loop, avoiding a descriptor lookup and sequence copy for every callback. +#[derive(Clone, Debug)] +pub struct PreparedFunction { + resolved: ResolvedFunction, + partial_args: Vec, +} + +impl PreparedFunction { + pub(crate) fn new(resolved: ResolvedFunction, partial_args: Vec) -> Self { + Self { + resolved, + partial_args, + } + } + + pub(crate) fn apply<'a, 'db>( + &self, + state: &mut crate::PureState<'a, 'db>, + args: &[Value], + ) -> Option + where + 'db: 'a, + { + let mut combined = Vec::with_capacity(self.partial_args.len() + args.len()); + combined.extend_from_slice(&self.partial_args); + combined.extend_from_slice(args); + self.resolved.apply(state, &combined) + } +} + +impl<'a> FunctionSequence<'a> { + pub(crate) fn parse(sequence: &'a [Value]) -> Self { + let (&descriptor, rest) = sequence + .split_first() + .expect("serialized FunctionContainer must include its descriptor"); + let (&argc, rest) = rest + .split_first() + .expect("serialized FunctionContainer must include its argument count"); + let argc = argc.index(); + assert!( + argc <= FUNCTION_REBUILD_MASK as usize, + "serialized FunctionContainer argument count exceeds its 31-bit encoding" + ); + let mask_words = argc.div_ceil(FUNCTION_REBUILD_MASK_BITS); + let expected = mask_words + .checked_add(argc) + .expect("serialized FunctionContainer length overflow"); + assert_eq!( + rest.len(), + expected, + "serialized FunctionContainer length does not match its argument count" + ); + let (masks, args) = rest.split_at(mask_words); + for mask in masks { + assert_eq!( + mask.rep() & !FUNCTION_REBUILD_MASK, + 0, + "serialized FunctionContainer mask uses its reserved high bit" + ); + } + if let Some(last) = masks.last() + && !argc.is_multiple_of(FUNCTION_REBUILD_MASK_BITS) + { + let used = argc % FUNCTION_REBUILD_MASK_BITS; + let used_mask = (1u32 << used) - 1; + assert_eq!( + last.rep() & !used_mask, + 0, + "serialized FunctionContainer mask sets bits beyond its argument count" + ); + } + Self { + descriptor, + masks, + args, + } + } + + pub(crate) fn descriptor(self) -> Value { + self.descriptor + } + + pub(crate) fn args(self) -> &'a [Value] { + self.args + } + + fn rebuilds_arg(self, index: usize) -> bool { + debug_assert!(index < self.args.len()); + self.masks[index / FUNCTION_REBUILD_MASK_BITS].rep() + & (1 << (index % FUNCTION_REBUILD_MASK_BITS)) + != 0 + } +} + +fn encode_function_sequence( + descriptor: Value, + partial_arcsorts: &[ArcSort], + args: &[Value], + out: &mut Vec, +) { + assert_eq!( + partial_arcsorts.len(), + args.len(), + "FunctionContainer sort and argument counts must match" + ); + assert!( + args.len() <= FUNCTION_REBUILD_MASK as usize, + "FunctionContainer argument count exceeds its 31-bit encoding" + ); + out.push(descriptor); + out.push(Value::from_usize(args.len())); + let mask_words = args.len().div_ceil(FUNCTION_REBUILD_MASK_BITS); + let mask_start = out.len(); + out.resize(mask_start + mask_words, Value::from_usize(0)); + for (index, sort) in partial_arcsorts.iter().enumerate() { + if sort.is_eq_sort() || sort.is_eq_container_sort() { + let word = mask_start + index / FUNCTION_REBUILD_MASK_BITS; + let bit = 1 << (index % FUNCTION_REBUILD_MASK_BITS); + out[word] = Value::from_usize(out[word].index() | bit); + } + } + out.extend_from_slice(args); +} + // implement hash and equality based on values only not arcsorts, since // arcsorts are not comparable and any two values that are equal must have the same sort @@ -69,6 +213,96 @@ impl ContainerValue for FunctionContainer { self.1.iter().map(|(_, v)| v).copied() } } + +impl SequenceContainerValue for FunctionContainer { + fn encode_sequence(&self, base_values: &BaseValues, out: &mut Vec) { + let partial_arcsorts = self + .1 + .iter() + .map(|(sort, _)| sort.clone()) + .collect::>(); + let args = self.1.iter().map(|(_, value)| *value).collect::>(); + let descriptor = base_values.get(ResolvedFunction { + id: self.0.clone(), + partial_arcsorts: partial_arcsorts.clone(), + name: self.2.clone(), + panic_id: self.3, + }); + encode_function_sequence(descriptor, &partial_arcsorts, &args, out); + } + + fn decode_sequence(sequence: &[Value], base_values: &BaseValues) -> Self { + let sequence = FunctionSequence::parse(sequence); + let ResolvedFunction { + id, + partial_arcsorts, + name, + panic_id, + } = base_values.unwrap(sequence.descriptor()); + assert_eq!( + partial_arcsorts.len(), + sequence.args().len(), + "FunctionContainer descriptor arity does not match its sequence" + ); + for (index, sort) in partial_arcsorts.iter().enumerate() { + assert_eq!( + sequence.rebuilds_arg(index), + sort.is_eq_sort() || sort.is_eq_container_sort(), + "FunctionContainer rebuild mask does not match its descriptor" + ); + } + Self( + id, + partial_arcsorts + .into_iter() + .zip(sequence.args().iter().copied()) + .collect(), + name, + panic_id, + ) + } + + fn sequence_values(sequence: &[Value]) -> &[Value] { + FunctionSequence::parse(sequence); + sequence + } + + fn visit_sequence_values(sequence: &[Value], visitor: &mut dyn FnMut(Value)) { + let sequence = FunctionSequence::parse(sequence); + for (index, value) in sequence.args().iter().copied().enumerate() { + if sequence.rebuilds_arg(index) { + visitor(value); + } + } + } + + fn rebuild_sequence( + sequence: &[Value], + _base_values: &BaseValues, + rebuilder: &dyn ValueRebuilder, + out: &mut Vec, + ) -> bool { + let parsed = FunctionSequence::parse(sequence); + let prefix_len = sequence.len() - parsed.args().len(); + out.extend_from_slice(&sequence[..prefix_len]); + let mut changed = false; + for (index, old) in parsed.args().iter().copied().enumerate() { + let new = if parsed.rebuilds_arg(index) { + rebuilder.rebuild_val(old) + } else { + old + }; + changed |= old != new; + out.push(new); + } + if changed { + true + } else { + out.clear(); + false + } + } +} #[derive(Debug)] pub struct FunctionSort { name: String, @@ -174,10 +408,10 @@ impl Sort for FunctionSort { } fn register_type(&self, backend: &mut egglog_bridge::EGraph) { - backend.register_container_ty::(); backend .base_values_mut() .register_type::(); + backend.register_sequence_container_ty::(); } fn as_arc_any(self: Arc) -> Arc { @@ -189,9 +423,15 @@ impl Sort for FunctionSort { } fn is_eq_container_sort(&self) -> bool { - self.inputs - .iter() - .any(|s| s.is_eq_sort() || s.is_eq_container_sort()) + // The sequence's rebuildable children are its captured prefix, whose + // sorts depend on the particular wrapped target and are not known when + // an outer container fixes its rebuild/indexing policy. Remaining + // inputs therefore cannot soundly answer this question: for example, + // `(Math, i64) -> Math` captured as `UnstableFn (i64) Math` still + // contains a rebuildable `Math`. Conservatively index every function + // identity in outer containers; the function sequence's mask keeps its + // own occurrence index and rebuild work precise. + true } fn serialized_name(&self, container_values: &ContainerValues, value: Value) -> String { @@ -437,11 +677,10 @@ struct Ctor { function: Arc, } -// `Ctor` (`unstable-fn "name" [...]`) builds a `FunctionContainer` and -// interns it via `register_container`. Container interning is idempotent, -// so it's safe in every context; declaring `State = PureState` -// permits this primitive inside rule queries, actions, and global -// contexts alike. +// `Ctor` (`unstable-fn "name" [...]`) serializes the descriptor and partial +// arguments directly into the sequence table. Container interning is +// idempotent, so it's safe in every context; declaring `State = PureState` +// permits this primitive inside rule queries, actions, and global contexts. impl Primitive for Ctor { fn name(&self) -> &str { &self.name @@ -463,24 +702,15 @@ impl PurePrim for Ctor { args: &[Value], ) -> Option { let (rf, args) = args.split_first().unwrap(); - let ResolvedFunction { - id, - partial_arcsorts, - name, - panic_id, - } = state.base_values().unwrap(*rf); + let resolved: ResolvedFunction = state.base_values().unwrap(*rf); self.function .partial_arcsorts .lock() .unwrap() - .extend(partial_arcsorts.iter().cloned()); - let args = partial_arcsorts - .iter() - .zip(args) - .map(|(b, x)| (b.clone(), *x)) - .collect(); - let y = FunctionContainer(id, args, name, panic_id); - Some(state.register_container(y)) + .extend(resolved.partial_arcsorts.iter().cloned()); + let mut key = Vec::new(); + encode_function_sequence(*rf, &resolved.partial_arcsorts, args, &mut key); + Some(state.register_container_sequence::(&key)) } } @@ -498,11 +728,14 @@ pub struct ResolvedFunction { /// than the calling thread unwinding. pub panic_id: ExternalFunctionId, } -// implement equality and hash based on id and arcsort names, since arcsorts are not comparable +// Implement equality and hash based on id, user-visible name, and ArcSort +// names. ArcSort trait objects are not directly comparable. The panic id is a +// build-site implementation detail and intentionally remains excluded. impl PartialEq for ResolvedFunction { fn eq(&self, other: &Self) -> bool { self.id == other.id + && self.name == other.name && self .partial_arcsorts .iter() @@ -521,6 +754,7 @@ impl Eq for ResolvedFunction {} impl Hash for ResolvedFunction { fn hash(&self, state: &mut H) { self.id.hash(state); + self.name.hash(state); for s in &self.partial_arcsorts { s.name().hash(state); } @@ -588,12 +822,20 @@ impl PurePrim for Apply { args: &[Value], ) -> Option { let (fc_val, args) = args.split_first().unwrap(); - let fc = state - .container_values() - .get_val::(*fc_val) - .unwrap() - .clone(); - state.apply_function(&fc, args) + state.apply_function_value(*fc_val, args) + } +} + +impl ResolvedFunction { + pub(crate) fn apply<'a, 'db>( + &self, + state: &mut crate::PureState<'a, 'db>, + args: &[Value], + ) -> Option + where + 'db: 'a, + { + apply_resolved_function(&self.id, self.panic_id, state, args) } } @@ -610,42 +852,238 @@ impl FunctionContainer { where 'db: 'a, { - let ctx = state.ctx(); let args: Vec<_> = self.1.iter().map(|(_, x)| x).chain(args).copied().collect(); - let can_mint = matches!(ctx, crate::Context::Write | crate::Context::Full); - let can_read = matches!(ctx, crate::Context::Read | crate::Context::Full); - let panic_id = self.3; - // On capability mismatch, trigger the egglog runtime panic - // pre-registered at the `unstable-fn` build site (see - // `BackendRule::prim`). The panic writes to the egraph's - // panic side channel and triggers early stop, so `run_rules` - // surfaces the misuse as an `Err`. - let mismatch = |state: &mut crate::PureState<'a, 'db>| -> Option { - state.call_external_func(panic_id, &[]) - }; - match &self.0 { - ResolvedFunctionId::Constructor(action) => { - if can_mint { - action.lookup_or_insert(state.raw_exec_state(), &args) - } else { - mismatch(state) - } + apply_resolved_function(&self.0, self.3, state, &args) + } +} + +fn apply_resolved_function<'a, 'db>( + id: &ResolvedFunctionId, + panic_id: ExternalFunctionId, + state: &mut crate::PureState<'a, 'db>, + args: &[Value], +) -> Option +where + 'db: 'a, +{ + let ctx = state.ctx(); + let can_mint = matches!(ctx, crate::Context::Write | crate::Context::Full); + let can_read = matches!(ctx, crate::Context::Read | crate::Context::Full); + // On capability mismatch, trigger the egglog runtime panic + // pre-registered at the `unstable-fn` build site (see + // `BackendRule::prim`). The panic writes to the egraph's + // panic side channel and triggers early stop, so `run_rules` + // surfaces the misuse as an `Err`. + let mismatch = |state: &mut crate::PureState<'a, 'db>| -> Option { + state.call_external_func(panic_id, &[]) + }; + match id { + ResolvedFunctionId::Constructor(action) => { + if can_mint { + action.lookup_or_insert(state.raw_exec_state(), args) + } else { + mismatch(state) } - ResolvedFunctionId::Function(action) => { - if can_read { - action.lookup(state.raw_exec_state(), &args) - } else { - mismatch(state) - } + } + ResolvedFunctionId::Function(action) => { + if can_read { + action.lookup(state.raw_exec_state(), args) + } else { + mismatch(state) } - ResolvedFunctionId::Primitive { context_ids } => { - // Pick the runtime id whose context matches the - // application ctx. - match context_ids[ctx] { - Some(id) => state.call_external_func(id, &args), - None => mismatch(state), - } + } + ResolvedFunctionId::Primitive { context_ids } => { + // Pick the runtime id whose context matches the + // application ctx. + match context_ids[ctx] { + Some(id) => state.call_external_func(id, args), + None => mismatch(state), } } } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + struct CountingRemap { + calls: AtomicUsize, + add: usize, + } + + impl ValueRebuilder for CountingRemap { + fn rebuild_val(&self, value: Value) -> Value { + self.calls.fetch_add(1, Ordering::Relaxed); + Value::from_usize(value.index() + self.add) + } + } + + fn value(index: usize) -> Value { + Value::from_usize(index) + } + + #[test] + fn sequence_masks_rebuildable_args_across_word_boundaries() { + let plain = I64Sort.to_arcsort(); + let eq: ArcSort = Arc::new(EqSort { + name: "MaskEq".to_owned(), + }); + let mut sorts = vec![plain; 33]; + for index in [0, 30, 31, 32] { + sorts[index] = eq.clone(); + } + let args = (100..133).map(value).collect::>(); + let descriptor = value(77); + let mut encoded = Vec::new(); + encode_function_sequence(descriptor, &sorts, &args, &mut encoded); + + assert_eq!(encoded[0], descriptor); + assert_eq!(encoded[1], value(33)); + assert_eq!(encoded[2], value((1 << 0) | (1 << 30))); + assert_eq!(encoded[3], value((1 << 0) | (1 << 1))); + assert_eq!(&encoded[4..], args); + + let mut visited = Vec::new(); + FunctionContainer::visit_sequence_values(&encoded, &mut |value| visited.push(value)); + assert_eq!(visited, [args[0], args[30], args[31], args[32]]); + + let remap = CountingRemap { + calls: AtomicUsize::new(0), + add: 1_000, + }; + let mut rebuilt = Vec::new(); + assert!(FunctionContainer::rebuild_sequence( + &encoded, + &BaseValues::default(), + &remap, + &mut rebuilt, + )); + assert_eq!(remap.calls.load(Ordering::Relaxed), 4); + assert_eq!(&rebuilt[..4], &encoded[..4]); + for index in 0..args.len() { + let expected = if [0, 30, 31, 32].contains(&index) { + value(args[index].index() + 1_000) + } else { + args[index] + }; + assert_eq!(rebuilt[4 + index], expected); + } + } + + #[test] + fn sequence_zero_args_is_descriptor_and_count_only() { + let descriptor = value(77); + let mut encoded = Vec::new(); + encode_function_sequence(descriptor, &[], &[], &mut encoded); + assert_eq!(encoded, [descriptor, value(0)]); + + let parsed = FunctionSequence::parse(&encoded); + assert_eq!(parsed.descriptor(), descriptor); + assert!(parsed.args().is_empty()); + + let mut visited = Vec::new(); + FunctionContainer::visit_sequence_values(&encoded, &mut |value| visited.push(value)); + assert!(visited.is_empty()); + + let remap = CountingRemap { + calls: AtomicUsize::new(0), + add: 1, + }; + let mut rebuilt = Vec::new(); + assert!(!FunctionContainer::rebuild_sequence( + &encoded, + &BaseValues::default(), + &remap, + &mut rebuilt, + )); + assert_eq!(remap.calls.load(Ordering::Relaxed), 0); + assert!(rebuilt.is_empty()); + } + + #[test] + fn sequence_rebuild_does_not_resolve_descriptor() { + // This descriptor is deliberately absent from the empty BaseValues. + // Raw sequence rebuild must use only the packed mask. + let encoded = [value(900), value(2), value(1), value(10), value(20)]; + let remap = CountingRemap { + calls: AtomicUsize::new(0), + add: 5, + }; + let mut rebuilt = Vec::new(); + assert!(FunctionContainer::rebuild_sequence( + &encoded, + &BaseValues::default(), + &remap, + &mut rebuilt, + )); + assert_eq!(remap.calls.load(Ordering::Relaxed), 1); + assert_eq!( + rebuilt, + [value(900), value(2), value(1), value(15), value(20)] + ); + + let unchanged = CountingRemap { + calls: AtomicUsize::new(0), + add: 0, + }; + rebuilt.clear(); + assert!(!FunctionContainer::rebuild_sequence( + &encoded, + &BaseValues::default(), + &unchanged, + &mut rebuilt, + )); + assert_eq!(unchanged.calls.load(Ordering::Relaxed), 1); + assert!(rebuilt.is_empty()); + } + + #[test] + fn sequence_codec_round_trips_descriptor_and_mixed_args() { + let mut base_values = BaseValues::default(); + base_values.register_type::(); + let eq: ArcSort = Arc::new(EqSort { + name: "RoundTripEq".to_owned(), + }); + let function = FunctionContainer( + ResolvedFunctionId::Primitive { + context_ids: EnumMap::default(), + }, + vec![(I64Sort.to_arcsort(), value(10)), (eq, value(20))], + "round-trip".to_owned(), + ExternalFunctionId::new(7), + ); + let mut encoded = Vec::new(); + function.encode_sequence(&base_values, &mut encoded); + assert_eq!(encoded[1], value(2)); + assert_eq!(encoded[2], value(2)); + assert_eq!( + FunctionContainer::decode_sequence(&encoded, &base_values), + function + ); + + let mut visited = Vec::new(); + FunctionContainer::visit_sequence_values(&encoded, &mut |value| visited.push(value)); + assert_eq!(visited, [value(20)]); + } + + #[test] + #[should_panic(expected = "reserved high bit")] + fn sequence_parser_rejects_mask_high_bit() { + FunctionSequence::parse(&[value(1), value(1), Value::new(1 << 31), value(2)]); + } + + #[test] + #[should_panic(expected = "bits beyond its argument count")] + fn sequence_parser_rejects_noncanonical_unused_bits() { + FunctionSequence::parse(&[value(1), value(1), value(2), value(2)]); + } + + #[test] + #[should_panic(expected = "length does not match its argument count")] + fn sequence_parser_rejects_trailing_values() { + FunctionSequence::parse(&[value(1), value(0), value(2)]); + } +} diff --git a/src/sort/map.rs b/src/sort/map.rs index 25a5659d3..8690d9f95 100644 --- a/src/sort/map.rs +++ b/src/sort/map.rs @@ -41,12 +41,12 @@ impl ContainerValue for MapContainer { } impl SequenceContainerValue for MapContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &BaseValues, out: &mut Vec) { out.push(map_rebuild_mask(self.do_rebuild_keys, self.do_rebuild_vals)); out.extend(self.data.iter().flat_map(|(key, value)| [*key, *value])); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &BaseValues) -> Self { let (&mask, data) = sequence .split_first() .expect("serialized MapContainer must include its rebuild mask"); @@ -98,6 +98,7 @@ impl SequenceContainerValue for MapContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -854,8 +855,11 @@ mod tests { ] { let map = map(rebuild_keys, rebuild_vals, &[(2, 20), (4, 40)]); let mut encoded = Vec::new(); - map.encode_sequence(&mut encoded); - assert_eq!(MapContainer::decode_sequence(&encoded), map); + map.encode_sequence(&BaseValues::default(), &mut encoded); + assert_eq!( + MapContainer::decode_sequence(&encoded, &BaseValues::default()), + map + ); assert_eq!( MapContainer::sequence_values(&encoded), &[value(2), value(20), value(4), value(40)] @@ -883,44 +887,51 @@ mod tests { assert_eq!(legacy, map(true, true, &[(1, 41), (3, 61)])); let mut encoded = Vec::new(); - original.encode_sequence(&mut encoded); + original.encode_sequence(&BaseValues::default(), &mut encoded); let mut rebuilt = Vec::new(); assert!(MapContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &rebuilder, &mut rebuilt )); - assert_eq!(MapContainer::decode_sequence(&rebuilt), legacy); + assert_eq!( + MapContainer::decode_sequence(&rebuilt, &BaseValues::default()), + legacy + ); } #[test] fn sequence_rebuild_respects_mask_and_false_leaves_output_empty() { let original = map(false, true, &[(2, 20), (4, 40)]); let mut encoded = Vec::new(); - original.encode_sequence(&mut encoded); + original.encode_sequence(&BaseValues::default(), &mut encoded); let mut rebuilt = Vec::new(); assert!(MapContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &Remap(vec![(value(2), value(3)), (value(20), value(21))]), &mut rebuilt, )); assert_eq!( - MapContainer::decode_sequence(&rebuilt), + MapContainer::decode_sequence(&rebuilt, &BaseValues::default()), map(false, true, &[(2, 21), (4, 40)]) ); let mut unchanged = Vec::new(); assert!(!MapContainer::rebuild_sequence( &rebuilt, + &BaseValues::default(), &Remap(vec![(value(2), value(3))]), &mut unchanged, )); assert!(unchanged.is_empty()); let mut non_rebuildable = Vec::new(); - map(false, false, &[(2, 20)]).encode_sequence(&mut non_rebuildable); + map(false, false, &[(2, 20)]).encode_sequence(&BaseValues::default(), &mut non_rebuildable); assert!(!MapContainer::rebuild_sequence( &non_rebuildable, + &BaseValues::default(), &Remap(vec![(value(2), value(3)), (value(20), value(21))]), &mut unchanged, )); diff --git a/src/sort/multiset.rs b/src/sort/multiset.rs index 0ba42bc20..fd68d50e0 100644 --- a/src/sort/multiset.rs +++ b/src/sort/multiset.rs @@ -106,7 +106,7 @@ impl ContainerValue for MultiSetContainer { } impl SequenceContainerValue for MultiSetContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &BaseValues, out: &mut Vec) { assert!( i64::try_from(self.data.len()).is_ok(), "multiset cardinality exceeds the i64 language representation" @@ -118,7 +118,7 @@ impl SequenceContainerValue for MultiSetContainer { } } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &BaseValues) -> Self { let (&header, data) = sequence .split_first() .expect("serialized MultiSetContainer must include its rebuild flag"); @@ -168,6 +168,7 @@ impl SequenceContainerValue for MultiSetContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -198,7 +199,7 @@ impl SequenceContainerValue for MultiSetContainer { do_rebuild: true, data: rebuilt, } - .encode_sequence(out); + .encode_sequence(_base_values, out); if out.as_slice() == sequence { out.clear(); return false; @@ -1078,13 +1079,13 @@ impl PurePrim for Map { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[0])?; + let function = state.prepare_function(args[0])?; // Copy before callbacks: they may intern containers and grow the // execution-local prediction storage backing the fast slice. let entries = multiset_entries(&state, args[1])?; let mut mapped_values = MultiSet::new(); for (v, count) in entries { - if let Some(mapped) = state.apply_function(&fc, &[v]) { + if let Some(mapped) = state.apply_prepared_function(&function, &[v]) { mapped_values.checked_insert_multiple_mut(mapped, count)?; } } @@ -1133,12 +1134,12 @@ impl FullPrim for FillIndex { mut state: crate::FullState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[1])?; + let resolved = state.resolve_function_value(args[1])?; let entries = multiset_entries(&state, args[0])? .into_iter() .map(|(value, count)| Some((value, i64::try_from(count).ok()?))) .collect::>>()?; - let action = match fc.0 { + let action = match resolved.id { ResolvedFunctionId::Constructor(a) | ResolvedFunctionId::Function(a) => a, // Primitive functions cannot be used with // unstable-multiset-fill-index, since they cannot be set. @@ -1193,9 +1194,9 @@ impl WritePrim for ClearIndex { mut state: crate::WriteState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[1])?; + let resolved = state.resolve_function_value(args[1])?; let entries = multiset_entries(&state, args[0])?; - let action = match fc.0 { + let action = match resolved.id { ResolvedFunctionId::Constructor(a) | ResolvedFunctionId::Function(a) => a, // Primitive functions cannot be used with // unstable-multiset-clear-index, since they cannot be deleted. @@ -1244,11 +1245,11 @@ impl PurePrim for FlatMap { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[0])?; + let function = state.prepare_function(args[0])?; let entries = multiset_entries(&state, args[1])?; let mut flattened = MultiSet::new(); for (v, count) in entries { - let mapped = state.apply_function(&fc, &[v]); + let mapped = state.apply_prepared_function(&function, &[v]); if let Some(mapped_ms) = mapped { for (mapped, mapped_count) in multiset_entries(&state, mapped_ms)? { flattened @@ -1301,11 +1302,11 @@ impl PurePrim for Filter { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[0])?; + let function = state.prepare_function(args[0])?; let entries = multiset_entries(&state, args[1])?; let mut filtered = Vec::with_capacity(entries.len()); for (v, count) in entries { - let mapped = state.apply_function(&fc, &[v]); + let mapped = state.apply_prepared_function(&function, &[v]); if mapped.is_some() == self.skip_empty { filtered.push((v, count)); } @@ -1400,7 +1401,7 @@ impl PurePrim for Reduce { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[0])?; + let function = state.prepare_function(args[0])?; let initial = args[1]; let entries = multiset_entries(&state, args[2])?; let mut acc = initial; @@ -1408,7 +1409,7 @@ impl PurePrim for Reduce { for (value, count) in entries { for _ in 0..count { if has_value { - acc = state.apply_function(&fc, &[acc, value])?; + acc = state.apply_prepared_function(&function, &[acc, value])?; } else { acc = value; has_value = true; @@ -1638,9 +1639,12 @@ mod tests { data, }; let mut encoded = Vec::new(); - multiset.encode_sequence(&mut encoded); + multiset.encode_sequence(&BaseValues::default(), &mut encoded); assert_eq!(encoded.len(), 1 + 2 * MULTISET_ENTRY_WIDTH); - assert_eq!(MultiSetContainer::decode_sequence(&encoded), multiset); + assert_eq!( + MultiSetContainer::decode_sequence(&encoded, &BaseValues::default()), + multiset + ); let mut children = Vec::new(); MultiSetContainer::visit_sequence_values(&encoded, &mut |value| children.push(value)); @@ -1649,13 +1653,14 @@ mod tests { let mut rebuilt = Vec::new(); assert!(MultiSetContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &Collapse { from: value(4), to: value(2), }, &mut rebuilt, )); - let rebuilt = MultiSetContainer::decode_sequence(&rebuilt); + let rebuilt = MultiSetContainer::decode_sequence(&rebuilt, &BaseValues::default()); assert_eq!( rebuilt.data.iter_counts().collect::>(), vec![(value(2), 8)] @@ -1673,9 +1678,12 @@ mod tests { data, }; let mut encoded = Vec::new(); - multiset.encode_sequence(&mut encoded); + multiset.encode_sequence(&BaseValues::default(), &mut encoded); assert_ne!(encoded[3], value(0)); - assert_eq!(MultiSetContainer::decode_sequence(&encoded), multiset); + assert_eq!( + MultiSetContainer::decode_sequence(&encoded, &BaseValues::default()), + multiset + ); } #[cfg(target_pointer_width = "64")] @@ -1686,7 +1694,7 @@ mod tests { encode_count(&mut encoded, i64::MAX as usize); encoded.push(value(2)); encode_count(&mut encoded, 1); - MultiSetContainer::decode_sequence(&encoded); + MultiSetContainer::decode_sequence(&encoded, &BaseValues::default()); } #[test] diff --git a/src/sort/pair.rs b/src/sort/pair.rs index 12c12f99a..b2a92b1f2 100644 --- a/src/sort/pair.rs +++ b/src/sort/pair.rs @@ -31,13 +31,13 @@ impl ContainerValue for PairContainer { } impl SequenceContainerValue for PairContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &BaseValues, out: &mut Vec) { let rebuild_mask = self.do_rebuild_first as usize | ((self.do_rebuild_second as usize) << 1); out.extend_from_slice(&[Value::from_usize(rebuild_mask), self.first, self.second]); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &BaseValues) -> Self { let [rebuild_mask, first, second] = sequence else { panic!("serialized PairContainer must contain a rebuild mask and two values"); }; @@ -83,6 +83,7 @@ impl SequenceContainerValue for PairContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -416,9 +417,12 @@ mod tests { second: value(20), }; let mut encoded = Vec::new(); - pair.encode_sequence(&mut encoded); + pair.encode_sequence(&BaseValues::default(), &mut encoded); assert_eq!(encoded, vec![value(2), value(10), value(20)]); - assert_eq!(PairContainer::decode_sequence(&encoded), pair); + assert_eq!( + PairContainer::decode_sequence(&encoded, &BaseValues::default()), + pair + ); assert_eq!(PairContainer::sequence_values(&encoded), &encoded[1..]); let mut visited = Vec::new(); @@ -433,6 +437,7 @@ mod tests { let mut rebuilt = Vec::new(); assert!(PairContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &rebuilder, &mut rebuilt, )); @@ -451,6 +456,7 @@ mod tests { let mut rebuilt = Vec::new(); assert!(!PairContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &rebuilder, &mut rebuilt, )); diff --git a/src/sort/set.rs b/src/sort/set.rs index 567e669a3..77d617795 100644 --- a/src/sort/set.rs +++ b/src/sort/set.rs @@ -25,12 +25,12 @@ impl ContainerValue for SetContainer { } impl SequenceContainerValue for SetContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &BaseValues, out: &mut Vec) { out.push(Value::from_usize(self.do_rebuild as usize)); out.extend(self.data.iter().copied()); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &BaseValues) -> Self { let (&header, data) = sequence .split_first() .expect("serialized SetContainer must include its rebuild flag"); @@ -65,6 +65,7 @@ impl SequenceContainerValue for SetContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -683,12 +684,16 @@ mod tests { data: [value(2), value(4), value(6)].into_iter().collect(), }; let mut encoded = Vec::new(); - set.encode_sequence(&mut encoded); - assert_eq!(SetContainer::decode_sequence(&encoded), set); + set.encode_sequence(&BaseValues::default(), &mut encoded); + assert_eq!( + SetContainer::decode_sequence(&encoded, &BaseValues::default()), + set + ); let mut rebuilt = Vec::new(); assert!(SetContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &Collapse { from: value(6), to: value(2), @@ -708,6 +713,7 @@ mod tests { let mut rebuilt = Vec::new(); assert!(!SetContainer::rebuild_sequence( &encoded, + &BaseValues::default(), &Collapse { from: value(2), to: value(4), diff --git a/src/sort/vec.rs b/src/sort/vec.rs index 6b5e6243e..f454fbd44 100644 --- a/src/sort/vec.rs +++ b/src/sort/vec.rs @@ -26,12 +26,12 @@ impl ContainerValue for VecContainer { } impl SequenceContainerValue for VecContainer { - fn encode_sequence(&self, out: &mut Vec) { + fn encode_sequence(&self, _base_values: &BaseValues, out: &mut Vec) { out.push(Value::from_usize(self.do_rebuild as usize)); out.extend_from_slice(&self.data); } - fn decode_sequence(sequence: &[Value]) -> Self { + fn decode_sequence(sequence: &[Value], _base_values: &BaseValues) -> Self { let (&header, data) = sequence .split_first() .expect("serialized VecContainer must include its rebuild flag"); @@ -64,6 +64,7 @@ impl SequenceContainerValue for VecContainer { fn rebuild_sequence( sequence: &[Value], + _base_values: &BaseValues, rebuilder: &dyn ValueRebuilder, out: &mut Vec, ) -> bool { @@ -663,7 +664,7 @@ impl PurePrim for VecMap { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state.value_to_owned_container::(args[0])?; + let function = state.prepare_function(args[0])?; // Copy before invoking the callback: it may intern another Vec and // grow this execution's local prediction storage. let input = state @@ -675,7 +676,7 @@ impl PurePrim for VecMap { })?; let mut new_data = Vec::with_capacity(input.len()); for v in input { - if let Some(mapped) = state.apply_function(&fc, &[v]) { + if let Some(mapped) = state.apply_prepared_function(&function, &[v]) { new_data.push(mapped); } } diff --git a/tests/container_rebuild.rs b/tests/container_rebuild.rs index 4c3131fa6..2ab017943 100644 --- a/tests/container_rebuild.rs +++ b/tests/container_rebuild.rs @@ -428,6 +428,120 @@ fn nested_set_rebuild_term_only() { .unwrap(); } +/// A sequence-backed `UnstableFn` canonicalizes its rebuildable partial args. +/// `UnstableFn` deliberately has no proof validator, so this runs in ordinary +/// mode rather than term/proof encoding. +#[test] +fn unstable_fn_rebuild() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor A () Math) + (constructor B () Math) + (constructor Apply (Math Math) Math) + (sort MathFn (UnstableFn (Math) Math)) + (constructor Holds (MathFn) Math) + (Holds (unstable-fn "Apply" (A))) + (Holds (unstable-fn "Apply" (B))) + (union (A) (B)) + (run 1) + (check (= (Holds (unstable-fn "Apply" (A))) + (Holds (unstable-fn "Apply" (B))))) + "#, + ) + .unwrap(); +} + +/// Container-valued partial arguments participate in the function sequence's +/// rebuild mask, so an inner Vec identity change canonicalizes the outer +/// function key without decoding its descriptor. +#[test] +fn unstable_fn_container_partial_rebuild() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor A () Math) + (constructor B () Math) + (constructor C () Math) + (sort MathVec (Vec Math)) + (constructor ApplyVec (MathVec Math) Math) + (sort MathFn (UnstableFn (Math) Math)) + (constructor Holds (MathFn) Math) + (let $fa (unstable-fn "ApplyVec" (vec-of (A)))) + (let $fb (unstable-fn "ApplyVec" (vec-of (B)))) + (Holds $fa) + (Holds $fb) + (union (A) (B)) + (run 1) + (check (= (Holds $fa) (Holds $fb))) + (let $result (unstable-app $fa (C))) + (let $expected (ApplyVec (vec-of (A)) (C))) + (run 1) + (check (= $result $expected)) + "#, + ) + .unwrap(); +} + +/// Rebuilding a captured Eq argument propagates through an outer Vec even when +/// the function's remaining input is a base sort. Function-sort rebuildability +/// cannot be inferred from its remaining inputs. +#[test] +fn nested_unstable_fn_dirty_propagates_to_outer_vec() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor A () Math) + (constructor B () Math) + (constructor Capture (Math i64) Math) + (sort IntFn (UnstableFn (i64) Math)) + (sort IntFnVec (Vec IntFn)) + (constructor Holds (IntFnVec) Math) + (Holds (vec-of (unstable-fn "Capture" (A)))) + (Holds (vec-of (unstable-fn "Capture" (B)))) + (union (A) (B)) + (run 1) + (check (= (Holds (vec-of (unstable-fn "Capture" (A)))) + (Holds (vec-of (unstable-fn "Capture" (B)))))) + "#, + ) + .unwrap(); +} + +/// A stable function identity remains readable after its sequence key is +/// rebuilt and the identity-to-row reverse map moves to the replacement row. +#[test] +fn unstable_fn_application_survives_rebuild_and_reverse_refresh() { + let mut egraph = EGraph::default(); + egraph + .parse_and_run_program( + None, + r#" + (sort Math) + (constructor A () Math) + (constructor B () Math) + (constructor C () Math) + (constructor Apply (Math Math) Math) + (sort MathFn (UnstableFn (Math) Math)) + (let $f (unstable-fn "Apply" (B))) + (union (A) (B)) + (run 1) + (let $result (unstable-app $f (C))) + (check (= $result (Apply (A) (C)))) + "#, + ) + .unwrap(); +} + /// Compact MultiSet count lanes are metadata, not child IDs. Nested rebuild /// propagation must follow only the semantic values in each encoded entry. #[test]