prima_core/
expr_pool.rs

1use dashmap::DashMap;
2use num_bigint::BigInt;
3use num_rational::BigRational;
4use num_traits::One;
5use std::cell::RefCell;
6use std::collections::hash_map::DefaultHasher;
7use std::collections::HashMap;
8use std::hash::{Hash, Hasher};
9use std::sync::{Mutex, OnceLock, RwLock};
10
11use crate::number::{Number, Real};
12use crate::symbol::SymbolId;
13use crate::value::IndeterminateForm;
14
15/// Handle to an expression in the symbolic world (spec §8.1). In-process hash-consing depends on
16/// creation order, so `ExprId` is **forbidden from cross-process serialization/caching** (ADR §6).
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
18pub struct ExprId(u32);
19
20impl ExprId {
21    /// Raw index into the expression store (spec §8.1). Exposed for hashable `ValueKey` construction
22    /// (spec §11.6); the index is process-local and must not cross process boundaries (ADR §6).
23    pub fn as_u32(self) -> u32 {
24        self.0
25    }
26
27    /// Reconstruct an `ExprId` from a raw index; the index must originate from the same process pool (spec §8.1).
28    pub fn from_u32(u: u32) -> ExprId {
29        ExprId(u)
30    }
31}
32
33/// Node in the symbolic world (spec §8.1). `Add`/`Mul` are stored as canonically ordered n-ary lists (spec §8.4),
34/// so equality is `ExprId` equality (O(1)).
35#[derive(Debug, Clone, PartialEq, Hash)]
36pub enum ExprData {
37    Symbol(SymbolId),
38    Integer(Box<BigInt>),
39    Rational(Box<BigRational>),
40    Real(Real),
41    Add(Box<[ExprId]>),
42    Mul(Box<[ExprId]>),
43    Pow { base: ExprId, exp: ExprId },
44    Apply { f: ExprId, args: Box<[ExprId]> },
45    Indeterminate(IndeterminateForm),
46}
47
48// Thread-local cache (spec §8.1): hit the local cache first, fall back to the global pool, write back on a hit.
49thread_local! {
50    static LOCAL_CACHE: RefCell<HashMap<u64, ExprId>> = RefCell::new(HashMap::new());
51}
52
53/// Process-wide shared hash-consing pool (spec §8.1/§12.4): maps content hash → `ExprId`.
54/// The central store is append-only (the symbolic layer is acyclic and resident), concurrency-safe.
55pub struct ExprPool {
56    global: DashMap<u64, ExprId>,
57    store: RwLock<Vec<ExprData>>,
58    alloc: Mutex<()>,
59}
60
61impl ExprPool {
62    pub fn new() -> ExprPool {
63        ExprPool {
64            global: DashMap::new(),
65            store: RwLock::new(Vec::new()),
66            alloc: Mutex::new(()),
67        }
68    }
69
70    /// Process-wide shared instance (`OnceLock`): the interpreter and the symbolic engine share one pool.
71    pub fn global() -> &'static ExprPool {
72        static POOL: OnceLock<ExprPool> = OnceLock::new();
73        POOL.get_or_init(ExprPool::new)
74    }
75
76    fn hash_data(data: &ExprData) -> u64 {
77        let mut h = DefaultHasher::new();
78        data.hash(&mut h);
79        h.finish()
80    }
81
82    /// Intern flow (spec §8.1): content hash → local cache → global pool → append-allocate and write
83    /// back to both caches. The same `ExprData` always yields the same `ExprId`.
84    pub fn intern(&self, data: ExprData) -> ExprId {
85        let key = Self::hash_data(&data);
86        let cached = LOCAL_CACHE.with(|c| c.borrow().get(&key).copied());
87        if let Some(id) = cached {
88            return id;
89        }
90        if let Some(id) = self.global.get(&key) {
91            let id = *id;
92            LOCAL_CACHE.with(|c| c.borrow_mut().insert(key, id));
93            return id;
94        }
95        let _guard = self.alloc.lock().unwrap();
96        if let Some(id) = self.global.get(&key) {
97            let id = *id;
98            LOCAL_CACHE.with(|c| c.borrow_mut().insert(key, id));
99            return id;
100        }
101        let mut store = self.store.write().unwrap();
102        let id = ExprId(store.len() as u32);
103        store.push(data);
104        self.global.insert(key, id);
105        LOCAL_CACHE.with(|c| c.borrow_mut().insert(key, id));
106        id
107    }
108
109    pub fn get(&self, id: ExprId) -> Option<ExprData> {
110        self.store.read().unwrap().get(id.0 as usize).cloned()
111    }
112
113    pub fn symbol(&self, id: SymbolId) -> ExprId {
114        self.intern(ExprData::Symbol(id))
115    }
116
117    pub fn integer(&self, n: i64) -> ExprId {
118        self.intern(ExprData::Integer(Box::new(BigInt::from(n))))
119    }
120
121    pub fn real(&self, x: f64) -> ExprId {
122        self.intern(ExprData::Real(Real::F64(x)))
123    }
124
125    pub fn number(&self, n: &Number) -> ExprId {
126        match n {
127            Number::Integer(i) => self.intern(ExprData::Integer(Box::new(i.clone()))),
128            Number::Rational(r) => {
129                if *r.denom() == BigInt::one() {
130                    self.intern(ExprData::Integer(Box::new(r.numer().clone())))
131                } else {
132                    self.intern(ExprData::Rational(Box::new(r.clone())))
133                }
134            }
135            Number::Real(r) => self.intern(ExprData::Real(*r)),
136            Number::Complex { .. } => panic!("complex numbers cannot be interned as expression nodes yet"),
137            // Fixed-width collapsed layer interns to the exact/`Real` node (spec §6.1).
138            Number::I8(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
139            Number::I16(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
140            Number::I32(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
141            Number::I64(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
142            Number::I128(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
143            Number::U8(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
144            Number::U16(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
145            Number::U32(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
146            Number::U64(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
147            Number::U128(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
148            Number::Isize(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
149            Number::Usize(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
150            Number::BigFloat(f) => self.intern(ExprData::Real(Real::F64(*f))),
151        }
152    }
153
154    pub fn const_number(&self, id: ExprId) -> Option<Number> {
155        match self.get(id)? {
156            ExprData::Integer(i) => Some(Number::Integer(*i)),
157            ExprData::Rational(r) => Some(Number::Rational(*r)),
158            ExprData::Real(r) => Some(Number::Real(r)),
159            _ => None,
160        }
161    }
162
163    fn node_rank(&self, id: ExprId) -> u8 {
164        match self.get(id) {
165            Some(ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Real(_)) => 0,
166            Some(ExprData::Symbol(_)) => 1,
167            _ => 2,
168        }
169    }
170
171    pub fn is_const_zero(&self, id: ExprId) -> bool {
172        self.const_number(id).is_some_and(|n| n.is_zero())
173    }
174
175    pub fn is_const_one(&self, id: ExprId) -> bool {
176        self.const_number(id).is_some_and(|n| n.is_one())
177    }
178
179    /// Raw `Add` node (no simplification), stored in canonical order (spec §8.4).
180    pub fn add(&self, items: &[ExprId]) -> ExprId {
181        let mut v = items.to_vec();
182        v.sort_by_key(|&id| (self.node_rank(id), id));
183        self.intern(ExprData::Add(v.into_boxed_slice()))
184    }
185
186    /// Raw `Mul` node (no simplification), stored in canonical order (spec §8.4).
187    pub fn mul(&self, items: &[ExprId]) -> ExprId {
188        let mut v = items.to_vec();
189        v.sort_by_key(|&id| (self.node_rank(id), id));
190        self.intern(ExprData::Mul(v.into_boxed_slice()))
191    }
192
193    pub fn pow(&self, base: ExprId, exp: ExprId) -> ExprId {
194        self.intern(ExprData::Pow { base, exp })
195    }
196
197    pub fn apply(&self, f: ExprId, args: &[ExprId]) -> ExprId {
198        self.intern(ExprData::Apply { f, args: args.to_vec().into_boxed_slice() })
199    }
200
201    /// Level 0/1 addition simplification (spec §8.3): `Add` flattening, constant merging, `x+0→x`;
202    /// the result is sorted in canonical order (numbers/constants → symbols → composite nodes, spec §8.4).
203    pub fn add_n(&self, items: &[ExprId]) -> ExprId {
204        let mut flat = Vec::new();
205        for &it in items {
206            if let Some(ExprData::Add(inner)) = self.get(it) {
207                flat.extend_from_slice(&inner);
208            } else {
209                flat.push(it);
210            }
211        }
212        let mut const_sum: Option<Number> = None;
213        let mut rest = Vec::new();
214        for &it in &flat {
215            if let Some(n) = self.const_number(it) {
216                const_sum = Some(match const_sum {
217                    Some(acc) => acc + n,
218                    None => n,
219                });
220            } else {
221                rest.push(it);
222            }
223        }
224        if let Some(n) = const_sum.filter(|n| !n.is_zero()) {
225            rest.push(self.number(&n));
226        }
227        if rest.is_empty() {
228            return self.integer(0);
229        }
230        if rest.len() == 1 {
231            return rest[0];
232        }
233        rest.sort_by_key(|&id| (self.node_rank(id), id));
234        self.intern(ExprData::Add(rest.into_boxed_slice()))
235    }
236
237    /// Level 0/1 multiplication simplification (spec §8.3): `Mul` flattening, constant merging, `0*x→0`, `1*x→x`.
238    pub fn mul_n(&self, items: &[ExprId]) -> ExprId {
239        let mut flat = Vec::new();
240        for &it in items {
241            if let Some(ExprData::Mul(inner)) = self.get(it) {
242                flat.extend_from_slice(&inner);
243            } else {
244                flat.push(it);
245            }
246        }
247        let mut const_prod: Option<Number> = None;
248        let mut rest = Vec::new();
249        for &it in &flat {
250            if let Some(n) = self.const_number(it) {
251                if n.is_zero() {
252                    return self.integer(0);
253                }
254                const_prod = Some(match const_prod {
255                    Some(acc) => acc * n,
256                    None => n,
257                });
258            } else {
259                rest.push(it);
260            }
261        }
262        if let Some(n) = const_prod.filter(|n| !n.is_one()) {
263            rest.push(self.number(&n));
264        }
265        if rest.is_empty() {
266            return self.integer(1);
267        }
268        if rest.len() == 1 {
269            return rest[0];
270        }
271        rest.sort_by_key(|&id| (self.node_rank(id), id));
272        self.intern(ExprData::Mul(rest.into_boxed_slice()))
273    }
274
275    pub fn add2(&self, a: ExprId, b: ExprId) -> ExprId {
276        self.add_n(&[a, b])
277    }
278
279    pub fn mul2(&self, a: ExprId, b: ExprId) -> ExprId {
280        self.mul_n(&[a, b])
281    }
282
283    /// Level 0/1 power simplification (spec §8.3): `x^0→1`, `x^1→x`, `1^x→1`, plus constant folding at the same level.
284    pub fn pow2(&self, base: ExprId, exp: ExprId) -> ExprId {
285        if self.is_const_zero(exp) {
286            return self.integer(1);
287        }
288        if self.is_const_one(exp) {
289            return base;
290        }
291        if self.is_const_one(base) {
292            return self.integer(1);
293        }
294        if let (Some(b), Some(e)) = (self.const_number(base), self.const_number(exp))
295            && let Some(r) = b.pow(&e)
296        {
297            return self.number(&r);
298        }
299        self.pow(base, exp)
300    }
301
302    pub fn sub2(&self, a: ExprId, b: ExprId) -> ExprId {
303        self.add2(a, self.mul2(self.integer(-1), b))
304    }
305
306    pub fn div2(&self, a: ExprId, b: ExprId) -> ExprId {
307        self.mul2(a, self.pow2(b, self.integer(-1)))
308    }
309}
310
311impl Default for ExprPool {
312    fn default() -> Self {
313        Self::new()
314    }
315}
prima_core/expr_pool.rs

prima_core/
expr_pool.rs
