tang_expr/
lib.rs

1//! tang-expr — RISC expression graph for symbolic computation.
2//!
3//! Builds computation graphs from generic `Scalar` code via a thread-local
4//! graph. Enables symbolic differentiation, sparsity detection, simplification,
5//! and multi-backend compilation (CPU closures, WGSL shaders).
6//!
7//! # Quick start
8//!
9//! ```
10//! use tang::Vec3;
11//! use tang_expr::{trace, ExprId};
12//!
13//! let (mut g, dot) = trace(|| {
14//!     let a = Vec3::new(ExprId::var(0), ExprId::var(1), ExprId::var(2));
15//!     let b = Vec3::new(ExprId::var(3), ExprId::var(4), ExprId::var(5));
16//!     a.dot(b)
17//! });
18//!
19//! // Evaluate with concrete values
20//! let result: f64 = g.eval(dot, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
21//! assert!((result - 32.0).abs() < 1e-10);
22//!
23//! // Symbolic differentiation
24//! let ddot_dx0 = g.diff(dot, 0);
25//! let ddot_dx0 = g.simplify(ddot_dx0);
26//! ```
27//!
28//! # Integer dialect
29//!
30//! [`int::IntGraph`] is a separate, bit-exact 32-bit integer graph
31//! (`i32`/`u32`/`bool`, wrapping arithmetic, `mul_wide`, structured loops)
32//! with three lowerings that compute the same bits: the CPU interpreter
33//! ([`int::IntGraph::eval`] / [`int::Program`]), WGSL and Rust source.
34//!
35//! ```
36//! use tang_expr::int::{IType, IntGraph};
37//!
38//! // sum of i*i for i in 0..n
39//! let mut g = IntGraph::new();
40//! let n = g.param(0, IType::U32);
41//! let zero = g.u32(0);
42//! let s = g.for_range(n, &[zero], |g, i, s| {
43//!     let sq = g.mul(i, i);
44//!     vec![g.add(s[0], sq)]
45//! });
46//! assert_eq!(g.eval(&s, &[10], &[]), vec![285]);
47//!
48//! let wgsl = g.to_wgsl_fn("sum_sq", &s);
49//! assert!(wgsl.contains("fn sum_sq"));
50//! ```
51
52pub mod codegen;
53pub mod compile;
54pub mod diff;
55pub mod display;
56pub mod eval;
57pub mod graph;
58pub mod int;
59pub mod node;
60pub mod reverse;
61mod scalar;
62pub mod simplify;
63pub mod sparsity;
64pub mod split;
65pub mod wgsl;
66
67pub use graph::ExprGraph;
68pub use node::{ExprId, Node};
69
70use std::cell::RefCell;
71
72thread_local! {
73    static GRAPH: RefCell<ExprGraph> = RefCell::new(ExprGraph::new());
74}
75
76/// Access the thread-local graph.
77pub fn with_graph<F, R>(f: F) -> R
78where
79    F: FnOnce(&mut ExprGraph) -> R,
80{
81    GRAPH.with(|g| f(&mut g.borrow_mut()))
82}
83
84/// Run a closure with a fresh graph, returning the graph and result.
85///
86/// Installs a new empty graph, runs `f` (which builds the expression via
87/// `ExprId` arithmetic / `Scalar` calls), then extracts the graph.
88pub fn trace<F, R>(f: F) -> (ExprGraph, R)
89where
90    F: FnOnce() -> R,
91{
92    // Swap in a fresh graph
93    GRAPH.with(|g| {
94        let old = std::mem::take(&mut *g.borrow_mut());
95        let result = f();
96        let graph = std::mem::replace(&mut *g.borrow_mut(), old);
97        (graph, result)
98    })
99}
100
101impl ExprId {
102    /// Create a variable node in the thread-local graph.
103    #[inline]
104    pub fn var(n: u16) -> Self {
105        with_graph(|g| g.var(n))
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use tang::{Scalar, Vec3};
113
114    #[test]
115    fn trace_vec3_dot() {
116        let (g, dot) = trace(|| {
117            let a = Vec3::new(ExprId::var(0), ExprId::var(1), ExprId::var(2));
118            let b = Vec3::new(ExprId::var(3), ExprId::var(4), ExprId::var(5));
119            a.dot(b)
120        });
121
122        // Evaluate: [1,2,3] . [4,5,6] = 4+10+18 = 32
123        let result: f64 = g.eval(dot, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
124        assert!((result - 32.0).abs() < 1e-10);
125    }
126
127    #[test]
128    fn trace_vec3_norm() {
129        let (g, norm) = trace(|| {
130            let v = Vec3::new(ExprId::var(0), ExprId::var(1), ExprId::var(2));
131            v.norm()
132        });
133
134        // norm([3,4,0]) = 5
135        let result: f64 = g.eval(norm, &[3.0, 4.0, 0.0]);
136        assert!((result - 5.0).abs() < 1e-10);
137    }
138
139    #[test]
140    fn trace_isolation() {
141        // Traces should be isolated
142        let (g1, _) = trace(|| {
143            let _x = ExprId::var(0);
144        });
145        let (g2, _) = trace(|| {
146            let _x = ExprId::var(0);
147        });
148        // Both graphs should have same size (3 pre-populated + 1 var)
149        assert_eq!(g1.len(), g2.len());
150    }
151
152    #[test]
153    fn from_f64_creates_lit() {
154        let (g, v) = trace(|| ExprId::from_f64(42.0));
155        let result: f64 = g.eval(v, &[]);
156        assert!((result - 42.0).abs() < 1e-10);
157    }
158
159    #[test]
160    fn scalar_constants() {
161        let (g, (zero, one, two)) = trace(|| {
162            let z: ExprId = Scalar::from_f64(0.0);
163            let o: ExprId = Scalar::from_f64(1.0);
164            let t: ExprId = Scalar::from_f64(2.0);
165            (z, o, t)
166        });
167        assert_eq!(g.eval::<f64>(zero, &[]), 0.0);
168        assert_eq!(g.eval::<f64>(one, &[]), 1.0);
169        assert_eq!(g.eval::<f64>(two, &[]), 2.0);
170    }
171}
tang_expr/lib.rs

tang_expr/
lib.rs
