1
use std::collections::VecDeque;
2
use uniplate::Uniplate;
3

            
4
use 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)]
8
pub struct ExpressionNodeId(usize);
9

            
10
impl ExpressionNodeId {
11
    /// Returns the arena slot index for this node id.
12
72591424205
    pub fn index(self) -> usize {
13
72591424205
        self.0
14
72591424205
    }
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)]
23
pub struct ExpressionArena {
24
    nodes: Vec<ExpressionArenaNode>,
25
    root: ExpressionNodeId,
26
}
27

            
28
#[derive(Clone, Debug)]
29
struct 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

            
51
impl ExpressionArena {
52
    /// Builds an arena from an expression tree.
53
2638728
    pub fn from_root(root: Expression) -> Self {
54
2638728
        let mut arena = Self {
55
2638728
            nodes: Vec::new(),
56
2638728
            root: ExpressionNodeId(0),
57
2638728
        };
58
2638728
        arena.root = arena.push_subtree(root, None, None, 0);
59
2638728
        arena
60
2638728
    }
61

            
62
    /// Returns the root node id.
63
151822139
    pub fn root(&self) -> ExpressionNodeId {
64
151822139
        self.root
65
151822139
    }
66

            
67
    /// Returns the expression payload stored at `id`.
68
66948211022
    pub fn expression(&self, id: ExpressionNodeId) -> &Expression {
69
66948211022
        &self.node(id).expr
70
66948211022
    }
71

            
72
    /// Returns the parent of `id`, or `None` for the root.
73
2748463744
    pub fn parent(&self, id: ExpressionNodeId) -> Option<ExpressionNodeId> {
74
2748463744
        self.node(id).parent
75
2748463744
    }
76

            
77
    /// Returns the direct expression children of `id`.
78
28443172306
    pub fn children(&self, id: ExpressionNodeId) -> &[ExpressionNodeId] {
79
28443172306
        &self.node(id).children
80
28443172306
    }
81

            
82
    /// Returns whether a direct child of `id` has the requested expression discriminant.
83
927250122
    pub(crate) fn has_direct_child_discriminant(
84
927250122
        &self,
85
927250122
        id: ExpressionNodeId,
86
927250122
        discriminant: usize,
87
927250122
    ) -> bool {
88
927250122
        self.node(id)
89
927250122
            .direct_child_discriminants
90
927250122
            .iter()
91
921817059
            .any(|&(candidate, count)| candidate == discriminant && count != 0)
92
927250122
    }
93

            
94
    /// Returns whether `id` is still reachable from the current root.
95
39274089636
    pub fn is_reachable(&self, id: ExpressionNodeId) -> bool {
96
39274089636
        self.node(id).reachable
97
39274089636
    }
98

            
99
    /// Returns the depth of `id` below the root.
100
3252694688
    pub fn depth(&self, id: ExpressionNodeId) -> usize {
101
3252694688
        self.node(id).depth
102
3252694688
    }
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
92860756
    pub fn preorder_path(&self, id: ExpressionNodeId) -> Vec<usize> {
110
92860756
        let mut path = Vec::with_capacity(self.depth(id));
111
92860756
        let mut current = id;
112
384582758
        while let Some(parent) = self.parent(current) {
113
291722002
            path.push(
114
291722002
                self.node(current)
115
291722002
                    .child_slot
116
291722002
                    .expect("non-root expression node must have a child slot"),
117
291722002
            );
118
291722002
            current = parent;
119
        }
120
92860756
        path.reverse();
121
92860756
        path
122
92860756
    }
123

            
124
    /// Returns the generation counter for `id`.
125
33124309372
    pub fn generation(&self, id: ExpressionNodeId) -> u32 {
126
33124309372
        self.node(id).generation
127
33124309372
    }
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
81135606
    pub fn replace_subtree(&mut self, id: ExpressionNodeId, replacement: Expression) {
140
81135606
        self.assert_valid_id(id);
141
81135606
        replacement.invalidate_cached_content_hash();
142
81135606
        replacement.meta_ref().clear_cached_domain();
143
81135606
        let old_discriminant = discriminant_from_value(self.expression(id));
144
81135606
        let new_discriminant = discriminant_from_value(&replacement);
145
81135606
        let parent = self.parent(id);
146

            
147
81135606
        let old_children = self.children(id).to_vec();
148
81135606
        for old_child in old_children {
149
265042995
            self.mark_subtree_unreachable(old_child);
150
        }
151

            
152
81135606
        let child_depth = self.depth(id) + 1;
153
81135606
        let children = replacement
154
81135606
            .children()
155
81135606
            .into_iter()
156
81135606
            .enumerate()
157
205915463
            .map(|(child_index, child)| {
158
205915463
                self.push_subtree(child, Some(id), Some(child_index), child_depth)
159
205915463
            })
160
            .collect::<Vec<_>>();
161
81135606
        let mut direct_child_discriminants = Vec::new();
162
81135606
        for &child_id in &children {
163
205915463
            let discriminant = discriminant_from_value(self.expression(child_id));
164
205915463
            if let Some((_, count)) = direct_child_discriminants
165
205915463
                .iter_mut()
166
394538105
                .find(|(candidate, _)| *candidate == discriminant)
167
            {
168
144115924
                *count += 1;
169
            } else {
170
61799539
                direct_child_discriminants.push((discriminant, 1));
171
            }
172
        }
173

            
174
81135606
        let node = self.node_mut(id);
175
81135606
        node.expr = replacement;
176
81135606
        node.payload_child_generations = vec![0; children.len()];
177
81135606
        node.children = children;
178
81135606
        node.direct_child_discriminants = direct_child_discriminants;
179
81135606
        node.generation = node.generation.wrapping_add(1);
180

            
181
81135606
        if old_discriminant != new_discriminant
182
40672417
            && let Some(parent_id) = parent
183
        {
184
40439041
            self.replace_direct_child_discriminant(parent_id, old_discriminant, new_discriminant);
185
40696565
        }
186
81135606
    }
187

            
188
    /// Appends top-level constraints to the root expression.
189
3872131
    pub fn add_root_children(&mut self, children: Vec<Expression>) -> Vec<ExpressionNodeId> {
190
3872131
        let root = self.root;
191
3872131
        assert!(
192
3872131
            matches!(self.expression(root), Expression::Root(..)),
193
            "arena root is not an Expression::Root"
194
        );
195
3872131
        self.rebuild_payload_from_children(root);
196
3872131
        let first_new_child_index = self.children(root).len();
197

            
198
3872131
        let new_children = children
199
3872131
            .into_iter()
200
3872131
            .enumerate()
201
11855778
            .map(|(offset, child)| {
202
11855778
                self.push_subtree(child, Some(root), Some(first_new_child_index + offset), 1)
203
11855778
            })
204
            .collect::<Vec<_>>();
205
3872131
        self.node_mut(root)
206
3872131
            .children
207
3872131
            .extend(new_children.iter().copied());
208
3872131
        for &child_id in &new_children {
209
11855778
            let child_discriminant = discriminant_from_value(self.expression(child_id));
210
11855778
            self.increment_direct_child_discriminant(root, child_discriminant);
211
        }
212

            
213
3872131
        let appended: Vec<_> = new_children
214
3872131
            .iter()
215
11855778
            .map(|&child| self.direct_child_expression(child))
216
            .collect();
217
3872131
        let generations = new_children
218
3872131
            .iter()
219
11855778
            .map(|&child| self.generation(child))
220
            .collect::<Vec<_>>();
221
3872131
        let root_node = self.node_mut(root);
222
3872131
        let Expression::Root(_, payload) = &mut root_node.expr else {
223
            unreachable!()
224
        };
225
3872131
        payload.extend(appended);
226
3872131
        root_node.payload_child_generations.extend(generations);
227
3872131
        root_node.generation = root_node.generation.wrapping_add(1);
228
3872131
        self.invalidate_expression_hashes_to_root(root);
229
3872131
        new_children
230
3872131
    }
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
147276748
    pub fn rebuild_payload_from_children(&mut self, id: ExpressionNodeId) {
238
147276748
        let changed: Vec<_> = self
239
147276748
            .children(id)
240
147276748
            .iter()
241
147276748
            .enumerate()
242
1436099530
            .filter_map(|(slot, &child)| {
243
1436099530
                let generation = self.generation(child);
244
1436099530
                (self.node(id).payload_child_generations.get(slot) != Some(&generation))
245
1436099530
                    .then_some((slot, child, generation))
246
1436099530
            })
247
            .collect();
248
147276748
        for (slot, child, generation) in changed {
249
169831945
            let child_expr = self.direct_child_expression(child);
250
169831945
            if !self
251
169831945
                .node_mut(id)
252
169831945
                .expr
253
169831945
                .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
169831945
            }
272
169831945
            self.node_mut(id).payload_child_generations[slot] = generation;
273
        }
274
147276748
        let node = self.node_mut(id);
275
147276748
        node.expr.meta_ref().clear_cached_domain();
276
147276748
        node.expr.invalidate_cached_content_hash();
277
147276748
        node.generation = node.generation.wrapping_add(1);
278
147276748
    }
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
1459577772
    pub fn sync_payload_for_changed_child(
287
1459577772
        &mut self,
288
1459577772
        parent_id: ExpressionNodeId,
289
1459577772
        child_id: ExpressionNodeId,
290
1459577772
    ) {
291
1459577772
        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
1459577772
        let child_expr = self.direct_child_expression(child_id);
297
1459577772
        let replaced = self
298
1459577772
            .node_mut(parent_id)
299
1459577772
            .expr
300
1459577772
            .try_replace_child_at(index, child_expr);
301
1459577772
        if !replaced {
302
            self.rebuild_payload_from_children(parent_id);
303
            return;
304
1459577772
        }
305

            
306
1459577772
        let generation = self.generation(child_id);
307
1459577772
        let node = self.node_mut(parent_id);
308
1459577772
        node.payload_child_generations[index] = generation;
309
1459577772
        node.expr.meta_ref().clear_cached_domain();
310
1459577772
        node.expr.invalidate_cached_content_hash();
311
1459577772
        node.generation = node.generation.wrapping_add(1);
312
1459577772
    }
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
1459577780
    fn direct_child_index(
319
1459577780
        &self,
320
1459577780
        parent_id: ExpressionNodeId,
321
1459577780
        child_id: ExpressionNodeId,
322
1459577780
    ) -> Option<usize> {
323
1459577780
        let child = self.node(child_id);
324
1459577780
        if child.parent != Some(parent_id) {
325
2
            return None;
326
1459577778
        }
327

            
328
1459577778
        let index = child.child_slot?;
329
1459577778
        (self.children(parent_id).get(index).copied() == Some(child_id)).then_some(index)
330
1459577780
    }
