radiate_expr/
compile.rs

1use crate::ops::{BinaryOp, UnaryOp};
2use crate::{Expr, ExprNode};
3use radiate_utils::AnyValue;
4
5impl Expr {
6    /// Walks the tree bottom-up and rewrites algebraically equivalent shapes
7    /// into the smallest possible form. Specifically:
8    ///
9    /// - Pure-literal subtrees fold (`Lit(2) + Lit(3)` → `Lit(5)`)
10    /// - `Add` / `Sub` / `Mul` / `Div` with one literal operand fuses into a
11    ///   `Unary(Affine)` (`x * 5 + 3` → `Affine { scale: 5, bias: 3 }`)
12    /// - Nested affines collapse: `s2 * (s1*x + b1) + b2` → `Affine(s2*s1, s2*b1 + b2)`
13    pub fn compile(self) -> Expr {
14        let name = self.name;
15        let kind = compile_kind(self.node);
16        let id = self.id;
17        Expr {
18            name,
19            id,
20            node: kind,
21        }
22    }
23}
24
25fn compile_kind(kind: ExprNode) -> ExprNode {
26    match kind {
27        ExprNode::Literal(_) | ExprNode::Selector(_) | ExprNode::Schedule(_) => kind,
28
29        ExprNode::Unary { child: u, op } => {
30            let child = u;
31            let child = child.compile();
32            match op {
33                UnaryOp::Affine { scale, bias } => fuse_affine(child, scale, bias).node,
34                other_op => ExprNode::Unary {
35                    child: Box::new(child),
36                    op: other_op,
37                },
38            }
39        }
40
41        ExprNode::Trinary {
42            first,
43            second,
44            third,
45            op,
46        } => ExprNode::Trinary {
47            first: Box::new((*first).compile()),
48            second: Box::new((*second).compile()),
49            third: Box::new((*third).compile()),
50            op,
51        },
52
53        ExprNode::Binary {
54            lhs: lhs_box,
55            rhs: rhs_box,
56            op,
57        } => {
58            let lhs = (*lhs_box).compile();
59            let rhs = (*rhs_box).compile();
60            reduce_binary(lhs, rhs, op).node
61        }
62
63        ExprNode::Reduce { mut child, rollup } => {
64            let old_child = std::mem::replace(
65                &mut child,
66                Box::new(Expr::new(ExprNode::Literal(AnyValue::Null))),
67            );
68            let new_child = old_child.compile();
69            ExprNode::Reduce {
70                child: Box::new(new_child),
71                rollup,
72            }
73        }
74        ExprNode::Rolling { mut child, buffer } => {
75            let old_child = std::mem::replace(
76                &mut child,
77                Box::new(Expr::new(ExprNode::Literal(AnyValue::Null))),
78            );
79            let new_child = old_child.compile();
80            ExprNode::Rolling {
81                child: Box::new(new_child),
82                buffer,
83            }
84        }
85    }
86}
87
88fn reduce_binary(lhs: Expr, rhs: Expr, op: BinaryOp) -> Expr {
89    if let (ExprNode::Literal(l), ExprNode::Literal(r)) = (&lhs.node, &rhs.node)
90        && let Some(folded) = fold_literals(l, r, op)
91    {
92        return Expr::new(ExprNode::Literal(folded));
93    }
94
95    match op {
96        BinaryOp::Add => {
97            if let ExprNode::Literal(v) = &rhs.node
98                && let Some(k) = v.extract::<f32>()
99            {
100                return fuse_affine(lhs, 1.0, k);
101            }
102            if let ExprNode::Literal(v) = &lhs.node
103                && let Some(k) = v.extract::<f32>()
104            {
105                return fuse_affine(rhs, 1.0, k);
106            }
107        }
108        BinaryOp::Sub => {
109            if let ExprNode::Literal(v) = &rhs.node
110                && let Some(k) = v.extract::<f32>()
111            {
112                return fuse_affine(lhs, 1.0, -k);
113            }
114            if let ExprNode::Literal(v) = &lhs.node
115                && let Some(k) = v.extract::<f32>()
116            {
117                return fuse_affine(rhs, -1.0, k);
118            }
119        }
120        BinaryOp::Mul => {
121            if let ExprNode::Literal(v) = &rhs.node
122                && let Some(s) = v.extract::<f32>()
123            {
124                return fuse_affine(lhs, s, 0.0);
125            }
126            if let ExprNode::Literal(v) = &lhs.node
127                && let Some(s) = v.extract::<f32>()
128            {
129                return fuse_affine(rhs, s, 0.0);
130            }
131        }
132        BinaryOp::Div => {
133            if let ExprNode::Literal(v) = &rhs.node
134                && let Some(d) = v.extract::<f32>()
135                && d != 0.0
136                && d.is_finite()
137            {
138                return fuse_affine(lhs, 1.0 / d, 0.0);
139            }
140        }
141        _ => {}
142    }
143
144    Expr::new(ExprNode::Binary {
145        lhs: Box::new(lhs),
146        rhs: Box::new(rhs),
147        op,
148    })
149}
150
151fn fold_literals(
152    l: &AnyValue<'static>,
153    r: &AnyValue<'static>,
154    op: BinaryOp,
155) -> Option<AnyValue<'static>> {
156    let a = l.extract::<f32>()?;
157    let b = r.extract::<f32>()?;
158    let result = match op {
159        BinaryOp::Add => a + b,
160        BinaryOp::Sub => a - b,
161        BinaryOp::Mul => a * b,
162        BinaryOp::Div if b != 0.0 => a / b,
163        _ => return None,
164    };
165    if result.is_finite() {
166        Some(AnyValue::Float32(result))
167    } else {
168        None
169    }
170}
171
172/// Construct `Unary(Affine(scale * child + bias))`, collapsing nested affines.
173/// `scale * (s2 * x + b2) + bias = (scale * s2) * x + (scale * b2 + bias)`.
174///
175/// Shared between the `.affine(...)` builder and the compile-pass binary-fusion
176/// rewriters so both produce the same fused shape.
177fn fuse_affine(child: Expr, scale: f32, bias: f32) -> Expr {
178    if let ExprNode::Unary { child: inner, op } = child.node {
179        if matches!(op, UnaryOp::Affine { .. }) {
180            let UnaryOp::Affine {
181                scale: s2,
182                bias: b2,
183            } = op
184            else {
185                unreachable!()
186            };
187
188            return Expr::new(ExprNode::Unary {
189                child: inner,
190                op: UnaryOp::Affine {
191                    scale: scale * s2,
192                    bias: scale * b2 + bias,
193                },
194            });
195        }
196
197        return Expr::new(ExprNode::Unary {
198            child: Box::new(Expr::new(ExprNode::Unary { child: inner, op })),
199            op: UnaryOp::Affine { scale, bias },
200        });
201    }
202
203    Expr::new(ExprNode::Unary {
204        child: Box::new(child),
205        op: UnaryOp::Affine { scale, bias },
206    })
207}
radiate_expr/compile.rs

radiate_expr/
compile.rs
