Skip to main content

formualizer_eval/engine/arena/
ast.rs

1/// AST arena with structural sharing and deduplication
2/// Stores formula AST nodes efficiently with content-addressable storage
3use super::string_interner::{StringId, StringInterner};
4use super::value_ref::ValueRef;
5use formualizer_parse::parser::{ExternalRefKind, TableSpecifier};
6use rustc_hash::FxHashMap;
7use std::collections::hash_map::DefaultHasher;
8use std::fmt;
9use std::hash::{Hash, Hasher};
10
11/// Reference to an AST node in the arena
12#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
13pub struct AstNodeId(u32);
14
15impl AstNodeId {
16    pub fn as_u32(self) -> u32 {
17        self.0
18    }
19
20    pub(crate) const fn from_u32(raw: u32) -> Self {
21        Self(raw)
22    }
23}
24
25impl fmt::Display for AstNodeId {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        write!(f, "AstNode({})", self.0)
28    }
29}
30
31#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
32pub struct TableSpecId(u32);
33
34impl TableSpecId {
35    pub fn as_u32(self) -> u32 {
36        self.0
37    }
38}
39
40/// Compact representation of AST nodes in the arena
41#[derive(Debug, Clone, PartialEq, Eq, Hash)]
42pub enum AstNodeData {
43    /// Literal value
44    Literal(ValueRef),
45
46    /// Explicitly omitted function argument.
47    Omitted,
48
49    /// Cell or range reference
50    Reference {
51        original_id: StringId,    // Original reference string
52        ref_type: CompactRefType, // Compact reference representation
53    },
54
55    /// Unary operation
56    UnaryOp { op_id: StringId, expr_id: AstNodeId },
57
58    /// Binary operation
59    BinaryOp {
60        op_id: StringId,
61        left_id: AstNodeId,
62        right_id: AstNodeId,
63    },
64
65    /// Function call
66    Function {
67        name_id: StringId,
68        args_offset: u32, // Index into args array
69        args_count: u16,  // Number of arguments
70    },
71
72    /// Array literal
73    Array {
74        rows: u16,
75        cols: u16,
76        elements_offset: u32, // Index into elements array
77    },
78}
79
80/// Identifies a sheet either by stable registry id or by unresolved name.
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub enum SheetKey {
83    Id(u16),
84    Name(StringId),
85}
86
87/// Compact representation of reference types
88#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
89pub enum CompactRefType {
90    Cell {
91        sheet: Option<SheetKey>,
92        row: u32,
93        col: u32,
94        row_abs: bool,
95        col_abs: bool,
96    },
97    Range {
98        sheet: Option<SheetKey>,
99        start_row: u32,
100        start_col: u32,
101        end_row: u32,
102        end_col: u32,
103        start_row_abs: bool,
104        start_col_abs: bool,
105        end_row_abs: bool,
106        end_col_abs: bool,
107    },
108    External {
109        raw_id: StringId,
110        book_id: StringId,
111        sheet_id: StringId,
112        kind: ExternalRefKind,
113    },
114    NamedRange(StringId),
115    Table {
116        name_id: StringId,
117        specifier_id: Option<TableSpecId>,
118    },
119    /// 3D cell reference (`Sheet1:Sheet3!A1`).
120    Cell3D {
121        sheet_first: StringId,
122        sheet_last: StringId,
123        row: u32,
124        col: u32,
125        row_abs: bool,
126        col_abs: bool,
127    },
128    /// 3D range reference (`Sheet1:Sheet3!A1:B2`).
129    Range3D {
130        sheet_first: StringId,
131        sheet_last: StringId,
132        start_row: u32,
133        start_col: u32,
134        end_row: u32,
135        end_col: u32,
136        start_row_abs: bool,
137        start_col_abs: bool,
138        end_row_abs: bool,
139        end_col_abs: bool,
140    },
141}
142
143/// Arena entry containing structural node data plus canonical metadata.
144///
145/// Phase 1 keeps metadata at its default value for existing raw interning
146/// paths. Future canonical interning paths populate `meta` via the arena
147/// canonicalization helpers.
148#[derive(Debug, Clone, PartialEq, Eq, Hash)]
149pub(crate) struct AstNodeEntry {
150    pub(crate) data: AstNodeData,
151    pub(crate) meta: AstNodeMetadata,
152}
153
154/// Canonical metadata associated with an arena AST node.
155#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
156pub(crate) struct AstNodeMetadata {
157    pub(crate) canonical_hash: u64,
158    pub(crate) labels: CanonicalLabels,
159    pub(crate) reference_returning_admission: ReferenceReturningAdmission,
160}
161
162#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
163pub(crate) struct ReferenceReturningAdmission(u8);
164
165impl ReferenceReturningAdmission {
166    pub(crate) const fn new(safe: bool, scalar: bool) -> Self {
167        Self((safe as u8) | ((scalar as u8) << 1))
168    }
169
170    pub(crate) const fn safe(self) -> bool {
171        self.0 & 1 != 0
172    }
173
174    pub(crate) const fn scalar(self) -> bool {
175        self.0 & 2 != 0
176    }
177}
178
179/// Compact bitset labels for arena-native canonicalization.
180#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
181pub(crate) struct CanonicalLabels {
182    pub(crate) flags: u64,
183    pub(crate) rejects: u64,
184}
185
186#[allow(dead_code)]
187impl CanonicalLabels {
188    pub(crate) const FLAG_RELATIVE_ONLY: u64 = 1 << 0;
189    pub(crate) const FLAG_ABSOLUTE_ONLY: u64 = 1 << 1;
190    pub(crate) const FLAG_MIXED_ANCHORS: u64 = 1 << 2;
191    pub(crate) const FLAG_VOLATILE: u64 = 1 << 3;
192    pub(crate) const FLAG_DYNAMIC: u64 = 1 << 4;
193    pub(crate) const FLAG_CONTAINS_STRUCTURED_REF: u64 = 1 << 5;
194    pub(crate) const FLAG_NEEDS_PLACEMENT_REWRITE: u64 = 1 << 6;
195    pub(crate) const FLAG_CONTAINS_NAME: u64 = 1 << 7;
196    pub(crate) const FLAG_CONTAINS_TABLE: u64 = 1 << 8;
197    pub(crate) const FLAG_CONTAINS_RANGE: u64 = 1 << 9;
198    pub(crate) const FLAG_CONTAINS_ARRAY: u64 = 1 << 10;
199    pub(crate) const FLAG_CONTAINS_LET_LAMBDA: u64 = 1 << 11;
200    pub(crate) const FLAG_CONTAINS_FUNCTION: u64 = 1 << 12;
201    pub(crate) const FLAG_EXPLICIT_SHEET: u64 = 1 << 13;
202    pub(crate) const FLAG_CURRENT_SHEET: u64 = 1 << 14;
203
204    // Reject bits mirror `CanonicalRejectReason` variants in
205    // `formula_plane/template_canonical.rs` by kind/variant order.
206    pub(crate) const REJECT_INVALID_PLACEMENT_ANCHOR: u64 = 1 << 0;
207    pub(crate) const REJECT_DYNAMIC_REFERENCE: u64 = 1 << 1;
208    pub(crate) const REJECT_UNKNOWN_OR_CUSTOM_FUNCTION: u64 = 1 << 2;
209    pub(crate) const REJECT_LOCAL_ENVIRONMENT: u64 = 1 << 3;
210    pub(crate) const REJECT_PARSER_VOLATILE_FLAG: u64 = 1 << 4;
211    pub(crate) const REJECT_VOLATILE_FUNCTION: u64 = 1 << 5;
212    pub(crate) const REJECT_REFERENCE_RETURNING_FUNCTION: u64 = 1 << 6;
213    pub(crate) const REJECT_ARRAY_OR_SPILL_FUNCTION: u64 = 1 << 7;
214    pub(crate) const REJECT_ARRAY_LITERAL: u64 = 1 << 8;
215    pub(crate) const REJECT_SPILL_REFERENCE: u64 = 1 << 9;
216    pub(crate) const REJECT_SPILL_RESULT_REGION_OPERATOR: u64 = 1 << 10;
217    pub(crate) const REJECT_IMPLICIT_INTERSECTION_OPERATOR: u64 = 1 << 11;
218    pub(crate) const REJECT_CALL_EXPRESSION: u64 = 1 << 12;
219    // Bit 13 (REJECT_NAMED_REFERENCE) retired: named references canonicalize
220    // by identity and are accepted/rejected at read-projection time instead.
221    pub(crate) const REJECT_STRUCTURED_REFERENCE: u64 = 1 << 14;
222    pub(crate) const REJECT_STRUCTURED_REFERENCE_CURRENT_ROW: u64 = 1 << 15;
223    pub(crate) const REJECT_THREE_D_REFERENCE: u64 = 1 << 16;
224    pub(crate) const REJECT_EXTERNAL_REFERENCE: u64 = 1 << 17;
225    pub(crate) const REJECT_OPEN_RANGE_REFERENCE: u64 = 1 << 18;
226    pub(crate) const REJECT_WHOLE_AXIS_REFERENCE: u64 = 1 << 19;
227    pub(crate) const REJECT_UNSUPPORTED_REFERENCE: u64 = 1 << 20;
228
229    pub(crate) fn has_flag(self, flag: u64) -> bool {
230        self.flags & flag != 0
231    }
232
233    pub(crate) fn has_reject(self, reject: u64) -> bool {
234        self.rejects & reject != 0
235    }
236}
237
238/// Arena for storing AST nodes with deduplication
239pub struct AstArena {
240    /// Node storage
241    nodes: Vec<AstNodeEntry>,
242
243    /// Hash -> node index for deduplication
244    dedup_map: FxHashMap<u64, AstNodeId>,
245
246    /// Function arguments storage (flattened)
247    function_args: Vec<AstNodeId>,
248
249    /// Array elements storage (flattened)
250    array_elements: Vec<AstNodeId>,
251
252    /// String pool for operators and function names
253    strings: StringInterner,
254
255    /// Structured table specifiers
256    table_specs: Vec<TableSpecifier>,
257    table_spec_dedup: FxHashMap<u64, TableSpecId>,
258
259    /// Statistics
260    dedup_hits: usize,
261}
262
263impl AstArena {
264    pub fn new() -> Self {
265        Self {
266            nodes: Vec::new(),
267            dedup_map: FxHashMap::default(),
268            function_args: Vec::new(),
269            array_elements: Vec::new(),
270            strings: StringInterner::new(),
271            table_specs: Vec::new(),
272            table_spec_dedup: FxHashMap::default(),
273            dedup_hits: 0,
274        }
275    }
276
277    pub fn with_capacity(node_cap: usize) -> Self {
278        Self {
279            nodes: Vec::with_capacity(node_cap),
280            dedup_map: FxHashMap::with_capacity_and_hasher(node_cap, Default::default()),
281            function_args: Vec::with_capacity(node_cap * 2), // Assume avg 2 args
282            array_elements: Vec::with_capacity(node_cap),
283            strings: StringInterner::with_capacity(node_cap / 10),
284            table_specs: Vec::new(),
285            table_spec_dedup: FxHashMap::default(),
286            dedup_hits: 0,
287        }
288    }
289
290    /// Insert a node, deduplicating if it already exists.
291    ///
292    /// Phase 1 preserves raw/literal interning semantics: metadata is filled
293    /// with zeros and is not part of the dedup key.
294    pub fn insert(&mut self, node: AstNodeData) -> AstNodeId {
295        self.insert_entry(node, AstNodeMetadata::default())
296    }
297
298    /// Insert a node with precomputed canonical metadata.
299    ///
300    /// This is unused in Phase 1 and reserved for the future canonical
301    /// interning path. Deduplication remains structural on `AstNodeData` only;
302    /// if a matching node already exists, the existing entry (and metadata)
303    /// wins.
304    #[allow(dead_code)]
305    pub(crate) fn insert_with_meta(
306        &mut self,
307        node: AstNodeData,
308        meta: AstNodeMetadata,
309    ) -> AstNodeId {
310        self.insert_entry(node, meta)
311    }
312
313    fn insert_entry(&mut self, node: AstNodeData, meta: AstNodeMetadata) -> AstNodeId {
314        // Compute structural hash. Metadata is deliberately excluded.
315        let hash = self.hash_node(&node);
316
317        // Check for existing node
318        if let Some(&id) = self.dedup_map.get(&hash) {
319            // Verify it's actually the same (handle hash collisions)
320            if self.nodes[id.0 as usize].data == node {
321                self.dedup_hits += 1;
322                return id;
323            }
324        }
325
326        // Add new node
327        let id = AstNodeId(self.nodes.len() as u32);
328        self.nodes.push(AstNodeEntry { data: node, meta });
329        self.dedup_map.insert(hash, id);
330        id
331    }
332
333    /// Insert a literal node
334    pub fn insert_literal(&mut self, value: ValueRef) -> AstNodeId {
335        self.insert(AstNodeData::Literal(value))
336    }
337
338    /// Insert an explicitly omitted argument node.
339    pub(crate) fn insert_omitted(&mut self) -> AstNodeId {
340        self.insert(AstNodeData::Omitted)
341    }
342
343    /// Insert a reference node
344    pub fn insert_reference(&mut self, original: &str, ref_type: CompactRefType) -> AstNodeId {
345        let original_id = self.strings.intern(original);
346        self.insert(AstNodeData::Reference {
347            original_id,
348            ref_type,
349        })
350    }
351
352    /// Insert a unary operation node
353    pub fn insert_unary_op(&mut self, op: &str, expr: AstNodeId) -> AstNodeId {
354        let op_id = self.strings.intern(op);
355        self.insert(AstNodeData::UnaryOp {
356            op_id,
357            expr_id: expr,
358        })
359    }
360
361    /// Insert a binary operation node
362    pub fn insert_binary_op(&mut self, op: &str, left: AstNodeId, right: AstNodeId) -> AstNodeId {
363        let op_id = self.strings.intern(op);
364        self.insert(AstNodeData::BinaryOp {
365            op_id,
366            left_id: left,
367            right_id: right,
368        })
369    }
370
371    /// Insert a function call node
372    pub fn insert_function(&mut self, name: &str, args: Vec<AstNodeId>) -> AstNodeId {
373        let name_id = self.strings.intern(name);
374        let args_offset = self.function_args.len() as u32;
375        let args_count = args.len() as u16;
376
377        self.function_args.extend(args);
378
379        self.insert(AstNodeData::Function {
380            name_id,
381            args_offset,
382            args_count,
383        })
384    }
385
386    /// Insert an array literal node
387    pub fn insert_array(&mut self, rows: u16, cols: u16, elements: Vec<AstNodeId>) -> AstNodeId {
388        assert_eq!(
389            elements.len(),
390            (rows * cols) as usize,
391            "Array dimensions don't match element count"
392        );
393
394        let elements_offset = self.array_elements.len() as u32;
395        self.array_elements.extend(elements);
396
397        self.insert(AstNodeData::Array {
398            rows,
399            cols,
400            elements_offset,
401        })
402    }
403
404    /// Get a node by ID
405    pub fn get(&self, id: AstNodeId) -> Option<&AstNodeData> {
406        self.nodes.get(id.0 as usize).map(|entry| &entry.data)
407    }
408
409    /// Get an arena entry by ID.
410    #[allow(dead_code)]
411    pub(crate) fn entry(&self, id: AstNodeId) -> Option<&AstNodeEntry> {
412        self.nodes.get(id.0 as usize)
413    }
414
415    /// Get canonical metadata for a node by ID.
416    #[allow(dead_code)]
417    pub(crate) fn metadata(&self, id: AstNodeId) -> Option<AstNodeMetadata> {
418        self.entry(id).map(|entry| entry.meta)
419    }
420
421    /// Get function arguments for a function node
422    pub fn get_function_args(&self, id: AstNodeId) -> Option<&[AstNodeId]> {
423        match self.get(id)? {
424            AstNodeData::Function {
425                args_offset,
426                args_count,
427                ..
428            } => {
429                let start = *args_offset as usize;
430                let end = start + *args_count as usize;
431                Some(&self.function_args[start..end])
432            }
433            _ => None,
434        }
435    }
436
437    /// Get array elements for an array node
438    pub fn get_array_elements(&self, id: AstNodeId) -> Option<&[AstNodeId]> {
439        match self.get(id)? {
440            AstNodeData::Array {
441                rows,
442                cols,
443                elements_offset,
444            } => {
445                let start = *elements_offset as usize;
446                let count = (*rows * *cols) as usize;
447                let end = start + count;
448                Some(&self.array_elements[start..end])
449            }
450            _ => None,
451        }
452    }
453
454    pub fn get_array_elements_info(&self, id: AstNodeId) -> Option<(u16, u16, &[AstNodeId])> {
455        match self.get(id)? {
456            AstNodeData::Array { rows, cols, .. } => {
457                let elements = self.get_array_elements(id)?;
458                Some((*rows, *cols, elements))
459            }
460            _ => None,
461        }
462    }
463
464    /// Resolve a string ID to its content
465    pub fn resolve_string(&self, id: StringId) -> &str {
466        self.strings.resolve(id)
467    }
468
469    /// Get the string interner (for external use)
470    pub fn strings(&self) -> &StringInterner {
471        &self.strings
472    }
473
474    /// Get mutable access to the string interner
475    pub fn strings_mut(&mut self) -> &mut StringInterner {
476        &mut self.strings
477    }
478
479    pub fn intern_table_specifier(&mut self, specifier: &TableSpecifier) -> TableSpecId {
480        let hash = {
481            let mut hasher = DefaultHasher::new();
482            specifier.hash(&mut hasher);
483            hasher.finish()
484        };
485
486        if let Some(&id) = self.table_spec_dedup.get(&hash)
487            && self
488                .table_specs
489                .get(id.0 as usize)
490                .is_some_and(|existing| existing == specifier)
491        {
492            return id;
493        }
494
495        let id = TableSpecId(self.table_specs.len() as u32);
496        self.table_specs.push(specifier.clone());
497        self.table_spec_dedup.insert(hash, id);
498        id
499    }
500
501    pub fn resolve_table_specifier(&self, id: TableSpecId) -> Option<&TableSpecifier> {
502        self.table_specs.get(id.0 as usize)
503    }
504
505    /// Compute hash for a node
506    fn hash_node(&self, node: &AstNodeData) -> u64 {
507        let mut hasher = DefaultHasher::new();
508        node.hash(&mut hasher);
509        hasher.finish()
510    }
511
512    /// Get statistics about the arena
513    pub fn stats(&self) -> AstArenaStats {
514        AstArenaStats {
515            node_count: self.nodes.len(),
516            dedup_hits: self.dedup_hits,
517            string_count: self.strings.len(),
518            table_spec_count: self.table_specs.len(),
519            total_args: self.function_args.len(),
520            total_array_elements: self.array_elements.len(),
521        }
522    }
523
524    /// Returns memory usage in bytes (approximate)
525    pub fn memory_usage(&self) -> usize {
526        self.nodes.capacity() * std::mem::size_of::<AstNodeEntry>()
527            + self.dedup_map.capacity() * (8 + 4) // hash + id
528            + self.function_args.capacity() * 4
529            + self.array_elements.capacity() * 4
530            + self.strings.memory_usage()
531            + self.table_specs.capacity() * std::mem::size_of::<TableSpecifier>()
532            + self.table_spec_dedup.capacity() * (8 + 4)
533    }
534
535    /// Clear all nodes from the arena
536    pub fn clear(&mut self) {
537        self.nodes.clear();
538        self.dedup_map.clear();
539        self.function_args.clear();
540        self.array_elements.clear();
541        self.strings.clear();
542        self.table_specs.clear();
543        self.table_spec_dedup.clear();
544        self.dedup_hits = 0;
545    }
546}
547
548impl Default for AstArena {
549    fn default() -> Self {
550        Self::new()
551    }
552}
553
554impl fmt::Debug for AstArena {
555    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
556        f.debug_struct("AstArena")
557            .field("nodes", &self.nodes.len())
558            .field("dedup_hits", &self.dedup_hits)
559            .field("strings", &self.strings.len())
560            .finish()
561    }
562}
563
564/// Statistics about the AST arena
565#[derive(Debug, Clone)]
566pub struct AstArenaStats {
567    pub node_count: usize,
568    pub dedup_hits: usize,
569    pub string_count: usize,
570    pub table_spec_count: usize,
571    pub total_args: usize,
572    pub total_array_elements: usize,
573}
574
575#[cfg(test)]
576mod tests {
577    use super::*;
578
579    #[test]
580    fn test_ast_arena_literal() {
581        let mut arena = AstArena::new();
582
583        let lit1 = arena.insert_literal(ValueRef::small_int(42).unwrap());
584        let lit2 = arena.insert_literal(ValueRef::boolean(true));
585
586        assert_ne!(lit1, lit2);
587
588        match arena.get(lit1) {
589            Some(AstNodeData::Literal(v)) => {
590                assert_eq!(v.as_small_int(), Some(42));
591            }
592            _ => panic!("Expected literal node"),
593        }
594    }
595
596    #[test]
597    fn test_ast_arena_deduplication() {
598        let mut arena = AstArena::new();
599
600        // Insert same literal twice
601        let lit1 = arena.insert_literal(ValueRef::small_int(42).unwrap());
602        let lit2 = arena.insert_literal(ValueRef::small_int(42).unwrap());
603
604        assert_eq!(lit1, lit2); // Should be deduplicated
605        assert_eq!(arena.stats().dedup_hits, 1);
606    }
607
608    #[test]
609    fn test_ast_arena_binary_op() {
610        let mut arena = AstArena::new();
611
612        let left = arena.insert_literal(ValueRef::small_int(1).unwrap());
613        let right = arena.insert_literal(ValueRef::small_int(2).unwrap());
614        let add = arena.insert_binary_op("+", left, right);
615
616        match arena.get(add) {
617            Some(AstNodeData::BinaryOp {
618                op_id,
619                left_id,
620                right_id,
621            }) => {
622                assert_eq!(arena.resolve_string(*op_id), "+");
623                assert_eq!(*left_id, left);
624                assert_eq!(*right_id, right);
625            }
626            _ => panic!("Expected binary op node"),
627        }
628    }
629
630    #[test]
631    fn test_ast_arena_function() {
632        let mut arena = AstArena::new();
633
634        let arg1 = arena.insert_literal(ValueRef::small_int(10).unwrap());
635        let arg2 = arena.insert_literal(ValueRef::small_int(20).unwrap());
636        let arg3 = arena.insert_literal(ValueRef::small_int(30).unwrap());
637
638        let func = arena.insert_function("SUM", vec![arg1, arg2, arg3]);
639
640        match arena.get(func) {
641            Some(AstNodeData::Function {
642                name_id,
643                args_count,
644                ..
645            }) => {
646                assert_eq!(arena.resolve_string(*name_id), "SUM");
647                assert_eq!(*args_count, 3);
648            }
649            _ => panic!("Expected function node"),
650        }
651
652        let args = arena.get_function_args(func).unwrap();
653        assert_eq!(args, &[arg1, arg2, arg3]);
654    }
655
656    #[test]
657    fn test_ast_arena_structural_sharing() {
658        let mut arena = AstArena::new();
659
660        // Create "A1" reference that will be shared
661        let a1_ref = arena.insert_reference(
662            "A1",
663            CompactRefType::Cell {
664                sheet: None,
665                row: 1,
666                col: 1,
667                row_abs: false,
668                col_abs: false,
669            },
670        );
671
672        // Create "A1 + 1"
673        let one = arena.insert_literal(ValueRef::small_int(1).unwrap());
674        let expr1 = arena.insert_binary_op("+", a1_ref, one);
675
676        // Create "A1 * 2"
677        let two = arena.insert_literal(ValueRef::small_int(2).unwrap());
678        let expr2 = arena.insert_binary_op("*", a1_ref, two);
679
680        // A1 reference should be shared
681        assert_eq!(arena.stats().node_count, 5); // A1, 1, +expr, 2, *expr
682
683        // Try to insert A1 again - should be deduplicated
684        let a1_ref2 = arena.insert_reference(
685            "A1",
686            CompactRefType::Cell {
687                sheet: None,
688                row: 1,
689                col: 1,
690                row_abs: false,
691                col_abs: false,
692            },
693        );
694        assert_eq!(a1_ref, a1_ref2);
695    }
696
697    #[test]
698    fn test_ast_arena_array() {
699        let mut arena = AstArena::new();
700
701        let elements = vec![
702            arena.insert_literal(ValueRef::small_int(1).unwrap()),
703            arena.insert_literal(ValueRef::small_int(2).unwrap()),
704            arena.insert_literal(ValueRef::small_int(3).unwrap()),
705            arena.insert_literal(ValueRef::small_int(4).unwrap()),
706        ];
707
708        let array = arena.insert_array(2, 2, elements.clone());
709
710        match arena.get(array) {
711            Some(AstNodeData::Array { rows, cols, .. }) => {
712                assert_eq!(*rows, 2);
713                assert_eq!(*cols, 2);
714            }
715            _ => panic!("Expected array node"),
716        }
717
718        let stored_elements = arena.get_array_elements(array).unwrap();
719        assert_eq!(stored_elements, &elements[..]);
720    }
721
722    #[test]
723    fn test_ast_arena_complex_expression() {
724        let mut arena = AstArena::new();
725
726        // Build: SUM(A1:A10) + IF(B1 > 0, C1, D1)
727
728        // A1:A10 range
729        let range = arena.insert_reference(
730            "A1:A10",
731            CompactRefType::Range {
732                sheet: None,
733                start_row: 1,
734                start_col: 1,
735                end_row: 10,
736                end_col: 1,
737                start_row_abs: false,
738                start_col_abs: false,
739                end_row_abs: false,
740                end_col_abs: false,
741            },
742        );
743
744        // SUM(A1:A10)
745        let sum = arena.insert_function("SUM", vec![range]);
746
747        // B1 reference
748        let b1 = arena.insert_reference(
749            "B1",
750            CompactRefType::Cell {
751                sheet: None,
752                row: 1,
753                col: 2,
754                row_abs: false,
755                col_abs: false,
756            },
757        );
758
759        // 0 literal
760        let zero = arena.insert_literal(ValueRef::small_int(0).unwrap());
761
762        // B1 > 0
763        let condition = arena.insert_binary_op(">", b1, zero);
764
765        // C1 and D1 references
766        let c1 = arena.insert_reference(
767            "C1",
768            CompactRefType::Cell {
769                sheet: None,
770                row: 1,
771                col: 3,
772                row_abs: false,
773                col_abs: false,
774            },
775        );
776        let d1 = arena.insert_reference(
777            "D1",
778            CompactRefType::Cell {
779                sheet: None,
780                row: 1,
781                col: 4,
782                row_abs: false,
783                col_abs: false,
784            },
785        );
786
787        // IF(B1 > 0, C1, D1)
788        let if_expr = arena.insert_function("IF", vec![condition, c1, d1]);
789
790        // Final: SUM(...) + IF(...)
791        let final_expr = arena.insert_binary_op("+", sum, if_expr);
792
793        // Verify structure
794        assert!(arena.get(final_expr).is_some());
795        // Note: zero literal gets deduplicated if used multiple times
796        // We have: range, sum, b1, zero, condition(>), c1, d1, if_expr, final_expr(+)
797        // That's 9 unique nodes (zero is deduplicated)
798        assert_eq!(arena.stats().node_count, 9); // All unique nodes except deduplicated zero
799    }
800
801    #[test]
802    fn test_ast_arena_string_deduplication() {
803        let mut arena = AstArena::new();
804
805        // Use same operator multiple times
806        let one = arena.insert_literal(ValueRef::small_int(1).unwrap());
807        let two = arena.insert_literal(ValueRef::small_int(2).unwrap());
808        let three = arena.insert_literal(ValueRef::small_int(3).unwrap());
809
810        let add1 = arena.insert_binary_op("+", one, two);
811        let add2 = arena.insert_binary_op("+", two, three);
812        let add3 = arena.insert_binary_op("+", one, three);
813
814        // "+" should be interned only once
815        assert_eq!(arena.strings().len(), 1);
816    }
817
818    #[test]
819    fn test_ast_arena_clear() {
820        let mut arena = AstArena::new();
821
822        arena.insert_literal(ValueRef::small_int(1).unwrap());
823        arena.insert_literal(ValueRef::small_int(2).unwrap());
824        let left = arena.insert_literal(ValueRef::small_int(3).unwrap());
825        let right = arena.insert_literal(ValueRef::small_int(4).unwrap());
826        arena.insert_binary_op("+", left, right);
827
828        assert_eq!(arena.stats().node_count, 5);
829
830        arena.clear();
831
832        assert_eq!(arena.stats().node_count, 0);
833        assert_eq!(arena.strings().len(), 0);
834    }
835}