331

            
332
1309510644
    fn increment_direct_child_discriminant(&mut self, id: ExpressionNodeId, discriminant: usize) {
333
1309510644
        let counts = &mut self.node_mut(id).direct_child_discriminants;
334
1309510644
        if let Some((_, count)) = counts
335
1309510644
            .iter_mut()
336
595640828
            .find(|(candidate, _)| *candidate == discriminant)
337
        {
338
392824874
            *count += 1;
339
        } else {
340
916685770
            counts.push((discriminant, 1));
341
        }
342
1309510644
    }
343

            
344
40439041
    fn replace_direct_child_discriminant(
345
40439041
        &mut self,
346
40439041
        id: ExpressionNodeId,
347
40439041
        old_discriminant: usize,
348
40439041
        new_discriminant: usize,
349
40439041
    ) {
350
40439041
        let counts = &mut self.node_mut(id).direct_child_discriminants;
351
40439041
        let old_index = counts
352
40439041
            .iter()
353
56177769
            .position(|(candidate, _)| *candidate == old_discriminant)
354
40439041
            .expect("old direct-child discriminant must be indexed");
355
40439041
        counts[old_index].1 -= 1;
356
40439041
        if counts[old_index].1 == 0 {
357
30584513
            counts.swap_remove(old_index);
358
9854528
        }
359

            
360
40439041
        if let Some((_, count)) = counts
361
40439041
            .iter_mut()
362
42423314
            .find(|(candidate, _)| *candidate == new_discriminant)
363
        {
364
12607987
            *count += 1;
365
        } else {
366
27831054
            counts.push((new_discriminant, 1));
367
        }
368
40439041
    }
