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::HashMap;
7use std::collections::hash_map::DefaultHasher;
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 { .. } => {
137                panic!("complex numbers cannot be interned as expression nodes yet")
138            }
139            // Fixed-width collapsed layer interns to the exact/`Real` node (spec §6.1).
140            Number::I8(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
141            Number::I16(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
142            Number::I32(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
143            Number::I64(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
144            Number::I128(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
145            Number::U8(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
146            Number::U16(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
147            Number::U32(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
148            Number::U64(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
149            Number::U128(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
150            Number::Isize(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
151            Number::Usize(v) => self.intern(ExprData::Integer(Box::new(BigInt::from(*v)))),
152            Number::BigFloat(f) => self.intern(ExprData::Real(Real::F64(*f))),
153        }
154    }
155
156    pub fn const_number(&self, id: ExprId) -> Option<Number> {
157        match self.get(id)? {
158            ExprData::Integer(i) => Some(Number::Integer(*i)),
159            ExprData::Rational(r) => Some(Number::Rational(*r)),
160            ExprData::Real(r) => Some(Number::Real(r)),
161            _ => None,
162        }
163    }
164
165    fn node_rank(&self, id: ExprId) -> u8 {
166        match self.get(id) {
167            Some(ExprData::Integer(_) | ExprData::Rational(_) | ExprData::Real(_)) => 0,
168            Some(ExprData::Symbol(_)) => 1,
169            _ => 2,
170        }
171    }
172
173    pub fn is_const_zero(&self, id: ExprId) -> bool {
174        self.const_number(id).is_some_and(|n| n.is_zero())
175    }
176
177    pub fn is_const_one(&self, id: ExprId) -> bool {
178        self.const_number(id).is_some_and(|n| n.is_one())
179    }
180
181    /// Raw `Add` node (no simplification), stored in canonical order (spec §8.4).
182    pub fn add(&self, items: &[ExprId]) -> ExprId {
183        let mut v = items.to_vec();
184        v.sort_by_key(|&id| (self.node_rank(id), id));
185        self.intern(ExprData::Add(v.into_boxed_slice()))
186    }
187
188    /// Raw `Mul` node (no simplification), stored in canonical order (spec §8.4).
189    pub fn mul(&self, items: &[ExprId]) -> ExprId {
190        let mut v = items.to_vec();
191        v.sort_by_key(|&id| (self.node_rank(id), id));
192        self.intern(ExprData::Mul(v.into_boxed_slice()))
193    }
194
195    pub fn pow(&self, base: ExprId, exp: ExprId) -> ExprId {
196        self.intern(ExprData::Pow { base, exp })
197    }
198
199    pub fn apply(&self, f: ExprId, args: &[ExprId]) -> ExprId {
200        self.intern(ExprData::Apply {
201            f,
202            args: args.to_vec().into_boxed_slice(),
203        })
204    }
205
206    /// Level 0/1 addition simplification (spec §8.3): `Add` flattening, constant merging, `x+0→x`;
207    /// the result is sorted in canonical order (numbers/constants → symbols → composite nodes, spec §8.4).
208    pub fn add_n(&self, items: &[ExprId]) -> ExprId {
209        let mut flat = Vec::new();
210        for &it in items {
211            if let Some(ExprData::Add(inner)) = self.get(it) {
212                flat.extend_from_slice(&inner);
213            } else {
214                flat.push(it);
215            }
216        }
217        let mut const_sum: Option<Number> = None;
218        let mut rest = Vec::new();
219        for &it in &flat {
220            if let Some(n) = self.const_number(it) {
221                const_sum = Some(match const_sum {
222                    Some(acc) => acc + n,
223                    None => n,
224                });
225            } else {
226                rest.push(it);
227            }
228        }
229        if let Some(n) = const_sum.filter(|n| !n.is_zero()) {
230            rest.push(self.number(&n));
231        }
232        if rest.is_empty() {
233            return self.integer(0);
234        }
235        if rest.len() == 1 {
236            return rest[0];
237        }
238        rest.sort_by_key(|&id| (self.node_rank(id), id));
239        self.intern(ExprData::Add(rest.into_boxed_slice()))
240    }
241
242    /// Level 0/1 multiplication simplification (spec §8.3): `Mul` flattening, constant merging, `0*x→0`, `1*x→x`.
243    pub fn mul_n(&self, items: &[ExprId]) -> ExprId {
244        let mut flat = Vec::new();
245        for &it in items {
246            if let Some(ExprData::Mul(inner)) = self.get(it) {
247                flat.extend_from_slice(&inner);
248            } else {
249                flat.push(it);
250            }
251        }
252        let mut const_prod: Option<Number> = None;
253        let mut rest = Vec::new();
254        for &it in &flat {
255            if let Some(n) = self.const_number(it) {
256                if n.is_zero() {
257                    return self.integer(0);
258                }
259                const_prod = Some(match const_prod {
260                    Some(acc) => acc * n,
261                    None => n,
262                });
263            } else {
264                rest.push(it);
265            }
266        }
267        if let Some(n) = const_prod.filter(|n| !n.is_one()) {
268            rest.push(self.number(&n));
269        }
270        if rest.is_empty() {
271            return self.integer(1);
272        }
273        if rest.len() == 1 {
274            return rest[0];
275        }
276        rest.sort_by_key(|&id| (self.node_rank(id), id));
277        self.intern(ExprData::Mul(rest.into_boxed_slice()))
278    }
279
280    pub fn add2(&self, a: ExprId, b: ExprId) -> ExprId {
281        self.add_n(&[a, b])
282    }
283
284    pub fn mul2(&self, a: ExprId, b: ExprId) -> ExprId {
285        self.mul_n(&[a, b])
286    }
287
288    /// Level 0/1 power simplification (spec §8.3): `x^0→1`, `x^1→x`, `1^x→1`, plus constant folding at the same level.
289    pub fn pow2(&self, base: ExprId, exp: ExprId) -> ExprId {
290        if self.is_const_zero(exp) {
291            return self.integer(1);
292        }
293        if self.is_const_one(exp) {
294            return base;
295        }
296        if self.is_const_one(base) {
297            return self.integer(1);
298        }
299        if let (Some(b), Some(e)) = (self.const_number(base), self.const_number(exp))
300            && let Some(r) = b.pow(&e)
301        {
302            return self.number(&r);
303        }
304        self.pow(base, exp)
305    }
306
307    pub fn sub2(&self, a: ExprId, b: ExprId) -> ExprId {
308        self.add2(a, self.mul2(self.integer(-1), b))
309    }
310
311    pub fn div2(&self, a: ExprId, b: ExprId) -> ExprId {
312        self.mul2(a, self.pow2(b, self.integer(-1)))
313    }
314}
315
316impl Default for ExprPool {
317    fn default() -> Self {
318        Self::new()
319    }
320}
prima_core/expr_pool.rs

prima_core/
expr_pool.rs
