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 let mut symbols: Vec<_> = referenced_symbols.into_iter().collect();
49 symbols.sort_by_key(|b| Reverse(b.1));
50
51 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 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 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}