369

            
370
    /// Clears cached expression content hashes from `id` through the root.
371
3872131
    fn invalidate_expression_hashes_to_root(&mut self, id: ExpressionNodeId) {
372
3872131
        let mut node = Some(id);
373
7744262
        while let Some(node_id) = node {
374
3872131
            let parent = self.node(node_id).parent;
375
3872131
            self.expression(node_id).invalidate_cached_content_hash();
376
3872131
            node = parent;
377
        }
378
3872131
    }
379

            
380
    /// Rebuilds the expression tree reachable from the arena root.
381
6
    pub fn into_root_expression(self) -> Expression {
382
6
        self.expression_from(self.root)
383
6
    }
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
2638702
    pub(crate) fn into_synced_root_expression(self) -> Expression {
391
2638702
        self.nodes
392
2638702
            .into_iter()
393
2638702
            .nth(self.root.0)
394
2638702
            .expect("arena must contain its root node")
395
2638702
            .expr
396
2638702
    }
397

            
398
1641265495
    fn direct_child_expression(&self, id: ExpressionNodeId) -> Expression {
399
1641265495
        self.expression(id).clone()
400
1641265495
    }
401

            
402
    /// Rebuilds the expression tree reachable from `id`.
403
62629
    pub fn expression_from(&self, id: ExpressionNodeId) -> Expression {
404
62629
        let node = self.node(id);
405
62629
        let children = node
406
62629
            .children
407
62629
            .iter()
408
52543
            .map(|child| self.expression_from(*child))
409
            .collect::<VecDeque<_>>();
410

            
411
62629
        let rebuilt = node.expr.with_children(children);
412
62629
        rebuilt.invalidate_cached_content_hash();
413
62629
        rebuilt
414
62629
    }
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
2
    pub fn reachable_subtree_ids(&self, id: ExpressionNodeId) -> Vec<ExpressionNodeId> {
430
6
        fn collect(
431
6
            arena: &ExpressionArena,
432
6
            node_id: ExpressionNodeId,
433
6
            nodes: &mut Vec<ExpressionNodeId>,
434
6
        ) {
435
6
            if !arena.is_reachable(node_id) {
436
                return;
437
6
            }
438
6
            nodes.push(node_id);
439
6
            for &child in arena.children(node_id) {
440
4
                collect(arena, child, nodes);
441
            }
442
6
        }
443

            
444
2
        let mut nodes = Vec::new();
445
2
        collect(self, id, &mut nodes);
446
2
        nodes
447
2
    }
