ocas_atom/lib.rs
1//! Symbolic expression tree and rewriting for oCAS.
2//!
3//! This crate provides the [`Atom`] type: an arena-backed, tagged-union
4//! representation of symbolic expressions. Atoms are immutable, copyable
5//! references into an [`Arena`] and form the core data structure used by the
6//! parser, printer, and rewrite engine.
7
8use std::cell::RefCell;
9use std::collections::HashMap;
10use std::sync::{Mutex, OnceLock};
11
12use ocas_core::arena::Arena;
13
14pub mod normalize;
15
16/// An interned symbolic name (variable, function, or constant).
17///
18/// Symbols are deduplicated globally and live for the remainder of the
19/// process. This keeps [`Atom`] small and comparable by identity.
20///
21/// # Example
22///
23/// ```
24/// use ocas_atom::Symbol;
25///
26/// let x = Symbol::new("x");
27/// let also_x = Symbol::new("x");
28/// assert_eq!(x, also_x);
29/// assert_eq!(x.as_str(), "x");
30/// ```
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
32pub struct Symbol(&'static str);
33
34impl Symbol {
35 /// Create a symbol from a name, interning it globally.
36 ///
37 /// # Example
38 ///
39 /// ```
40 /// use ocas_atom::Symbol;
41 ///
42 /// let sym = Symbol::new("my_var");
43 /// assert_eq!(sym.as_str(), "my_var");
44 /// ```
45 pub fn new(name: &str) -> Self {
46 Self(intern(name))
47 }
48
49 /// Return the symbol's string representation.
50 ///
51 /// # Example
52 ///
53 /// ```
54 /// use ocas_atom::Symbol;
55 ///
56 /// let sym = Symbol::new("y");
57 /// assert_eq!(sym.as_str(), "y");
58 /// ```
59 pub fn as_str(&self) -> &str {
60 self.0
61 }
62}
63
64fn intern(name: &str) -> &'static str {
65 static TABLE: OnceLock<Mutex<HashMap<String, &'static str>>> = OnceLock::new();
66 let table = TABLE.get_or_init(|| Mutex::new(HashMap::new()));
67 let mut table = table.lock().expect("symbol interner lock poisoned");
68 table.entry(name.to_owned()).or_insert_with(|| {
69 let boxed = name.to_owned().into_boxed_str();
70 Box::leak(boxed)
71 })
72}
73
74/// A reference to an expression node allocated in an arena.
75///
76/// `Atom` is a small copyable handle. The actual node data lives in the
77/// [`Arena`] and is freed when the arena is dropped.
78///
79/// # Example
80///
81/// ```
82/// use ocas_atom::AtomArena;
83/// use ocas_core::arena::Arena;
84///
85/// let arena = Arena::new();
86/// let ctx = AtomArena::new(&arena);
87/// let x = ctx.var("x");
88/// let two = ctx.num(2);
89/// let expr = ctx.pow(x, two);
90/// assert_eq!(expr.to_string(), "x^2");
91/// ```
92#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
93pub struct Atom<'a>(&'a AtomNode<'a>);
94
95impl<'a> Atom<'a> {
96 /// Access the underlying node data.
97 ///
98 /// # Example
99 ///
100 /// ```
101 /// use ocas_atom::{AtomArena, AtomNode};
102 /// use ocas_core::arena::Arena;
103 ///
104 /// let arena = Arena::new();
105 /// let ctx = AtomArena::new(&arena);
106 /// let x = ctx.var("x");
107 /// assert!(matches!(x.node(), AtomNode::Var(_)));
108 /// ```
109 pub fn node(&self) -> &'a AtomNode<'a> {
110 self.0
111 }
112
113 /// Returns the direct children of this atom, in left-to-right order.
114 ///
115 /// `Num`, `Var`, and `Pow` report no children through this API; use
116 /// [`Self::binary_children`] for the two operands of `Pow`.
117 ///
118 /// # Example
119 ///
120 /// ```
121 /// use ocas_atom::AtomArena;
122 /// use ocas_core::arena::Arena;
123 ///
124 /// let arena = Arena::new();
125 /// let ctx = AtomArena::new(&arena);
126 /// let x = ctx.var("x");
127 /// let y = ctx.var("y");
128 /// let sum = ctx.add(&[x, y, ctx.num(1)]);
129 /// assert_eq!(sum.children().len(), 3);
130 /// ```
131 pub fn children(&self) -> &'a [Atom<'a>] {
132 match self.node() {
133 AtomNode::Num(_) | AtomNode::Var(_) => &[],
134 AtomNode::Fun(_, args) | AtomNode::Add(args) | AtomNode::Mul(args) => args,
135 AtomNode::Pow(base, exp) => {
136 // This function cannot return a dynamically-allocated slice,
137 // so callers that need the two-element slice should use
138 // [`Self::binary_children`]. For now, `Pow` reports no children
139 // through this API to keep the return type a plain slice.
140 let _ = (base, exp);
141 &[]
142 }
143 }
144 }
145
146 /// If this atom is a binary operator, return its two operands.
147 ///
148 /// # Example
149 ///
150 /// ```
151 /// use ocas_atom::AtomArena;
152 /// use ocas_core::arena::Arena;
153 ///
154 /// let arena = Arena::new();
155 /// let ctx = AtomArena::new(&arena);
156 /// let x = ctx.var("x");
157 /// let y = ctx.var("y");
158 /// let power = ctx.pow(x, y);
159 /// let (base, exp) = power.binary_children().unwrap();
160 /// assert_eq!(base.to_string(), "x");
161 /// assert_eq!(exp.to_string(), "y");
162 /// ```
163 pub fn binary_children(&self) -> Option<(Atom<'a>, Atom<'a>)> {
164 match self.node() {
165 AtomNode::Pow(base, exp) => Some((*base, *exp)),
166 _ => None,
167 }
168 }
169}
170
171/// The concrete data stored for each expression node.
172///
173/// # Example
174///
175/// ```
176/// use ocas_atom::{AtomArena, AtomNode};
177/// use ocas_core::arena::Arena;
178///
179/// let arena = Arena::new();
180/// let ctx = AtomArena::new(&arena);
181/// let x = ctx.var("x");
182/// match x.node() {
183/// AtomNode::Var(s) => assert_eq!(s.as_str(), "x"),
184/// _ => panic!("expected variable"),
185/// }
186/// ```
187#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
188pub enum AtomNode<'a> {
189 /// A 64-bit signed integer literal.
190 Num(i64),
191 /// A named variable or constant.
192 Var(Symbol),
193 /// A named function applied to a list of arguments.
194 Fun(Symbol, &'a [Atom<'a>]),
195 /// A sum of sub-expressions.
196 Add(&'a [Atom<'a>]),
197 /// A product of sub-expressions.
198 Mul(&'a [Atom<'a>]),
199 /// A power with base and exponent.
200 Pow(Atom<'a>, Atom<'a>),
201}
202
203/// A context that allocates [`Atom`]s in an [`Arena`].
204///
205/// All construction methods are immutable from the caller's perspective;
206/// mutation happens through the arena's interior mutability. Identical
207/// sub-expressions are hash-consed so that structural equality implies
208/// pointer equality.
209///
210/// # Example
211///
212/// ```
213/// use ocas_atom::AtomArena;
214/// use ocas_core::arena::Arena;
215///
216/// let arena = Arena::new();
217/// let ctx = AtomArena::new(&arena);
218/// let x = ctx.var("x");
219/// let y = ctx.var("y");
220/// let sum = ctx.add(&[x, y]);
221/// assert_eq!(sum.to_string(), "x + y");
222/// ```
223pub struct AtomArena<'a> {
224 arena: &'a Arena,
225 cons_table: RefCell<HashMap<AtomNode<'a>, Atom<'a>>>,
226}
227
228impl<'a> AtomArena<'a> {
229 /// Create an `AtomArena` backed by the given arena.
230 ///
231 /// # Example
232 ///
233 /// ```
234 /// use ocas_atom::AtomArena;
235 /// use ocas_core::arena::Arena;
236 ///
237 /// let arena = Arena::new();
238 /// let ctx = AtomArena::new(&arena);
239 /// let n = ctx.num(42);
240 /// assert_eq!(n.to_string(), "42");
241 /// ```
242 pub fn new(arena: &'a Arena) -> Self {
243 Self {
244 arena,
245 cons_table: RefCell::new(HashMap::new()),
246 }
247 }
248
249 fn intern(&self, candidate: AtomNode<'a>) -> Atom<'a> {
250 let mut table = self.cons_table.borrow_mut();
251 *table
252 .entry(candidate)
253 .or_insert_with(|| Atom(self.arena.allocate_with(|| candidate)))
254 }
255
256 /// Create an integer literal atom.
257 ///
258 /// # Example
259 ///
260 /// ```
261 /// use ocas_atom::AtomArena;
262 /// use ocas_core::arena::Arena;
263 ///
264 /// let arena = Arena::new();
265 /// let ctx = AtomArena::new(&arena);
266 /// let n = ctx.num(7);
267 /// assert_eq!(n.to_string(), "7");
268 /// ```
269 pub fn num(&self, value: i64) -> Atom<'a> {
270 self.intern(AtomNode::Num(value))
271 }
272
273 /// Create a variable atom from a name.
274 ///
275 /// # Example
276 ///
277 /// ```
278 /// use ocas_atom::AtomArena;
279 /// use ocas_core::arena::Arena;
280 ///
281 /// let arena = Arena::new();
282 /// let ctx = AtomArena::new(&arena);
283 /// let x = ctx.var("x");
284 /// assert_eq!(x.to_string(), "x");
285 /// ```
286 pub fn var(&self, name: &str) -> Atom<'a> {
287 self.intern(AtomNode::Var(Symbol::new(name)))
288 }
289
290 /// Create a function application atom.
291 ///
292 /// # Panics
293 ///
294 /// Panics in debug mode if `args` is empty.
295 ///
296 /// # Example
297 ///
298 /// ```
299 /// use ocas_atom::AtomArena;
300 /// use ocas_core::arena::Arena;
301 ///
302 /// let arena = Arena::new();
303 /// let ctx = AtomArena::new(&arena);
304 /// let x = ctx.var("x");
305 /// let f = ctx.fun("sin", &[x]);
306 /// assert_eq!(f.to_string(), "sin(x)");
307 /// ```
308 pub fn fun(&self, name: &str, args: &[Atom<'a>]) -> Atom<'a> {
309 debug_assert!(!args.is_empty(), "Fun node requires at least one argument");
310 let slice = self.arena.allocate_slice(args);
311 self.intern(AtomNode::Fun(Symbol::new(name), slice))
312 }
313
314 /// Create an addition atom.
315 ///
316 /// # Panics
317 ///
318 /// Panics in debug mode if `args` is empty.
319 ///
320 /// # Example
321 ///
322 /// ```
323 /// use ocas_atom::AtomArena;
324 /// use ocas_core::arena::Arena;
325 ///
326 /// let arena = Arena::new();
327 /// let ctx = AtomArena::new(&arena);
328 /// let x = ctx.var("x");
329 /// let y = ctx.var("y");
330 /// let sum = ctx.add(&[x, y]);
331 /// assert_eq!(sum.to_string(), "x + y");
332 /// ```
333 pub fn add(&self, args: &[Atom<'a>]) -> Atom<'a> {
334 debug_assert!(!args.is_empty(), "Add node requires at least one argument");
335 let slice = self.arena.allocate_slice(args);
336 self.intern(AtomNode::Add(slice))
337 }
338
339 /// Create a multiplication atom.
340 ///
341 /// # Panics
342 ///
343 /// Panics in debug mode if `args` is empty.
344 ///
345 /// # Example
346 ///
347 /// ```
348 /// use ocas_atom::AtomArena;
349 /// use ocas_core::arena::Arena;
350 ///
351 /// let arena = Arena::new();
352 /// let ctx = AtomArena::new(&arena);
353 /// let x = ctx.var("x");
354 /// let y = ctx.var("y");
355 /// let product = ctx.mul(&[x, y]);
356 /// assert_eq!(product.to_string(), "x*y");
357 /// ```
358 pub fn mul(&self, args: &[Atom<'a>]) -> Atom<'a> {
359 debug_assert!(!args.is_empty(), "Mul node requires at least one argument");
360 let slice = self.arena.allocate_slice(args);
361 self.intern(AtomNode::Mul(slice))
362 }
363
364 /// Create a power atom.
365 ///
366 /// # Example
367 ///
368 /// ```
369 /// use ocas_atom::AtomArena;
370 /// use ocas_core::arena::Arena;
371 ///
372 /// let arena = Arena::new();
373 /// let ctx = AtomArena::new(&arena);
374 /// let x = ctx.var("x");
375 /// let p = ctx.pow(x, ctx.num(3));
376 /// assert_eq!(p.to_string(), "x^3");
377 /// ```
378 pub fn pow(&self, base: Atom<'a>, exp: Atom<'a>) -> Atom<'a> {
379 self.intern(AtomNode::Pow(base, exp))
380 }
381}
382
383impl std::fmt::Display for Atom<'_> {
384 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
385 match self.node() {
386 AtomNode::Num(n) => write!(f, "{n}"),
387 AtomNode::Var(s) => write!(f, "{}", s.as_str()),
388 AtomNode::Fun(name, args) => {
389 write!(f, "{}(", name.as_str())?;
390 for (i, arg) in args.iter().enumerate() {
391 if i > 0 {
392 write!(f, ", ")?;
393 }
394 write!(f, "{arg}")?;
395 }
396 write!(f, ")")
397 }
398 AtomNode::Add(args) => {
399 for (i, arg) in args.iter().enumerate() {
400 if i > 0 {
401 write!(f, " + ")?;
402 }
403 write_parenthesized(arg, f)?;
404 }
405 Ok(())
406 }
407 AtomNode::Mul(args) => {
408 for (i, arg) in args.iter().enumerate() {
409 if i > 0 {
410 write!(f, "*")?;
411 }
412 write_parenthesized(arg, f)?;
413 }
414 Ok(())
415 }
416 AtomNode::Pow(base, exp) => {
417 write_parenthesized(base, f)?;
418 write!(f, "^")?;
419 write_parenthesized(exp, f)
420 }
421 }
422 }
423}
424
425fn write_parenthesized(atom: &Atom<'_>, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
426 match atom.node() {
427 AtomNode::Num(_) | AtomNode::Var(_) => write!(f, "{atom}"),
428 _ => write!(f, "({atom})"),
429 }
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435
436 #[test]
437 fn construct_num() {
438 let arena = Arena::new();
439 let ctx = AtomArena::new(&arena);
440 let n = ctx.num(42);
441 assert_eq!(n.to_string(), "42");
442 assert!(matches!(n.node(), AtomNode::Num(42)));
443 }
444
445 #[test]
446 fn construct_var() {
447 let arena = Arena::new();
448 let ctx = AtomArena::new(&arena);
449 let x = ctx.var("x");
450 assert_eq!(x.to_string(), "x");
451 assert!(matches!(x.node(), AtomNode::Var(s) if s.as_str() == "x"));
452 }
453
454 #[test]
455 fn construct_add() {
456 let arena = Arena::new();
457 let ctx = AtomArena::new(&arena);
458 let x = ctx.var("x");
459 let y = ctx.var("y");
460 let sum = ctx.add(&[x, y]);
461 assert_eq!(sum.to_string(), "x + y");
462 }
463
464 #[test]
465 fn construct_mul() {
466 let arena = Arena::new();
467 let ctx = AtomArena::new(&arena);
468 let x = ctx.var("x");
469 let two = ctx.num(2);
470 let prod = ctx.mul(&[x, two]);
471 assert_eq!(prod.to_string(), "x*2");
472 }
473
474 #[test]
475 fn construct_fun() {
476 let arena = Arena::new();
477 let ctx = AtomArena::new(&arena);
478 let x = ctx.var("x");
479 let sin = ctx.fun("sin", &[x]);
480 assert_eq!(sin.to_string(), "sin(x)");
481 }
482
483 #[test]
484 fn fun_with_multiple_args() {
485 let arena = Arena::new();
486 let ctx = AtomArena::new(&arena);
487 let x = ctx.var("x");
488 let y = ctx.var("y");
489 let f = ctx.fun("f", &[x, y]);
490 assert_eq!(f.to_string(), "f(x, y)");
491 }
492
493 #[test]
494 fn children_returns_direct_subexpressions() {
495 let arena = Arena::new();
496 let ctx = AtomArena::new(&arena);
497 let x = ctx.var("x");
498 let y = ctx.var("y");
499 let sum = ctx.add(&[x, y]);
500 assert_eq!(sum.children(), &[x, y]);
501 assert_eq!(x.children(), &[]);
502 }
503
504 #[test]
505 fn nested_expression_prints_with_parentheses() {
506 let arena = Arena::new();
507 let ctx = AtomArena::new(&arena);
508 let x = ctx.var("x");
509 let y = ctx.var("y");
510 let sum = ctx.add(&[x, y]);
511 let two = ctx.num(2);
512 let squared = ctx.pow(sum, two);
513 assert_eq!(squared.to_string(), "(x + y)^2");
514 }
515
516 #[test]
517 fn atom_equality_uses_structure() {
518 let arena = Arena::new();
519 let ctx = AtomArena::new(&arena);
520 let x = ctx.var("x");
521 let y = ctx.var("y");
522 let a = ctx.add(&[x, y]);
523 let b = ctx.add(&[x, y]);
524 let c = ctx.add(&[y, x]);
525 assert_eq!(a, b);
526 assert_ne!(a, c);
527 }
528
529 #[test]
530 fn symbol_identity_is_preserved() {
531 let a = Symbol::new("x");
532 let b = Symbol::new("x");
533 let c = Symbol::new("y");
534 assert_eq!(a, b);
535 assert_ne!(a, c);
536 assert_eq!(a.as_str(), "x");
537 }
538
539 #[test]
540 fn atom_is_copyable() {
541 let arena = Arena::new();
542 let ctx = AtomArena::new(&arena);
543 let x = ctx.var("x");
544 let copied = x;
545 assert_eq!(x, copied);
546 }
547
548 #[test]
549 fn hash_consing_reuses_identical_nodes() {
550 let arena = Arena::new();
551 let ctx = AtomArena::new(&arena);
552 let x = ctx.var("x");
553 let y = ctx.var("y");
554 let a = ctx.add(&[x, y]);
555 let b = ctx.add(&[x, y]);
556 // Hash-consing should return the same arena pointer.
557 assert!(std::ptr::eq(a.node(), b.node()));
558 }
559
560 #[test]
561 fn hash_consing_distinguishes_different_nodes() {
562 let arena = Arena::new();
563 let ctx = AtomArena::new(&arena);
564 let x = ctx.var("x");
565 let y = ctx.var("y");
566 let a = ctx.add(&[x, y]);
567 let b = ctx.add(&[y, x]);
568 assert!(!std::ptr::eq(a.node(), b.node()));
569 }
570}
571
572#[cfg(test)]
573mod proptests {
574 use super::*;
575 use ocas_core::arena::Arena;
576 use proptest::prelude::*;
577
578 /// Owned expression tree used for property-test generation.
579 #[derive(Debug, Clone)]
580 enum PropExpr {
581 Num(i64),
582 Var(&'static str),
583 Fun(&'static str, Vec<PropExpr>),
584 Add(Vec<PropExpr>),
585 Mul(Vec<PropExpr>),
586 Pow(Box<PropExpr>, Box<PropExpr>),
587 }
588
589 fn build_atom<'a>(ctx: &AtomArena<'a>, expr: &PropExpr) -> Atom<'a> {
590 match expr {
591 PropExpr::Num(n) => ctx.num(*n),
592 PropExpr::Var(name) => ctx.var(name),
593 PropExpr::Fun(name, args) => {
594 let atoms: Vec<Atom<'a>> = args.iter().map(|a| build_atom(ctx, a)).collect();
595 ctx.fun(name, &atoms)
596 }
597 PropExpr::Add(args) => {
598 let atoms: Vec<Atom<'a>> = args.iter().map(|a| build_atom(ctx, a)).collect();
599 ctx.add(&atoms)
600 }
601 PropExpr::Mul(args) => {
602 let atoms: Vec<Atom<'a>> = args.iter().map(|a| build_atom(ctx, a)).collect();
603 ctx.mul(&atoms)
604 }
605 PropExpr::Pow(base, exp) => ctx.pow(build_atom(ctx, base), build_atom(ctx, exp)),
606 }
607 }
608
609 fn prop_expr() -> impl Strategy<Value = PropExpr> {
610 let leaf = prop_oneof![
611 (-100..100i64).prop_map(PropExpr::Num),
612 Just(PropExpr::Var("x")),
613 Just(PropExpr::Var("y")),
614 Just(PropExpr::Var("z")),
615 ];
616 leaf.prop_recursive(4, 64, 4, |inner| {
617 prop_oneof![
618 inner.clone().prop_map(|e| PropExpr::Fun("sin", vec![e])),
619 inner.clone().prop_map(|e| PropExpr::Fun("cos", vec![e])),
620 prop::collection::vec(inner.clone(), 1..4).prop_map(PropExpr::Add),
621 prop::collection::vec(inner.clone(), 1..4).prop_map(PropExpr::Mul),
622 (inner.clone(), inner.clone())
623 .prop_map(|(b, e)| PropExpr::Pow(Box::new(b), Box::new(e))),
624 ]
625 })
626 }
627
628 proptest! {
629 #[test]
630 fn normalize_is_idempotent(expr in prop_expr()) {
631 let arena = Arena::new();
632 let ctx = AtomArena::new(&arena);
633 let atom = build_atom(&ctx, &expr);
634 let once = normalize::normalize(&ctx, atom);
635 let twice = normalize::normalize(&ctx, once);
636 assert_eq!(once.to_string(), twice.to_string());
637 }
638
639 #[test]
640 fn add_identity(expr in prop_expr()) {
641 let arena = Arena::new();
642 let ctx = AtomArena::new(&arena);
643 let atom = build_atom(&ctx, &expr);
644 let zero = ctx.num(0);
645 let with_zero = ctx.add(&[atom, zero]);
646 let normalized = normalize::normalize(&ctx, with_zero);
647 assert_eq!(normalized.to_string(), normalize::normalize(&ctx, atom).to_string());
648 }
649
650 #[test]
651 fn mul_identity(expr in prop_expr()) {
652 let arena = Arena::new();
653 let ctx = AtomArena::new(&arena);
654 let atom = build_atom(&ctx, &expr);
655 let one = ctx.num(1);
656 let with_one = ctx.mul(&[atom, one]);
657 let normalized = normalize::normalize(&ctx, with_one);
658 assert_eq!(normalized.to_string(), normalize::normalize(&ctx, atom).to_string());
659 }
660 }
661}