diff --git a/src/sort/vec.rs b/src/sort/vec.rs index 22076fdd4..6b5e6243e 100644 --- a/src/sort/vec.rs +++ b/src/sort/vec.rs @@ -1,4 +1,6 @@ use crate::Write; +use crate::constraint::AllEqualTypeConstraint; +use crate::numeric_id::NumericId; use std::any::TypeId; use std::iter::zip; @@ -23,6 +25,70 @@ impl ContainerValue for VecContainer { } } +impl SequenceContainerValue for VecContainer { + fn encode_sequence(&self, 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 { + let (&header, data) = sequence + .split_first() + .expect("serialized VecContainer must include its rebuild flag"); + assert!( + header == Value::from_usize(0) || header == Value::from_usize(1), + "serialized VecContainer has an invalid rebuild flag" + ); + Self { + do_rebuild: header == Value::from_usize(1), + data: data.to_vec(), + } + } + + fn sequence_values(sequence: &[Value]) -> &[Value] { + sequence + .get(1..) + .expect("serialized VecContainer must include its rebuild flag") + } + + fn visit_sequence_values(sequence: &[Value], visitor: &mut dyn FnMut(Value)) { + let (&header, data) = sequence + .split_first() + .expect("serialized VecContainer must include its rebuild flag"); + match header.index() { + 0 => {} + 1 => data.iter().copied().for_each(visitor), + _ => panic!("serialized VecContainer has an invalid rebuild flag"), + } + } + + fn rebuild_sequence( + sequence: &[Value], + rebuilder: &dyn ValueRebuilder, + out: &mut Vec, + ) -> bool { + let (&header, data) = sequence + .split_first() + .expect("serialized VecContainer must include its rebuild flag"); + if header == Value::from_usize(0) { + return false; + } + assert_eq!( + header, + Value::from_usize(1), + "serialized VecContainer has an invalid rebuild flag" + ); + out.push(header); + out.extend_from_slice(data); + if rebuilder.rebuild_slice(&mut out[1..]) { + true + } else { + out.clear(); + false + } + } +} + #[derive(Clone, Debug)] pub struct VecSort { name: String, @@ -111,6 +177,10 @@ impl ContainerSort for VecSort { &self.name } + fn register_type(&self, backend: &mut egglog_bridge::EGraph) { + backend.register_sequence_container_ty::(); + } + fn inner_sorts(&self) -> Vec { vec![self.element.clone()] } @@ -180,21 +250,67 @@ impl ContainerSort for VecSort { do_rebuild: self.ctx.is_eq_container_sort(), data: xs .collect() } }, vec_of_validator); - add_primitive!(eg, "vec-append" = {self.clone(): VecSort} [xs: @VecContainer (arc)] -> @VecContainer (arc) { VecContainer { - do_rebuild: self.ctx.is_eq_container_sort(), - data: xs.flat_map(|x| x.data).collect() - } }); - - add_primitive!(eg, "vec-push" = |mut xs: @VecContainer (arc), x: # (self.element())| -> @VecContainer (arc) {{ xs.data.push(x); xs }}); - add_primitive!(eg, "vec-pop" = |mut xs: @VecContainer (arc) | -> @VecContainer (arc) {{ xs.data.pop(); xs }}); - - add_primitive_with_validator!(eg, "vec-length" = |xs: @VecContainer (arc)| -> i64 { xs.data.len() as i64 }, vec_length_validator); - add_primitive_with_validator!(eg, "vec-contains" = |xs: @VecContainer (arc), x: # (self.element())| -?> () { ( xs.data.contains(&x)).then_some(()) }, vec_contains_validator); - add_primitive_with_validator!(eg, "vec-not-contains" = |xs: @VecContainer (arc), x: # (self.element())| -?> () { (!xs.data.contains(&x)).then_some(()) }, vec_not_contains_validator); + eg.add_pure_primitive( + VecAppend { + name: "vec-append".into(), + vec: arc.clone(), + }, + None, + ); + + for (name, op) in [ + ("vec-push", VecEditOp::Push), + ("vec-pop", VecEditOp::Pop), + ("vec-set", VecEditOp::Set), + ("vec-remove", VecEditOp::Remove), + ] { + eg.add_pure_primitive( + VecEdit { + name: name.into(), + vec: arc.clone(), + element: self.element(), + op, + }, + None, + ); + } - add_primitive_with_validator!(eg, "vec-get" = | xs: @VecContainer (arc), i: i64 | -?> # (self.element()) { xs.data.get(i as usize).copied() }, vec_get_validator); - add_primitive!(eg, "vec-set" = |mut xs: @VecContainer (arc), i: i64, x: # (self.element())| -?> @VecContainer (arc) {{ let idx = usize::try_from(i).ok()?; if idx >= xs.data.len() { None } else { xs.data[idx] = x; Some(xs) } }}); - add_primitive!(eg, "vec-remove" = |mut xs: @VecContainer (arc), i: i64 | -?> @VecContainer (arc) {{ let idx = usize::try_from(i).ok()?; if idx >= xs.data.len() { None } else { xs.data.remove(idx); Some(xs) } }}); + eg.add_pure_primitive( + VecRead { + name: "vec-length".into(), + vec: arc.clone(), + element: self.element(), + op: VecReadOp::Length, + }, + Some(Arc::new(vec_length_validator)), + ); + eg.add_pure_primitive( + VecRead { + name: "vec-contains".into(), + vec: arc.clone(), + element: self.element(), + op: VecReadOp::Contains, + }, + Some(Arc::new(vec_contains_validator)), + ); + eg.add_pure_primitive( + VecRead { + name: "vec-not-contains".into(), + vec: arc.clone(), + element: self.element(), + op: VecReadOp::NotContains, + }, + Some(Arc::new(vec_not_contains_validator)), + ); + eg.add_pure_primitive( + VecRead { + name: "vec-get".into(), + vec: arc.clone(), + element: self.element(), + op: VecReadOp::Get, + }, + Some(Arc::new(vec_get_validator)), + ); if self.element.is_eq_sort() { eg.add_write_primitive( Union { @@ -253,6 +369,233 @@ impl ContainerSort for VecSort { } } +#[derive(Clone, Copy)] +enum VecReadOp { + Length, + Contains, + NotContains, + Get, +} + +/// Vec reads with a sequence-slice fast path and a reconstruction fallback. +#[derive(Clone)] +struct VecRead { + name: String, + vec: ArcSort, + element: ArcSort, + op: VecReadOp, +} + +impl Primitive for VecRead { + fn name(&self) -> &str { + &self.name + } + + fn get_type_constraints(&self, span: &Span) -> Box { + let types = match self.op { + VecReadOp::Length => vec![self.vec.clone(), I64Sort.to_arcsort()], + VecReadOp::Contains | VecReadOp::NotContains => { + vec![ + self.vec.clone(), + self.element.clone(), + UnitSort.to_arcsort(), + ] + } + VecReadOp::Get => vec![self.vec.clone(), I64Sort.to_arcsort(), self.element.clone()], + }; + SimpleTypeConstraint::new(self.name(), types, span.clone()).into_box() + } +} + +impl PurePrim for VecRead { + fn apply<'a, 'db>(&self, state: crate::PureState<'a, 'db>, args: &[Value]) -> Option { + let [vec_id, rest @ ..] = args else { + return None; + }; + match self.op { + VecReadOp::Length => { + if !rest.is_empty() { + return None; + } + let len = state + .with_container_sequence::(*vec_id, |values| values.len()) + .or_else(|| { + state + .value_to_owned_container::(*vec_id) + .map(|vec| vec.data.len()) + })?; + Some(state.base_values().get::(len as i64)) + } + VecReadOp::Contains | VecReadOp::NotContains => { + let [needle] = rest else { return None }; + let contains = state + .with_container_sequence::(*vec_id, |values| { + values.contains(needle) + }) + .or_else(|| { + state + .value_to_owned_container::(*vec_id) + .map(|vec| vec.data.contains(needle)) + })?; + let succeeds = match self.op { + VecReadOp::Contains => contains, + VecReadOp::NotContains => !contains, + _ => unreachable!(), + }; + succeeds.then(|| state.base_values().get::<()>(())) + } + VecReadOp::Get => { + let [index] = rest else { return None }; + let index = usize::try_from(state.base_values().unwrap::(*index)).ok()?; + state + .with_container_sequence::(*vec_id, |values| { + values.get(index).copied() + }) + .or_else(|| { + state + .value_to_owned_container::(*vec_id) + .map(|vec| vec.data.get(index).copied()) + })? + } + } + } +} + +#[derive(Clone, Copy)] +enum VecEditOp { + Push, + Pop, + Set, + Remove, +} + +/// Vec updates with a serialized-slice fast path and a legacy slow fallback. +#[derive(Clone)] +struct VecEdit { + name: String, + vec: ArcSort, + element: ArcSort, + op: VecEditOp, +} + +impl Primitive for VecEdit { + fn name(&self) -> &str { + &self.name + } + + fn get_type_constraints(&self, span: &Span) -> Box { + let types = match self.op { + VecEditOp::Push => vec![self.vec.clone(), self.element.clone(), self.vec.clone()], + VecEditOp::Pop => vec![self.vec.clone(), self.vec.clone()], + VecEditOp::Set => vec![ + self.vec.clone(), + I64Sort.to_arcsort(), + self.element.clone(), + self.vec.clone(), + ], + VecEditOp::Remove => vec![self.vec.clone(), I64Sort.to_arcsort(), self.vec.clone()], + }; + SimpleTypeConstraint::new(self.name(), types, span.clone()).into_box() + } +} + +impl PurePrim for VecEdit { + fn apply<'a, 'db>( + &self, + mut state: crate::PureState<'a, 'db>, + args: &[Value], + ) -> Option { + let vec_id = *args.first()?; + let build_key = |data: &[Value]| -> Option> { + let mut key = Vec::with_capacity(data.len() + 2); + key.push(Value::from_usize(self.vec.is_eq_container_sort() as usize)); + match self.op { + VecEditOp::Push => { + let [_, value] = args else { return None }; + key.extend_from_slice(data); + key.push(*value); + } + VecEditOp::Pop => { + let [_] = args else { return None }; + key.extend_from_slice(data.get(..data.len().saturating_sub(1))?); + } + VecEditOp::Set => { + let [_, index, value] = args else { + return None; + }; + let index = usize::try_from(state.base_values().unwrap::(*index)).ok()?; + if index >= data.len() { + return None; + } + key.extend_from_slice(data); + key[index + 1] = *value; + } + VecEditOp::Remove => { + let [_, index] = args else { return None }; + let index = usize::try_from(state.base_values().unwrap::(*index)).ok()?; + if index >= data.len() { + return None; + } + key.extend_from_slice(&data[..index]); + key.extend_from_slice(&data[index + 1..]); + } + } + Some(key) + }; + + let key = state + .with_container_sequence::(vec_id, build_key) + .or_else(|| { + state + .value_to_owned_container::(vec_id) + .map(|vec| build_key(&vec.data)) + })??; + Some(state.register_container_sequence::(&key)) + } +} + +#[derive(Clone)] +struct VecAppend { + name: String, + vec: ArcSort, +} + +impl Primitive for VecAppend { + fn name(&self) -> &str { + &self.name + } + + fn get_type_constraints(&self, span: &Span) -> Box { + AllEqualTypeConstraint::new(self.name(), span.clone()) + .with_all_arguments_sort(self.vec.clone()) + .with_output_sort(self.vec.clone()) + .into_box() + } +} + +impl PurePrim for VecAppend { + fn apply<'a, 'db>( + &self, + mut state: crate::PureState<'a, 'db>, + args: &[Value], + ) -> Option { + let mut key = vec![Value::from_usize(self.vec.is_eq_container_sort() as usize)]; + for value in args { + if state + .with_container_sequence::(*value, |values| { + key.extend_from_slice(values); + }) + .is_none() + { + key.extend_from_slice( + &state.value_to_owned_container::(*value)?.data, + ); + } + } + Some(state.register_container_sequence::(&key)) + } +} + /** * Register a vec map primitive if the function matches the input and output vec. */ @@ -320,27 +663,28 @@ impl PurePrim for VecMap { mut state: crate::PureState<'a, 'db>, args: &[Value], ) -> Option { - let fc = state - .container_values() - .get_val::(args[0]) - .unwrap() - .clone(); - let vec = state - .container_values() - .get_val::(args[1]) - .unwrap() - .clone(); - let mut new_data = Vec::with_capacity(vec.data.len()); - for v in vec.data { + let fc = state.value_to_owned_container::(args[0])?; + // Copy before invoking the callback: it may intern another Vec and + // grow this execution's local prediction storage. + let input = state + .with_container_sequence::(args[1], <[Value]>::to_vec) + .or_else(|| { + state + .value_to_owned_container::(args[1]) + .map(|vec| vec.data) + })?; + let mut new_data = Vec::with_capacity(input.len()); + for v in input { if let Some(mapped) = state.apply_function(&fc, &[v]) { new_data.push(mapped); } } - let new_vec = VecContainer { - do_rebuild: self.output_vec.is_eq_container_sort(), - data: new_data, - }; - Some(state.register_container(new_vec)) + let mut key = Vec::with_capacity(new_data.len() + 1); + key.push(Value::from_usize( + self.output_vec.is_eq_container_sort() as usize + )); + key.extend_from_slice(&new_data); + Some(state.register_container_sequence::(&key)) } } @@ -378,16 +722,22 @@ impl WritePrim for Union { mut state: crate::WriteState<'a, 'db>, args: &[Value], ) -> Option { + // The union calls below mutate state, so materialize the two borrowed + // slices only after taking the allocation-free fast lookup path. let left = state - .container_values() - .get_val::(args[0])? - .clone() - .data; + .with_container_sequence::(args[0], <[Value]>::to_vec) + .or_else(|| { + state + .value_to_owned_container::(args[0]) + .map(|vec| vec.data) + })?; let right = state - .container_values() - .get_val::(args[1])? - .clone() - .data; + .with_container_sequence::(args[1], <[Value]>::to_vec) + .or_else(|| { + state + .value_to_owned_container::(args[1]) + .map(|vec| vec.data) + })?; if left.len() != right.len() { return None; } diff --git a/tests/container-before-table-rebuild.egg b/tests/container-before-table-rebuild.egg new file mode 100644 index 000000000..13d90aa8e --- /dev/null +++ b/tests/container-before-table-rebuild.egg @@ -0,0 +1,17 @@ +;; A sequence-backed container table must rebuild before an ordinary table +;; whose keys and values both refer to the affected e-classes. Otherwise the +;; two `image` rows collide while their Vec values are still distinct, and the +;; `:no-merge` function reports an illegal merge. +(sort E) +(constructor a () E) +(constructor b () E) + +(sort VE (Vec E)) +(function image (E) VE :no-merge) + +(set (image (a)) (vec-of (a))) +(set (image (b)) (vec-of (b))) + +(union (a) (b)) + +(check (= (image (a)) (image (b)))) diff --git a/tests/container-dirty-after-table-rebuild.egg b/tests/container-dirty-after-table-rebuild.egg new file mode 100644 index 000000000..68a4603f6 --- /dev/null +++ b/tests/container-dirty-after-table-rebuild.egg @@ -0,0 +1,17 @@ +;; Rebuilding this relation changes both its ordinary e-class key and the +;; contents of its stable Vec identity. The rebuilt replacement must be +;; published before dirty-container refresh scans for rows containing `$v`; +;; otherwise refresh sees and re-inserts the obsolete `(b, $v)` row. +(sort E) +(sort VE (Vec E)) + +(constructor a () E) +(constructor b () E) +(relation R (E VE)) + +(let $v (vec-of (b))) +(R (b) $v) + +(union (a) (b)) + +(check (R (a) (vec-of (a)))) diff --git a/tests/snapshots/files__proof_unsupported_files.snap b/tests/snapshots/files__proof_unsupported_files.snap index 13d42f9ce..dfe3b6378 100644 --- a/tests/snapshots/files__proof_unsupported_files.snap +++ b/tests/snapshots/files__proof_unsupported_files.snap @@ -8,6 +8,8 @@ before-proofs.egg bignum.egg bool.egg complex-merge-prim.egg +container-before-table-rebuild.egg +container-dirty-after-table-rebuild.egg container-fail.egg container-rebuild.egg cyk.egg diff --git a/tests/snapshots/files__shared_snapshot_container_before_table_rebuild.snap b/tests/snapshots/files__shared_snapshot_container_before_table_rebuild.snap new file mode 100644 index 000000000..6fa99c302 --- /dev/null +++ b/tests/snapshots/files__shared_snapshot_container_before_table_rebuild.snap @@ -0,0 +1,7 @@ +--- +source: tests/files.rs +expression: snapshot_content_across_treatments +--- +((a 1) + (b 1) + (image 1)) diff --git a/tests/snapshots/files__shared_snapshot_container_dirty_after_table_rebuild.snap b/tests/snapshots/files__shared_snapshot_container_dirty_after_table_rebuild.snap new file mode 100644 index 000000000..910194a0f --- /dev/null +++ b/tests/snapshots/files__shared_snapshot_container_dirty_after_table_rebuild.snap @@ -0,0 +1,7 @@ +--- +source: tests/files.rs +expression: snapshot_content_across_treatments +--- +((R 1) + (a 1) + (b 1))