448

            
449
1518064835
    fn push_subtree(
450
1518064835
        &mut self,
451
1518064835
        expr: Expression,
452
1518064835
        parent: Option<ExpressionNodeId>,
453
1518064835
        child_slot: Option<usize>,
454
1518064835
        depth: usize,
455
1518064835
    ) -> ExpressionNodeId {
456
1518064835
        expr.invalidate_cached_content_hash();
457
1518064835
        expr.meta_ref().clear_cached_domain();
458
1518064835
        let id = ExpressionNodeId(self.nodes.len());
459
1518064835
        let child_exprs = expr.children();
460
1518064835
        self.nodes.push(ExpressionArenaNode {
461
1518064835
            expr,
462
1518064835
            parent,
463
1518064835
            children: Vec::new(),
464
1518064835
            payload_child_generations: Vec::new(),
465
1518064835
            direct_child_discriminants: Vec::new(),
466
1518064835
            reachable: true,
467
1518064835
            child_slot,
468
1518064835
            depth,
469
1518064835
            generation: 0,
470
1518064835
        });
471

            
472
1518064835
        let children: Vec<_> = child_exprs
473
1518064835
            .into_iter()
474
1518064835
            .enumerate()
475
1297654866
            .map(|(child_index, child)| {
476
1297654866
                self.push_subtree(child, Some(id), Some(child_index), depth + 1)
477
1297654866
            })
478
            .collect();
479
1518064835
        self.nodes[id.0].payload_child_generations = vec![0; children.len()];
480
1518064835
        self.nodes[id.0].children = children;
481
1518064835
        let child_discriminants = self.nodes[id.0]
482
1518064835
            .children
483
1518064835
            .iter()
484
1297654866
            .map(|&child_id| discriminant_from_value(self.expression(child_id)))
485
            .collect::<Vec<_>>();
486
1518064835
        for child_discriminant in child_discriminants {
487
1297654866
            self.increment_direct_child_discriminant(id, child_discriminant);
488
        }
489

            
490
1518064835
        id
491
1518064835
    }
