conjure_cp_core/ast/
expression_arena.rs

1use std::collections::VecDeque;
2use uniplate::Uniplate;
3
4use super::{Expression, discriminant_from_value};
5
6/// Stable handle for an expression node stored in an [`ExpressionArena`].
7#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
8pub struct ExpressionNodeId(usize);
9
10impl ExpressionNodeId {
11    /// Returns the arena slot index for this node id.
12    pub fn index(self) -> usize {
13        self.0
14    }
15}
16
17/// Arena-backed representation of an [`Expression`] tree.
18///
19/// The arena keeps direct parent/child links so callers can jump to known nodes without walking
20/// down from the root each time. Subtree replacement preserves the replaced node's id and appends
21/// the replacement descendants; descendants of the old subtree become unreachable from the root.
22#[derive(Clone, Debug)]
23pub struct ExpressionArena {
24    nodes: Vec<ExpressionArenaNode>,
25    root: ExpressionNodeId,
26}
27
28#[derive(Clone, Debug)]
29struct ExpressionArenaNode {
30    /// The expression payload, replaced with an empty tombstone when this node becomes
31    /// unreachable.
32    expr: Expression,
33    parent: Option<ExpressionNodeId>,
34    children: Vec<ExpressionNodeId>,
35    /// Child generations already reflected in the stored payload.
36    payload_child_generations: Vec<u32>,
37    /// Counts of direct child expression variants.
38    ///
39    /// The rule prefilter asks this question extremely often. Keeping the small set of distinct
40    /// variants here avoids rescanning every element of wide matrix literals after each rewrite.
41    direct_child_discriminants: Vec<(usize, usize)>,
42    reachable: bool,
43    /// Position in the parent's expression-child sequence. `None` only for the root.
44    child_slot: Option<usize>,
45    /// Distance from the root, cached because the scheduler reads it frequently.
46    depth: usize,
47    /// Incremented when this node's rewrite-relevant content changes.
48    generation: u32,
49}
50
51impl ExpressionArena {
52    /// Builds an arena from an expression tree.
53    pub fn from_root(root: Expression) -> Self {
54        let mut arena = Self {
55            nodes: Vec::new(),
56            root: ExpressionNodeId(0),
57        };
58        arena.root = arena.push_subtree(root, None, None, 0);
59        arena
60    }
61
62    /// Returns the root node id.
63    pub fn root(&self) -> ExpressionNodeId {
64        self.root
65    }
66
67    /// Returns the expression payload stored at `id`.
68    pub fn expression(&self, id: ExpressionNodeId) -> &Expression {
69        &self.node(id).expr
70    }
71
72    /// Returns the parent of `id`, or `None` for the root.
73    pub fn parent(&self, id: ExpressionNodeId) -> Option<ExpressionNodeId> {
74        self.node(id).parent
75    }
76
77    /// Returns the direct expression children of `id`.
78    pub fn children(&self, id: ExpressionNodeId) -> &[ExpressionNodeId] {
79        &self.node(id).children
80    }
81
82    /// Returns whether a direct child of `id` has the requested expression discriminant.
83    pub(crate) fn has_direct_child_discriminant(
84        &self,
85        id: ExpressionNodeId,
86        discriminant: usize,
87    ) -> bool {
88        self.node(id)
89            .direct_child_discriminants
90            .iter()
91            .any(|&(candidate, count)| candidate == discriminant && count != 0)
92    }
93
94    /// Returns whether `id` is still reachable from the current root.
95    pub fn is_reachable(&self, id: ExpressionNodeId) -> bool {
96        self.node(id).reachable
97    }
98
99    /// Returns the depth of `id` below the root.
100    pub fn depth(&self, id: ExpressionNodeId) -> usize {
101        self.node(id).depth
102    }
103
104    /// Builds the current preorder path of `id`.
105    ///
106    /// Lexicographic ordering of these paths is the model preorder used by the rewriter. Paths
107    /// are only materialised for the relatively rare ordering operations; storing one on every
108    /// node makes building a deep arena quadratic in the tree depth.
109    pub fn preorder_path(&self, id: ExpressionNodeId) -> Vec<usize> {
110        let mut path = Vec::with_capacity(self.depth(id));
111        let mut current = id;
112        while let Some(parent) = self.parent(current) {
113            path.push(
114                self.node(current)
115                    .child_slot
116                    .expect("non-root expression node must have a child slot"),
117            );
118            current = parent;
119        }
120        path.reverse();
121        path
122    }
123
124    /// Returns the generation counter for `id`.
125    pub fn generation(&self, id: ExpressionNodeId) -> u32 {
126        self.node(id).generation
127    }
128
129    /// Records that rewrite-relevant content at `id` has changed.
130    pub fn bump_generation(&mut self, id: ExpressionNodeId) {
131        {
132            let node = self.node_mut(id);
133            node.generation = node.generation.wrapping_add(1);
134        }
135        self.expression(id).invalidate_cached_content_hash();
136    }
137
138    /// Replaces the subtree at `id` while preserving `id` itself.
139    pub fn replace_subtree(&mut self, id: ExpressionNodeId, replacement: Expression) {
140        self.assert_valid_id(id);
141        replacement.invalidate_cached_content_hash();
142        replacement.meta_ref().clear_cached_domain();
143        let old_discriminant = discriminant_from_value(self.expression(id));
144        let new_discriminant = discriminant_from_value(&replacement);
145        let parent = self.parent(id);
146
147        let old_children = self.children(id).to_vec();
148        for old_child in old_children {
149            self.mark_subtree_unreachable(old_child);
150        }
151
152        let child_depth = self.depth(id) + 1;
153        let children = replacement
154            .children()
155            .into_iter()
156            .enumerate()
157            .map(|(child_index, child)| {
158                self.push_subtree(child, Some(id), Some(child_index), child_depth)
159            })
160            .collect::<Vec<_>>();
161        let mut direct_child_discriminants = Vec::new();
162        for &child_id in &children {
163            let discriminant = discriminant_from_value(self.expression(child_id));
164            if let Some((_, count)) = direct_child_discriminants
165                .iter_mut()
166                .find(|(candidate, _)| *candidate == discriminant)
167            {
168                *count += 1;
169            } else {
170                direct_child_discriminants.push((discriminant, 1));
171            }
172        }
173
174        let node = self.node_mut(id);
175        node.expr = replacement;
176        node.payload_child_generations = vec![0; children.len()];
177        node.children = children;
178        node.direct_child_discriminants = direct_child_discriminants;
179        node.generation = node.generation.wrapping_add(1);
180
181        if old_discriminant != new_discriminant
182            && let Some(parent_id) = parent
183        {
184            self.replace_direct_child_discriminant(parent_id, old_discriminant, new_discriminant);
185        }
186    }
187
188    /// Appends top-level constraints to the root expression.
189    pub fn add_root_children(&mut self, children: Vec<Expression>) -> Vec<ExpressionNodeId> {
190        let root = self.root;
191        assert!(
192            matches!(self.expression(root), Expression::Root(..)),
193            "arena root is not an Expression::Root"
194        );
195        self.rebuild_payload_from_children(root);
196        let first_new_child_index = self.children(root).len();
197
198        let new_children = children
199            .into_iter()
200            .enumerate()
201            .map(|(offset, child)| {
202                self.push_subtree(child, Some(root), Some(first_new_child_index + offset), 1)
203            })
204            .collect::<Vec<_>>();
205        self.node_mut(root)
206            .children
207            .extend(new_children.iter().copied());
208        for &child_id in &new_children {
209            let child_discriminant = discriminant_from_value(self.expression(child_id));
210            self.increment_direct_child_discriminant(root, child_discriminant);
211        }
212
213        let appended: Vec<_> = new_children
214            .iter()
215            .map(|&child| self.direct_child_expression(child))
216            .collect();
217        let generations = new_children
218            .iter()
219            .map(|&child| self.generation(child))
220            .collect::<Vec<_>>();
221        let root_node = self.node_mut(root);
222        let Expression::Root(_, payload) = &mut root_node.expr else {
223            unreachable!()
224        };
225        payload.extend(appended);
226        root_node.payload_child_generations.extend(generations);
227        root_node.generation = root_node.generation.wrapping_add(1);
228        self.invalidate_expression_hashes_to_root(root);
229        new_children
230    }
231
232    /// Rebuilds the stored expression payload at `id` from its direct arena children.
233    ///
234    /// Only children whose generation changed are copied into the payload, so refreshing a
235    /// wide expression does not clone its unchanged siblings. Ancestor repair should walk
236    /// upward so deeper nodes are refreshed first.
237    pub fn rebuild_payload_from_children(&mut self, id: ExpressionNodeId) {
238        let changed: Vec<_> = self
239            .children(id)
240            .iter()
241            .enumerate()
242            .filter_map(|(slot, &child)| {
243                let generation = self.generation(child);
244                (self.node(id).payload_child_generations.get(slot) != Some(&generation))
245                    .then_some((slot, child, generation))
246            })
247            .collect();
248        for (slot, child, generation) in changed {
249            let child_expr = self.direct_child_expression(child);
250            if !self
251                .node_mut(id)
252                .expr
253                .try_replace_child_at(slot, child_expr)
254            {
255                // Some expression shapes do not support slot replacement.
256                let children = self
257                    .children(id)
258                    .iter()
259                    .map(|&child| self.direct_child_expression(child))
260                    .collect();
261                let rebuilt = self.expression(id).with_children(children);
262                let generations = self
263                    .children(id)
264                    .iter()
265                    .map(|&child| self.generation(child))
266                    .collect();
267                let node = self.node_mut(id);
268                node.expr = rebuilt;
269                node.payload_child_generations = generations;
270                break;
271            }
272            self.node_mut(id).payload_child_generations[slot] = generation;
273        }
274        let node = self.node_mut(id);
275        node.expr.meta_ref().clear_cached_domain();
276        node.expr.invalidate_cached_content_hash();
277        node.generation = node.generation.wrapping_add(1);
278    }
279
280    /// Syncs the parent payload after a direct child changed.
281    ///
282    /// Uses [`Uniplate::try_replace_child_at`](uniplate::Uniplate::try_replace_child_at) so
283    /// same-arity updates avoid cloning siblings. The child's stored slot gives its position in
284    /// the parent in O(1). Falls back to a full rebuild if the child is missing or in-place replace
285    /// fails (e.g. arity mismatch).
286    pub fn sync_payload_for_changed_child(
287        &mut self,
288        parent_id: ExpressionNodeId,
289        child_id: ExpressionNodeId,
290    ) {
291        let Some(index) = self.direct_child_index(parent_id, child_id) else {
292            self.rebuild_payload_from_children(parent_id);
293            return;
294        };
295
296        let child_expr = self.direct_child_expression(child_id);
297        let replaced = self
298            .node_mut(parent_id)
299            .expr
300            .try_replace_child_at(index, child_expr);
301        if !replaced {
302            self.rebuild_payload_from_children(parent_id);
303            return;
304        }
305
306        let generation = self.generation(child_id);
307        let node = self.node_mut(parent_id);
308        node.payload_child_generations[index] = generation;
309        node.expr.meta_ref().clear_cached_domain();
310        node.expr.invalidate_cached_content_hash();
311        node.generation = node.generation.wrapping_add(1);
312    }
313
314    /// Returns `child_id`'s position among `parent_id`'s direct children.
315    ///
316    /// Checking the stored parent and child slot keeps this lookup safe for unreachable nodes
317    /// whose relationship metadata is retained after subtree replacement.
318    fn direct_child_index(
319        &self,
320        parent_id: ExpressionNodeId,
321        child_id: ExpressionNodeId,
322    ) -> Option<usize> {
323        let child = self.node(child_id);
324        if child.parent != Some(parent_id) {
325            return None;
326        }
327
328        let index = child.child_slot?;
329        (self.children(parent_id).get(index).copied() == Some(child_id)).then_some(index)
330    }
331
332    fn increment_direct_child_discriminant(&mut self, id: ExpressionNodeId, discriminant: usize) {
333        let counts = &mut self.node_mut(id).direct_child_discriminants;
334        if let Some((_, count)) = counts
335            .iter_mut()
336            .find(|(candidate, _)| *candidate == discriminant)
337        {
338            *count += 1;
339        } else {
340            counts.push((discriminant, 1));
341        }
342    }
343
344    fn replace_direct_child_discriminant(
345        &mut self,
346        id: ExpressionNodeId,
347        old_discriminant: usize,
348        new_discriminant: usize,
349    ) {
350        let counts = &mut self.node_mut(id).direct_child_discriminants;
351        let old_index = counts
352            .iter()
353            .position(|(candidate, _)| *candidate == old_discriminant)
354            .expect("old direct-child discriminant must be indexed");
355        counts[old_index].1 -= 1;
356        if counts[old_index].1 == 0 {
357            counts.swap_remove(old_index);
358        }
359
360        if let Some((_, count)) = counts
361            .iter_mut()
362            .find(|(candidate, _)| *candidate == new_discriminant)
363        {
364            *count += 1;
365        } else {
366            counts.push((new_discriminant, 1));
367        }
368    }
369
370    /// Clears cached expression content hashes from `id` through the root.
371    fn invalidate_expression_hashes_to_root(&mut self, id: ExpressionNodeId) {
372        let mut node = Some(id);
373        while let Some(node_id) = node {
374            let parent = self.node(node_id).parent;
375            self.expression(node_id).invalidate_cached_content_hash();
376            node = parent;
377        }
378    }
379
380    /// Rebuilds the expression tree reachable from the arena root.
381    pub fn into_root_expression(self) -> Expression {
382        self.expression_from(self.root)
383    }
384
385    /// Moves out the root payload when callers have kept ancestor payloads synchronized.
386    ///
387    /// Unlike [`Self::into_root_expression`], this does not recursively rebuild the tree. It is
388    /// intended for the rewriter's settled arena, where every changed child has already been
389    /// propagated to the root.
390    pub(crate) fn into_synced_root_expression(self) -> Expression {
391        self.nodes
392            .into_iter()
393            .nth(self.root.0)
394            .expect("arena must contain its root node")
395            .expr
396    }
397
398    fn direct_child_expression(&self, id: ExpressionNodeId) -> Expression {
399        self.expression(id).clone()
400    }
401
402    /// Rebuilds the expression tree reachable from `id`.
403    pub fn expression_from(&self, id: ExpressionNodeId) -> Expression {
404        let node = self.node(id);
405        let children = node
406            .children
407            .iter()
408            .map(|child| self.expression_from(*child))
409            .collect::<VecDeque<_>>();
410
411        let rebuilt = node.expr.with_children(children);
412        rebuilt.invalidate_cached_content_hash();
413        rebuilt
414    }
415
416    /// Returns the number of slots currently allocated in the arena.
417    ///
418    /// This includes unreachable slots left behind by subtree replacement.
419    pub fn len(&self) -> usize {
420        self.nodes.len()
421    }
422
423    /// Returns true when the arena contains no nodes.
424    pub fn is_empty(&self) -> bool {
425        self.nodes.is_empty()
426    }
427
428    /// Returns reachable node ids under `id` in rewriter preorder.
429    pub fn reachable_subtree_ids(&self, id: ExpressionNodeId) -> Vec<ExpressionNodeId> {
430        fn collect(
431            arena: &ExpressionArena,
432            node_id: ExpressionNodeId,
433            nodes: &mut Vec<ExpressionNodeId>,
434        ) {
435            if !arena.is_reachable(node_id) {
436                return;
437            }
438            nodes.push(node_id);
439            for &child in arena.children(node_id) {
440                collect(arena, child, nodes);
441            }
442        }
443
444        let mut nodes = Vec::new();
445        collect(self, id, &mut nodes);
446        nodes
447    }
448
449    fn push_subtree(
450        &mut self,
451        expr: Expression,
452        parent: Option<ExpressionNodeId>,
453        child_slot: Option<usize>,
454        depth: usize,
455    ) -> ExpressionNodeId {
456        expr.invalidate_cached_content_hash();
457        expr.meta_ref().clear_cached_domain();
458        let id = ExpressionNodeId(self.nodes.len());
459        let child_exprs = expr.children();
460        self.nodes.push(ExpressionArenaNode {
461            expr,
462            parent,
463            children: Vec::new(),
464            payload_child_generations: Vec::new(),
465            direct_child_discriminants: Vec::new(),
466            reachable: true,
467            child_slot,
468            depth,
469            generation: 0,
470        });
471
472        let children: Vec<_> = child_exprs
473            .into_iter()
474            .enumerate()
475            .map(|(child_index, child)| {
476                self.push_subtree(child, Some(id), Some(child_index), depth + 1)
477            })
478            .collect();
479        self.nodes[id.0].payload_child_generations = vec![0; children.len()];
480        self.nodes[id.0].children = children;
481        let child_discriminants = self.nodes[id.0]
482            .children
483            .iter()
484            .map(|&child_id| discriminant_from_value(self.expression(child_id)))
485            .collect::<Vec<_>>();
486        for child_discriminant in child_discriminants {
487            self.increment_direct_child_discriminant(id, child_discriminant);
488        }
489
490        id
491    }
492
493    fn mark_subtree_unreachable(&mut self, id: ExpressionNodeId) {
494        if !self.node(id).reachable {
495            return;
496        }
497
498        let children = std::mem::take(&mut self.node_mut(id).children);
499        let node = self.node_mut(id);
500        node.expr = Expression::Root(super::Metadata::new(), Vec::new());
501        node.direct_child_discriminants.clear();
502        node.reachable = false;
503        node.generation = node.generation.wrapping_add(1);
504
505        for child in children {
506            self.mark_subtree_unreachable(child);
507        }
508    }
509
510    fn node(&self, id: ExpressionNodeId) -> &ExpressionArenaNode {
511        self.nodes
512            .get(id.0)
513            .unwrap_or_else(|| panic!("invalid expression node id: {id:?}"))
514    }
515
516    fn node_mut(&mut self, id: ExpressionNodeId) -> &mut ExpressionArenaNode {
517        self.nodes
518            .get_mut(id.0)
519            .unwrap_or_else(|| panic!("invalid expression node id: {id:?}"))
520    }
521
522    fn assert_valid_id(&self, id: ExpressionNodeId) {
523        let _ = self.node(id);
524    }
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530    use crate::ast::{Metadata, Moo};
531
532    fn int(value: i32) -> Expression {
533        value.into()
534    }
535
536    fn eq(left: Expression, right: Expression) -> Expression {
537        Expression::Eq(Metadata::new(), Moo::new(left), Moo::new(right))
538    }
539
540    fn root(exprs: Vec<Expression>) -> Expression {
541        Expression::Root(Metadata::new(), exprs)
542    }
543
544    #[test]
545    fn round_trips_expression_tree() {
546        let expr = eq(int(1), eq(int(2), int(3)));
547        let arena = ExpressionArena::from_root(expr.clone());
548
549        assert_eq!(arena.into_root_expression(), expr);
550    }
551
552    #[test]
553    fn records_parent_and_child_links() {
554        let arena = ExpressionArena::from_root(eq(int(1), int(2)));
555        let root = arena.root();
556        let children = arena.children(root);
557
558        assert_eq!(children.len(), 2);
559        assert_eq!(arena.parent(root), None);
560        assert_eq!(arena.parent(children[0]), Some(root));
561        assert_eq!(arena.parent(children[1]), Some(root));
562    }
563
564    #[test]
565    fn direct_child_discriminants_update_without_losing_duplicate_counts() {
566        let mut arena = ExpressionArena::from_root(root(vec![int(1), int(2)]));
567        let root_id = arena.root();
568        let children = arena.children(root_id).to_vec();
569        let atomic = discriminant_from_value(arena.expression(children[0]));
570        let equality = discriminant_from_value(&eq(int(1), int(2)));
571
572        assert!(arena.has_direct_child_discriminant(root_id, atomic));
573        assert!(!arena.has_direct_child_discriminant(root_id, equality));
574
575        arena.replace_subtree(children[0], eq(int(3), int(4)));
576        assert!(arena.has_direct_child_discriminant(root_id, atomic));
577        assert!(arena.has_direct_child_discriminant(root_id, equality));
578        assert!(arena.has_direct_child_discriminant(children[0], atomic));
579
580        arena.replace_subtree(children[1], eq(int(5), int(6)));
581        assert!(!arena.has_direct_child_discriminant(root_id, atomic));
582        assert!(arena.has_direct_child_discriminant(root_id, equality));
583    }
584
585    #[test]
586    fn replacement_marks_old_descendants_unreachable_and_paths_new_children() {
587        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(3)]));
588        let root_id = arena.root();
589        let first_child = arena.children(root_id)[0];
590        let old_left = arena.children(first_child)[0];
591
592        arena.replace_subtree(first_child, eq(int(4), int(5)));
593
594        assert!(!arena.is_reachable(old_left));
595        assert_eq!(arena.preorder_path(first_child), &[0]);
596
597        let new_children = arena.children(first_child);
598        assert_eq!(arena.preorder_path(new_children[0]), &[0, 0]);
599        assert_eq!(arena.preorder_path(new_children[1]), &[0, 1]);
600        assert!(
601            arena
602                .reachable_subtree_ids(first_child)
603                .contains(&new_children[0])
604        );
605        assert!(!arena.reachable_subtree_ids(first_child).contains(&old_left));
606    }
607
608    #[test]
609    fn added_root_children_get_preorder_paths() {
610        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
611        let added = arena.add_root_children(vec![int(2), int(3)]);
612
613        assert_eq!(added.len(), 2);
614        assert_eq!(arena.preorder_path(added[0]), &[1]);
615        assert_eq!(arena.preorder_path(added[1]), &[2]);
616        assert_eq!(arena.parent(added[0]), Some(arena.root()));
617        assert_eq!(arena.parent(added[1]), Some(arena.root()));
618    }
619
620    #[test]
621    fn obtains_direct_child_indices_from_preorder_paths() {
622        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(3)]));
623        let root_id = arena.root();
624        let eq_id = arena.children(root_id)[0];
625        let eq_children = arena.children(eq_id).to_vec();
626
627        assert_eq!(arena.direct_child_index(root_id, eq_id), Some(0));
628        assert_eq!(arena.direct_child_index(eq_id, eq_children[0]), Some(0));
629        assert_eq!(arena.direct_child_index(eq_id, eq_children[1]), Some(1));
630        assert_eq!(arena.direct_child_index(root_id, eq_children[0]), None);
631        assert_eq!(arena.direct_child_index(eq_id, root_id), None);
632
633        let added = arena.add_root_children(vec![int(4)]);
634        assert_eq!(arena.direct_child_index(root_id, added[0]), Some(2));
635    }
636
637    #[test]
638    fn direct_child_index_rejects_unreachable_replaced_children() {
639        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2))]));
640        let root_id = arena.root();
641        let eq_id = arena.children(root_id)[0];
642        let old_left = arena.children(eq_id)[0];
643
644        arena.replace_subtree(eq_id, eq(int(3), int(4)));
645
646        assert!(!arena.is_reachable(old_left));
647        assert_eq!(arena.direct_child_index(eq_id, old_left), None);
648        assert_eq!(
649            arena.direct_child_index(eq_id, arena.children(eq_id)[0]),
650            Some(0)
651        );
652    }
653
654    #[test]
655    fn adding_root_children_bumps_root_generation() {
656        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
657        let root_id = arena.root();
658        let before = arena.generation(root_id);
659
660        arena.add_root_children(vec![int(2)]);
661
662        assert_ne!(arena.generation(root_id), before);
663    }
664
665    #[test]
666    fn replaces_subtree_without_changing_replaced_node_id() {
667        let mut arena = ExpressionArena::from_root(eq(int(1), int(2)));
668        let root = arena.root();
669        let left = arena.children(root)[0];
670
671        arena.replace_subtree(left, eq(int(3), int(4)));
672
673        assert_eq!(arena.children(root)[0], left);
674        assert_eq!(arena.parent(arena.children(left)[0]), Some(left));
675        assert_eq!(arena.into_root_expression(), eq(eq(int(3), int(4)), int(2)));
676    }
677
678    #[test]
679    fn rebuilds_ancestor_payload_after_child_replacement() {
680        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2))]));
681        let root_id = arena.root();
682        let eq_id = arena.children(root_id)[0];
683        let left = arena.children(eq_id)[0];
684        let old_eq_hash = arena.expression(eq_id).cached_content_hash();
685        let old_root_hash = arena.expression(root_id).cached_content_hash();
686
687        arena.replace_subtree(left, int(3));
688        arena.rebuild_payload_from_children(eq_id);
689        arena.rebuild_payload_from_children(root_id);
690
691        assert_eq!(arena.expression(eq_id), &eq(int(3), int(2)));
692        assert_eq!(arena.expression(root_id), &root(vec![eq(int(3), int(2))]));
693        assert_ne!(arena.expression(eq_id).cached_content_hash(), old_eq_hash);
694        assert_ne!(
695            arena.expression(root_id).cached_content_hash(),
696            old_root_hash
697        );
698    }
699
700    #[test]
701    fn incremental_rebuild_handles_multiple_batches_and_changed_arity() {
702        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
703        let root_id = arena.root();
704        let eq_id = arena.children(root_id)[0];
705        let children = arena.children(eq_id).to_vec();
706        arena.replace_subtree(children[0], int(3));
707        arena.replace_subtree(children[1], int(4));
708        arena.rebuild_payload_from_children(eq_id);
709        arena.rebuild_payload_from_children(root_id);
710        assert_eq!(
711            arena.expression(root_id),
712            &root(vec![eq(int(3), int(4)), int(9)])
713        );
714
715        arena.replace_subtree(eq_id, root(vec![int(5), int(6), int(7)]));
716        let child = arena.children(eq_id)[2];
717        arena.replace_subtree(child, int(8));
718        arena.rebuild_payload_from_children(eq_id);
719        arena.rebuild_payload_from_children(root_id);
720        assert_eq!(
721            arena.expression(root_id),
722            &root(vec![root(vec![int(5), int(6), int(8)]), int(9)])
723        );
724
725        let added = arena.add_root_children(vec![int(10)]);
726        arena.replace_subtree(added[0], int(11));
727        arena.rebuild_payload_from_children(root_id);
728        assert_eq!(
729            arena.expression(root_id),
730            &root(vec![root(vec![int(5), int(6), int(8)]), int(9), int(11)])
731        );
732    }
733
734    #[test]
735    fn sync_and_rebuild_agree_after_child_replacement() {
736        let mut sync_arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
737        let mut rebuild_arena = sync_arena.clone();
738
739        for arena in [&mut sync_arena, &mut rebuild_arena] {
740            let root_id = arena.root();
741            let eq_id = arena.children(root_id)[0];
742            let left = arena.children(eq_id)[0];
743            arena.replace_subtree(left, int(3));
744        }
745
746        let sync_root = sync_arena.root();
747        let sync_eq = sync_arena.children(sync_root)[0];
748        let sync_left = sync_arena.children(sync_eq)[0];
749        sync_arena.sync_payload_for_changed_child(sync_eq, sync_left);
750        sync_arena.sync_payload_for_changed_child(sync_root, sync_eq);
751
752        let rebuild_root = rebuild_arena.root();
753        let rebuild_eq = rebuild_arena.children(rebuild_root)[0];
754        rebuild_arena.rebuild_payload_from_children(rebuild_eq);
755        rebuild_arena.rebuild_payload_from_children(rebuild_root);
756
757        assert_eq!(
758            sync_arena.into_root_expression(),
759            rebuild_arena.into_root_expression()
760        );
761    }
762
763    #[test]
764    fn syncs_same_arity_ancestor_payload_without_full_rebuild() {
765        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
766        let root_id = arena.root();
767        let eq_id = arena.children(root_id)[0];
768        let left = arena.children(eq_id)[0];
769        let untouched = arena.children(root_id)[1];
770        let untouched_before = arena.expression(untouched).clone();
771
772        arena.replace_subtree(left, int(3));
773        arena.sync_payload_for_changed_child(eq_id, left);
774        arena.sync_payload_for_changed_child(root_id, eq_id);
775
776        assert_eq!(arena.expression(eq_id), &eq(int(3), int(2)));
777        assert_eq!(
778            arena.expression(root_id),
779            &root(vec![eq(int(3), int(2)), int(9)])
780        );
781        // Sibling payload slot is left in place (same expression value / identity of content).
782        assert_eq!(arena.expression(untouched), &untouched_before);
783        assert_eq!(
784            arena.into_root_expression(),
785            root(vec![eq(int(3), int(2)), int(9)])
786        );
787    }
788
789    #[test]
790    fn syncs_wide_matrix_child_slot_in_place() {
791        use crate::ast::{AbstractLiteral, Domain, Range};
792        use crate::into_matrix_expr;
793
794        let elems: Vec<Expression> = (0..32).map(int).collect();
795        let matrix = into_matrix_expr![elems; Domain::int(vec![Range::Bounded(1, 32)])];
796        let mut arena =
797            ExpressionArena::from_root(Expression::And(Metadata::new(), Moo::new(matrix)));
798        let and_id = arena.root();
799        let matrix_id = arena.children(and_id)[0];
800        let target = arena.children(matrix_id)[17];
801
802        arena.replace_subtree(target, int(99));
803        arena.sync_payload_for_changed_child(matrix_id, target);
804        arena.sync_payload_for_changed_child(and_id, matrix_id);
805
806        let Expression::And(_, inner) = arena.expression(and_id) else {
807            panic!("expected And");
808        };
809        let Expression::AbstractLiteral(_, AbstractLiteral::Matrix(elems, _)) = inner.as_ref()
810        else {
811            panic!("expected matrix");
812        };
813        assert_eq!(elems.len(), 32);
814        assert_eq!(elems[17], int(99));
815        assert_eq!(elems[0], int(0));
816        assert_eq!(elems[31], int(31));
817    }
818
819    #[test]
820    fn appending_constraints_preserves_pending_child_updates() {
821        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
822        let root_id = arena.root();
823        let first = arena.children(root_id)[0];
824        arena.replace_subtree(first, int(2));
825        arena.add_root_children(vec![int(3)]);
826        assert_eq!(arena.expression(root_id), &root(vec![int(2), int(3)]));
827        arena.replace_subtree(first, int(4));
828        arena.rebuild_payload_from_children(root_id);
829        assert_eq!(arena.expression(root_id), &root(vec![int(4), int(3)]));
830    }
831
832    #[test]
833    fn appends_root_children() {
834        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
835
836        arena.add_root_children(vec![int(2), int(3)]);
837
838        assert_eq!(
839            arena.into_root_expression(),
840            root(vec![int(1), int(2), int(3)])
841        );
842    }
843}
conjure_cp_core/ast/expression_arena.rs

conjure_cp_core/ast/
expression_arena.rs
