Skip to main content

rue_hir/
environment.rs

1use std::cmp::Reverse;
2
3use indexmap::IndexMap;
4use num_bigint::BigUint;
5
6use crate::SymbolId;
7
8#[derive(Debug, Default, Clone)]
9#[must_use]
10pub enum Environment {
11    #[default]
12    Nil,
13    Leaf(SymbolId),
14    Pair(Box<Environment>, Box<Environment>),
15}
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum PathError {
19    SymbolNotFound,
20}
21
22impl Environment {
23    pub fn sequential(symbols: &[SymbolId], nil_terminated: bool) -> Self {
24        let mut result = Self::Nil;
25
26        for &symbol in symbols.iter().rev() {
27            if matches!(result, Self::Nil) && !nil_terminated {
28                result = Self::Leaf(symbol);
29                continue;
30            }
31
32            result = Self::Pair(Box::new(Self::Leaf(symbol)), Box::new(result));
33        }
34
35        result
36    }
37
38    pub fn tree(referenced_symbols: IndexMap<SymbolId, usize>) -> Self {
39        if referenced_symbols.is_empty() {
40            return Self::Nil;
41        }
42
43        if referenced_symbols.len() == 1 {
44            return Self::Leaf(*referenced_symbols.keys().next().unwrap());
45        }
46
47        // Sort symbols by reference count in descending order
48        let mut symbols: Vec<_> = referenced_symbols.into_iter().collect();
49        symbols.sort_by_key(|b| Reverse(b.1));
50
51        // Split symbols into two groups, trying to balance:
52        // 1. Total reference counts between groups
53        // 2. More frequently referenced symbols in the first group (shorter paths)
54        let mut split_idx = symbols.len().div_ceil(2);
55        let mut min_diff = i64::MAX;
56
57        for i in 1..symbols.len() {
58            let mut new_first_sum = 0usize;
59            let mut new_rest_sum = 0usize;
60
61            for (j, (_, count)) in symbols.iter().enumerate() {
62                if j < i {
63                    new_first_sum += *count;
64                } else {
65                    new_rest_sum += *count;
66                }
67            }
68
69            #[allow(clippy::cast_possible_wrap)]
70            let diff = (new_first_sum as i64 - new_rest_sum as i64).abs();
71            if diff < min_diff {
72                min_diff = diff;
73                split_idx = i;
74            }
75        }
76
77        // Create first and rest groups
78        let first_group: IndexMap<_, _> = symbols
79            .iter()
80            .take(split_idx)
81            .map(|(k, v)| (*k, *v))
82            .collect();
83        let rest_group: IndexMap<_, _> = symbols
84            .iter()
85            .skip(split_idx)
86            .map(|(k, v)| (*k, *v))
87            .collect();
88
89        // Recursively build the tree
90        Self::Pair(
91            Box::new(Self::tree(first_group)),
92            Box::new(Self::tree(rest_group)),
93        )
94    }
95
96    pub fn path(&self, symbol_id: SymbolId) -> Result<BigUint, PathError> {
97        let ops = self
98            .get_pair_ops(symbol_id)
99            .ok_or(PathError::SymbolNotFound)?;
100        let mut path = BigUint::from(1u8);
101        for op in ops.into_iter().rev() {
102            path <<= 1;
103            if matches!(op, PairOp::Rest) {
104                path += 1u8;
105            }
106        }
107        Ok(path)
108    }
109
110    fn get_pair_ops(&self, symbol_id: SymbolId) -> Option<Vec<PairOp>> {
111        match self {
112            Self::Nil => None,
113            Self::Leaf(leaf) => {
114                if *leaf == symbol_id {
115                    Some(vec![])
116                } else {
117                    None
118                }
119            }
120            Self::Pair(first, rest) => {
121                if let Some(mut ops) = first.get_pair_ops(symbol_id) {
122                    ops.insert(0, PairOp::First);
123                    Some(ops)
124                } else if let Some(mut ops) = rest.get_pair_ops(symbol_id) {
125                    ops.insert(0, PairOp::Rest);
126                    Some(ops)
127                } else {
128                    None
129                }
130            }
131        }
132    }
133}
134
135#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
136pub enum PairOp {
137    First,
138    Rest,
139}