492

            
493
1507168529
    fn mark_subtree_unreachable(&mut self, id: ExpressionNodeId) {
494
1507168529
        if !self.node(id).reachable {
495
            return;
496
1507168529
        }
497

            
498
1507168529
        let children = std::mem::take(&mut self.node_mut(id).children);
499
1507168529
        let node = self.node_mut(id);
500
1507168529
        node.expr = Expression::Root(super::Metadata::new(), Vec::new());
501
1507168529
        node.direct_child_discriminants.clear();
502
1507168529
        node.reachable = false;
503
1507168529
        node.generation = node.generation.wrapping_add(1);
504

            
505
1507168529
        for child in children {
506
1242125534
            self.mark_subtree_unreachable(child);
507
        }
508
1507168529
    }
509

            
510
179497829097
    fn node(&self, id: ExpressionNodeId) -> &ExpressionArenaNode {
511
179497829097
        self.nodes
512
179497829097
            .get(id.0)
513
            .unwrap_or_else(|| panic!("invalid expression node id: {id:?}"))
514
179497829097
    }
515

            
516
7859262793
    fn node_mut(&mut self, id: ExpressionNodeId) -> &mut ExpressionArenaNode {
517
7859262793
        self.nodes
518
7859262793
            .get_mut(id.0)
519
            .unwrap_or_else(|| panic!("invalid expression node id: {id:?}"))
520
7859262793
    }
521

            
522
81135606
    fn assert_valid_id(&self, id: ExpressionNodeId) {
523
81135606
        let _ = self.node(id);
524
81135606
    }
525
}
526

            
527
#[cfg(test)]
528
mod tests {
529
    use super::*;
530
    use crate::ast::{Metadata, Moo};
531

            
532
135
    fn int(value: i32) -> Expression {
533
135
        value.into()
534
135
    }
535

            
536
25
    fn eq(left: Expression, right: Expression) -> Expression {
537
25
        Expression::Eq(Metadata::new(), Moo::new(left), Moo::new(right))
538
25
    }
539

            
540
24
    fn root(exprs: Vec<Expression>) -> Expression {
541
24
        Expression::Root(Metadata::new(), exprs)
542
24
    }
543

            
544
    #[test]
545
1
    fn round_trips_expression_tree() {
546
1
        let expr = eq(int(1), eq(int(2), int(3)));
547
1
        let arena = ExpressionArena::from_root(expr.clone());
548

            
549
1
        assert_eq!(arena.into_root_expression(), expr);
550
1
    }
551

            
552
    #[test]
553
1
    fn records_parent_and_child_links() {
554
1
        let arena = ExpressionArena::from_root(eq(int(1), int(2)));
555
1
        let root = arena.root();
556
1
        let children = arena.children(root);
557

            
558
1
        assert_eq!(children.len(), 2);
559
1
        assert_eq!(arena.parent(root), None);
560
1
        assert_eq!(arena.parent(children[0]), Some(root));
561
1
        assert_eq!(arena.parent(children[1]), Some(root));
562
1
    }
563

            
564
    #[test]
565
1
    fn direct_child_discriminants_update_without_losing_duplicate_counts() {
566
1
        let mut arena = ExpressionArena::from_root(root(vec![int(1), int(2)]));
567
1
        let root_id = arena.root();
568
1
        let children = arena.children(root_id).to_vec();
569
1
        let atomic = discriminant_from_value(arena.expression(children[0]));
570
1
        let equality = discriminant_from_value(&eq(int(1), int(2)));
571

            
572
1
        assert!(arena.has_direct_child_discriminant(root_id, atomic));
573
1
        assert!(!arena.has_direct_child_discriminant(root_id, equality));
574

            
575
1
        arena.replace_subtree(children[0], eq(int(3), int(4)));
576
1
        assert!(arena.has_direct_child_discriminant(root_id, atomic));
577
1
        assert!(arena.has_direct_child_discriminant(root_id, equality));
578
1
        assert!(arena.has_direct_child_discriminant(children[0], atomic));
579

            
580
1
        arena.replace_subtree(children[1], eq(int(5), int(6)));
581
1
        assert!(!arena.has_direct_child_discriminant(root_id, atomic));
582
1
        assert!(arena.has_direct_child_discriminant(root_id, equality));
583
1
    }
584

            
585
    #[test]
586
1
    fn replacement_marks_old_descendants_unreachable_and_paths_new_children() {
587
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(3)]));
588
1
        let root_id = arena.root();
589
1
        let first_child = arena.children(root_id)[0];
590
1
        let old_left = arena.children(first_child)[0];
591

            
592
1
        arena.replace_subtree(first_child, eq(int(4), int(5)));
593

            
594
1
        assert!(!arena.is_reachable(old_left));
595
1
        assert_eq!(arena.preorder_path(first_child), &[0]);
596

            
597
1
        let new_children = arena.children(first_child);
598
1
        assert_eq!(arena.preorder_path(new_children[0]), &[0, 0]);
599
1
        assert_eq!(arena.preorder_path(new_children[1]), &[0, 1]);
600
1
        assert!(
601
1
            arena
602
1
                .reachable_subtree_ids(first_child)
603
1
                .contains(&new_children[0])
604
        );
605
1
        assert!(!arena.reachable_subtree_ids(first_child).contains(&old_left));
606
1
    }
607

            
608
    #[test]
609
1
    fn added_root_children_get_preorder_paths() {
610
1
        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
611
1
        let added = arena.add_root_children(vec![int(2), int(3)]);
612

            
613
1
        assert_eq!(added.len(), 2);
614
1
        assert_eq!(arena.preorder_path(added[0]), &[1]);
615
1
        assert_eq!(arena.preorder_path(added[1]), &[2]);
616
1
        assert_eq!(arena.parent(added[0]), Some(arena.root()));
617
1
        assert_eq!(arena.parent(added[1]), Some(arena.root()));
618
1
    }
619

            
620
    #[test]
621
1
    fn obtains_direct_child_indices_from_preorder_paths() {
622
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(3)]));
623
1
        let root_id = arena.root();
624
1
        let eq_id = arena.children(root_id)[0];
625
1
        let eq_children = arena.children(eq_id).to_vec();
626

            
627
1
        assert_eq!(arena.direct_child_index(root_id, eq_id), Some(0));
628
1
        assert_eq!(arena.direct_child_index(eq_id, eq_children[0]), Some(0));
629
1
        assert_eq!(arena.direct_child_index(eq_id, eq_children[1]), Some(1));
630
1
        assert_eq!(arena.direct_child_index(root_id, eq_children[0]), None);
631
1
        assert_eq!(arena.direct_child_index(eq_id, root_id), None);
632

            
633
1
        let added = arena.add_root_children(vec![int(4)]);
634
1
        assert_eq!(arena.direct_child_index(root_id, added[0]), Some(2));
635
1
    }
636

            
637
    #[test]
638
1
    fn direct_child_index_rejects_unreachable_replaced_children() {
639
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2))]));
640
1
        let root_id = arena.root();
641
1
        let eq_id = arena.children(root_id)[0];
642
1
        let old_left = arena.children(eq_id)[0];
643

            
644
1
        arena.replace_subtree(eq_id, eq(int(3), int(4)));
645

            
646
1
        assert!(!arena.is_reachable(old_left));
647
1
        assert_eq!(arena.direct_child_index(eq_id, old_left), None);
648
1
        assert_eq!(
649
1
            arena.direct_child_index(eq_id, arena.children(eq_id)[0]),
650
1
            Some(0)
651
        );
652
1
    }
653

            
654
    #[test]
655
1
    fn adding_root_children_bumps_root_generation() {
656
1
        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
657
1
        let root_id = arena.root();
658
1
        let before = arena.generation(root_id);
659

            
660
1
        arena.add_root_children(vec![int(2)]);
661

            
662
1
        assert_ne!(arena.generation(root_id), before);
663
1
    }
664

            
665
    #[test]
666
1
    fn replaces_subtree_without_changing_replaced_node_id() {
667
1
        let mut arena = ExpressionArena::from_root(eq(int(1), int(2)));
668
1
        let root = arena.root();
669
1
        let left = arena.children(root)[0];
670

            
671
1
        arena.replace_subtree(left, eq(int(3), int(4)));
672

            
673
1
        assert_eq!(arena.children(root)[0], left);
674
1
        assert_eq!(arena.parent(arena.children(left)[0]), Some(left));
675
1
        assert_eq!(arena.into_root_expression(), eq(eq(int(3), int(4)), int(2)));
676
1
    }
677

            
678
    #[test]
679
1
    fn rebuilds_ancestor_payload_after_child_replacement() {
680
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2))]));
681
1
        let root_id = arena.root();
682
1
        let eq_id = arena.children(root_id)[0];
683
1
        let left = arena.children(eq_id)[0];
684
1
        let old_eq_hash = arena.expression(eq_id).cached_content_hash();
685
1
        let old_root_hash = arena.expression(root_id).cached_content_hash();
686

            
687
1
        arena.replace_subtree(left, int(3));
688
1
        arena.rebuild_payload_from_children(eq_id);
689
1
        arena.rebuild_payload_from_children(root_id);
690

            
691
1
        assert_eq!(arena.expression(eq_id), &eq(int(3), int(2)));
692
1
        assert_eq!(arena.expression(root_id), &root(vec![eq(int(3), int(2))]));
693
1
        assert_ne!(arena.expression(eq_id).cached_content_hash(), old_eq_hash);
694
1
        assert_ne!(
695
1
            arena.expression(root_id).cached_content_hash(),
696
1
            old_root_hash
697
        );
698
1
    }
699

            
700
    #[test]
701
1
    fn incremental_rebuild_handles_multiple_batches_and_changed_arity() {
702
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
703
1
        let root_id = arena.root();
704
1
        let eq_id = arena.children(root_id)[0];
705
1
        let children = arena.children(eq_id).to_vec();
706
1
        arena.replace_subtree(children[0], int(3));
707
1
        arena.replace_subtree(children[1], int(4));
708
1
        arena.rebuild_payload_from_children(eq_id);
709
1
        arena.rebuild_payload_from_children(root_id);
710
1
        assert_eq!(
711
1
            arena.expression(root_id),
712
1
            &root(vec![eq(int(3), int(4)), int(9)])
713
        );
714

            
715
1
        arena.replace_subtree(eq_id, root(vec![int(5), int(6), int(7)]));
716
1
        let child = arena.children(eq_id)[2];
717
1
        arena.replace_subtree(child, int(8));
718
1
        arena.rebuild_payload_from_children(eq_id);
719
1
        arena.rebuild_payload_from_children(root_id);
720
1
        assert_eq!(
721
1
            arena.expression(root_id),
722
1
            &root(vec![root(vec![int(5), int(6), int(8)]), int(9)])
723
        );
724

            
725
1
        let added = arena.add_root_children(vec![int(10)]);
726
1
        arena.replace_subtree(added[0], int(11));
727
1
        arena.rebuild_payload_from_children(root_id);
728
1
        assert_eq!(
729
1
            arena.expression(root_id),
730
1
            &root(vec![root(vec![int(5), int(6), int(8)]), int(9), int(11)])
731
        );
732
1
    }
733

            
734
    #[test]
735
1
    fn sync_and_rebuild_agree_after_child_replacement() {
736
1
        let mut sync_arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
737
1
        let mut rebuild_arena = sync_arena.clone();
738

            
739
1
        for arena in [&mut sync_arena, &mut rebuild_arena] {
740
2
            let root_id = arena.root();
741
2
            let eq_id = arena.children(root_id)[0];
742
2
            let left = arena.children(eq_id)[0];
743
2
            arena.replace_subtree(left, int(3));
744
        }
745

            
746
1
        let sync_root = sync_arena.root();
747
1
        let sync_eq = sync_arena.children(sync_root)[0];
748
1
        let sync_left = sync_arena.children(sync_eq)[0];
749
1
        sync_arena.sync_payload_for_changed_child(sync_eq, sync_left);
750
1
        sync_arena.sync_payload_for_changed_child(sync_root, sync_eq);
751

            
752
1
        let rebuild_root = rebuild_arena.root();
753
1
        let rebuild_eq = rebuild_arena.children(rebuild_root)[0];
754
1
        rebuild_arena.rebuild_payload_from_children(rebuild_eq);
755
1
        rebuild_arena.rebuild_payload_from_children(rebuild_root);
756

            
757
1
        assert_eq!(
758
1
            sync_arena.into_root_expression(),
759
1
            rebuild_arena.into_root_expression()
760
        );
761
1
    }
762

            
763
    #[test]
764
1
    fn syncs_same_arity_ancestor_payload_without_full_rebuild() {
765
1
        let mut arena = ExpressionArena::from_root(root(vec![eq(int(1), int(2)), int(9)]));
766
1
        let root_id = arena.root();
767
1
        let eq_id = arena.children(root_id)[0];
768
1
        let left = arena.children(eq_id)[0];
769
1
        let untouched = arena.children(root_id)[1];
770
1
        let untouched_before = arena.expression(untouched).clone();
771

            
772
1
        arena.replace_subtree(left, int(3));
773
1
        arena.sync_payload_for_changed_child(eq_id, left);
774
1
        arena.sync_payload_for_changed_child(root_id, eq_id);
775

            
776
1
        assert_eq!(arena.expression(eq_id), &eq(int(3), int(2)));
777
1
        assert_eq!(
778
1
            arena.expression(root_id),
779
1
            &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
1
        assert_eq!(arena.expression(untouched), &untouched_before);
783
1
        assert_eq!(
784
1
            arena.into_root_expression(),
785
1
            root(vec![eq(int(3), int(2)), int(9)])
786
        );
787
1
    }
788

            
789
    #[test]
790
1
    fn syncs_wide_matrix_child_slot_in_place() {
791
        use crate::ast::{AbstractLiteral, Domain, Range};
792
        use crate::into_matrix_expr;
793

            
794
1
        let elems: Vec<Expression> = (0..32).map(int).collect();
795
1
        let matrix = into_matrix_expr![elems; Domain::int(vec![Range::Bounded(1, 32)])];
796
        let mut arena =
797
1
            ExpressionArena::from_root(Expression::And(Metadata::new(), Moo::new(matrix)));
798
1
        let and_id = arena.root();
799
1
        let matrix_id = arena.children(and_id)[0];
800
1
        let target = arena.children(matrix_id)[17];
801

            
802
1
        arena.replace_subtree(target, int(99));
803
1
        arena.sync_payload_for_changed_child(matrix_id, target);
804
1
        arena.sync_payload_for_changed_child(and_id, matrix_id);
805

            
806
1
        let Expression::And(_, inner) = arena.expression(and_id) else {
807
            panic!("expected And");
808
        };
809
1
        let Expression::AbstractLiteral(_, AbstractLiteral::Matrix(elems, _)) = inner.as_ref()
810
        else {
811
            panic!("expected matrix");
812
        };
813
1
        assert_eq!(elems.len(), 32);
814
1
        assert_eq!(elems[17], int(99));
815
1
        assert_eq!(elems[0], int(0));
816
1
        assert_eq!(elems[31], int(31));
817
1
    }
818

            
819
    #[test]
820
1
    fn appending_constraints_preserves_pending_child_updates() {
821
1
        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
822
1
        let root_id = arena.root();
823
1
        let first = arena.children(root_id)[0];
824
1
        arena.replace_subtree(first, int(2));
825
1
        arena.add_root_children(vec![int(3)]);
826
1
        assert_eq!(arena.expression(root_id), &root(vec![int(2), int(3)]));
827
1
        arena.replace_subtree(first, int(4));
828
1
        arena.rebuild_payload_from_children(root_id);
829
1
        assert_eq!(arena.expression(root_id), &root(vec![int(4), int(3)]));
830
1
    }
831

            
832
    #[test]
833
1
    fn appends_root_children() {
834
1
        let mut arena = ExpressionArena::from_root(root(vec![int(1)]));
835

            
836
1
        arena.add_root_children(vec![int(2), int(3)]);
837

            
838
1
        assert_eq!(
839
1
            arena.into_root_expression(),
840
1
            root(vec![int(1), int(2), int(3)])
841
        );
842
1
    }
843
}