Skip to main content

fidget_core/context/
mod.rs

1//! Infrastructure for representing math expressions as trees and graphs
2//!
3//! There are two families of representations in this module:
4//!
5//! - A [`Tree`] is a free-floating math expression, which can be cloned
6//!   and has overloaded operators for ease of use.  It is **not** deduplicated;
7//!   two calls to [`Tree::constant(1.0)`](Tree::constant) will allocate two
8//!   different objects.
9//!   `Tree` objects are typically used when building up expressions; they
10//!   should be converted to `Node` objects (in a particular `Context`) after
11//!   they have been constructed.
12//! - A [`Context`] is an arena for unique (deduplicated) math expressions,
13//!   which are represented as [`Node`] handles.  Each `Node` is specific to a
14//!   particular context.  Only `Node` objects can be converted into
15//!   [`Function`](crate::eval::Function) objects for evaluation.
16//!
17//! In other words, the typical workflow is `Tree → (Context, Node) → Function`.
18mod indexed;
19mod op;
20mod tree;
21
22use indexed::{Index, IndexMap, IndexVec, define_index};
23pub use op::{BinaryOpcode, Op, UnaryOpcode};
24pub use tree::{Tree, TreeOp};
25
26use crate::var::Var;
27
28use std::collections::{BTreeMap, HashMap};
29use std::fmt::Write;
30use std::io::{BufRead, BufReader, Read};
31use std::sync::Arc;
32
33use nalgebra::Matrix4;
34use ordered_float::OrderedFloat;
35
36define_index!(Node, "An index in the `Context::ops` map");
37
38/// A `Context` holds a set of deduplicated constants, variables, and
39/// operations.
40///
41/// It should be used like an arena allocator: it grows over time, then frees
42/// all of its contents when dropped.  There is no reference counting within the
43/// context.
44///
45/// Items in the context are accessed with [`Node`] keys, which are simple
46/// handles into an internal map.  Inside the context, operations are
47/// represented with the [`Op`] type.
48#[derive(Debug, Default)]
49pub struct Context {
50    ops: IndexMap<Op, Node>,
51}
52
53impl Context {
54    /// Build a new empty context
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Clears the context
60    ///
61    /// All [`Node`] handles from this context are invalidated.
62    ///
63    /// ```
64    /// # use fidget_core::context::Context;
65    /// let mut ctx = Context::new();
66    /// let x = ctx.x();
67    /// ctx.clear();
68    /// assert!(ctx.eval_xyz(x, 1.0, 0.0, 0.0).is_err());
69    /// ```
70    pub fn clear(&mut self) {
71        self.ops.clear();
72    }
73
74    /// Returns the number of [`Op`] nodes in the context
75    ///
76    /// ```
77    /// # use fidget_core::context::Context;
78    /// let mut ctx = Context::new();
79    /// let x = ctx.x();
80    /// assert_eq!(ctx.len(), 1);
81    /// let y = ctx.y();
82    /// assert_eq!(ctx.len(), 2);
83    /// ctx.clear();
84    /// assert_eq!(ctx.len(), 0);
85    /// ```
86    pub fn len(&self) -> usize {
87        self.ops.len()
88    }
89
90    /// Checks whether the context is empty
91    pub fn is_empty(&self) -> bool {
92        self.ops.is_empty()
93    }
94
95    /// Checks whether the given [`Node`] is valid in this context
96    fn check_node(&self, node: Node) -> Result<(), BadNode> {
97        self.get_op(node).ok_or(BadNode).map(|_| ())
98    }
99
100    /// Looks up the constant associated with the given node.
101    ///
102    /// If the node is invalid for this tree, returns an error; if the node is
103    /// not a constant, returns `Ok(None)`.
104    pub fn get_const(&self, n: Node) -> Result<f32, ConstError> {
105        match self.get_op(n) {
106            Some(Op::Const(c)) => Ok(c.0),
107            Some(_) => Err(ConstError::NotAConst),
108            None => Err(ConstError::BadNode(BadNode)),
109        }
110    }
111
112    /// Looks up the [`Var`] associated with the given node.
113    ///
114    /// If the node is invalid for this tree or not an `Op::Input`, returns an
115    /// error.
116    pub fn get_var(&self, n: Node) -> Result<Var, VarError> {
117        match self.get_op(n) {
118            Some(Op::Input(v)) => Ok(*v),
119            Some(..) => Err(VarError::NotAVar(NotAVar)),
120            None => Err(VarError::BadNode(BadNode)),
121        }
122    }
123
124    ////////////////////////////////////////////////////////////////////////////
125    // Primitives
126    /// Constructs or finds a [`Var::X`] node
127    /// ```
128    /// # use fidget_core::context::Context;
129    /// let mut ctx = Context::new();
130    /// let x = ctx.x();
131    /// let v = ctx.eval_xyz(x, 1.0, 0.0, 0.0).unwrap();
132    /// assert_eq!(v, 1.0);
133    /// ```
134    pub fn x(&mut self) -> Node {
135        self.var(Var::X)
136    }
137
138    /// Constructs or finds a [`Var::Y`] node
139    pub fn y(&mut self) -> Node {
140        self.var(Var::Y)
141    }
142
143    /// Constructs or finds a [`Var::Z`] node
144    pub fn z(&mut self) -> Node {
145        self.var(Var::Z)
146    }
147
148    /// Constructs or finds a variable input node
149    ///
150    /// To make an anonymous variable, call this function with [`Var::new()`]:
151    ///
152    /// ```
153    /// # use fidget_core::{context::Context, var::Var};
154    /// # use std::collections::HashMap;
155    /// let mut ctx = Context::new();
156    /// let v1 = ctx.var(Var::new());
157    /// let v2 = ctx.var(Var::new());
158    /// assert_ne!(v1, v2);
159    ///
160    /// let mut vars = HashMap::new();
161    /// vars.insert(ctx.get_var(v1).unwrap(), 3.0);
162    /// assert_eq!(ctx.eval(v1, &vars).unwrap(), 3.0);
163    /// assert!(ctx.eval(v2, &vars).is_err()); // v2 isn't in the map
164    /// ```
165    pub fn var(&mut self, v: Var) -> Node {
166        self.ops.insert(Op::Input(v))
167    }
168
169    /// Returns a 3-element array of `X`, `Y`, `Z` nodes
170    pub fn axes(&mut self) -> [Node; 3] {
171        [self.x(), self.y(), self.z()]
172    }
173
174    /// Returns a node representing the given constant value.
175    /// ```
176    /// # let mut ctx = fidget_core::context::Context::new();
177    /// let v = ctx.constant(3.0);
178    /// assert_eq!(ctx.eval_xyz(v, 0.0, 0.0, 0.0).unwrap(), 3.0);
179    /// ```
180    pub fn constant(&mut self, f: f32) -> Node {
181        self.ops.insert(Op::Const(OrderedFloat(f)))
182    }
183
184    ////////////////////////////////////////////////////////////////////////////
185    // Helper functions to create nodes with constant folding
186    /// Find or create a [Node] for the given unary operation, with constant
187    /// folding.
188    fn op_unary(&mut self, a: Node, op: UnaryOpcode) -> Result<Node, BadNode> {
189        let op_a = *self.get_op(a).ok_or(BadNode)?;
190        let out = if let Op::Const(a) = op_a {
191            self.constant(op.eval(a.0))
192        } else {
193            self.ops.insert(Op::Unary(op, a))
194        };
195        Ok(out)
196    }
197    /// Find or create a [Node] for the given binary operation, with constant
198    /// folding.
199    fn op_binary(
200        &mut self,
201        a: Node,
202        b: Node,
203        op: BinaryOpcode,
204    ) -> Result<Node, BadNode> {
205        let op_a = *self.get_op(a).ok_or(BadNode)?;
206        let op_b = *self.get_op(b).ok_or(BadNode)?;
207        let out = if let (Op::Const(a), Op::Const(b)) = (op_a, op_b) {
208            self.constant(op.eval(a.0, b.0))
209        } else {
210            self.ops.insert(Op::Binary(op, a, b))
211        };
212        Ok(out)
213    }
214
215    /// Find or create a [Node] for the given commutative operation, with
216    /// constant folding; deduplication is encouraged by sorting `a` and `b`.
217    fn op_binary_commutative(
218        &mut self,
219        a: Node,
220        b: Node,
221        op: BinaryOpcode,
222    ) -> Result<Node, BadNode> {
223        self.op_binary(a.min(b), a.max(b), op)
224    }
225
226    /// Builds an addition node
227    /// ```
228    /// # let mut ctx = fidget_core::context::Context::new();
229    /// let x = ctx.x();
230    /// let op = ctx.add(x, 1.0).unwrap();
231    /// let v = ctx.eval_xyz(op, 1.0, 0.0, 0.0).unwrap();
232    /// assert_eq!(v, 2.0);
233    /// ```
234    pub fn add<A: IntoNode, B: IntoNode>(
235        &mut self,
236        a: A,
237        b: B,
238    ) -> Result<Node, BadNode> {
239        let a: Node = a.into_node(self)?;
240        let b: Node = b.into_node(self)?;
241        if a == b {
242            let two = self.constant(2.0);
243            self.mul(a, two)
244        } else {
245            match (self.get_const(a), self.get_const(b)) {
246                (Ok(0.0), _) => Ok(b),
247                (_, Ok(0.0)) => Ok(a),
248                _ => self.op_binary_commutative(a, b, BinaryOpcode::Add),
249            }
250        }
251    }
252
253    /// Builds an multiplication node
254    /// ```
255    /// # let mut ctx = fidget_core::context::Context::new();
256    /// let x = ctx.x();
257    /// let op = ctx.mul(x, 5.0).unwrap();
258    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
259    /// assert_eq!(v, 10.0);
260    /// ```
261    pub fn mul<A: IntoNode, B: IntoNode>(
262        &mut self,
263        a: A,
264        b: B,
265    ) -> Result<Node, BadNode> {
266        let a = a.into_node(self)?;
267        let b = b.into_node(self)?;
268        if a == b {
269            self.square(a)
270        } else {
271            match (self.get_const(a), self.get_const(b)) {
272                (Ok(1.0), _) => Ok(b),
273                (_, Ok(1.0)) => Ok(a),
274                (Ok(0.0), _) => Ok(a),
275                (_, Ok(0.0)) => Ok(b),
276                _ => self.op_binary_commutative(a, b, BinaryOpcode::Mul),
277            }
278        }
279    }
280
281    /// Builds an `min` node
282    /// ```
283    /// # let mut ctx = fidget_core::context::Context::new();
284    /// let x = ctx.x();
285    /// let op = ctx.min(x, 5.0).unwrap();
286    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
287    /// assert_eq!(v, 2.0);
288    /// ```
289    pub fn min<A: IntoNode, B: IntoNode>(
290        &mut self,
291        a: A,
292        b: B,
293    ) -> Result<Node, BadNode> {
294        let a = a.into_node(self)?;
295        let b = b.into_node(self)?;
296        if a == b {
297            Ok(a)
298        } else {
299            self.op_binary_commutative(a, b, BinaryOpcode::Min)
300        }
301    }
302    /// Builds an `max` node
303    /// ```
304    /// # let mut ctx = fidget_core::context::Context::new();
305    /// let x = ctx.x();
306    /// let op = ctx.max(x, 5.0).unwrap();
307    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
308    /// assert_eq!(v, 5.0);
309    /// ```
310    pub fn max<A: IntoNode, B: IntoNode>(
311        &mut self,
312        a: A,
313        b: B,
314    ) -> Result<Node, BadNode> {
315        let a = a.into_node(self)?;
316        let b = b.into_node(self)?;
317        if a == b {
318            Ok(a)
319        } else {
320            self.op_binary_commutative(a, b, BinaryOpcode::Max)
321        }
322    }
323
324    /// Builds an `and` node
325    ///
326    /// If both arguments are non-zero, returns the right-hand argument.
327    /// Otherwise, returns zero.
328    ///
329    /// This node can be simplified using a tracing evaluator:
330    /// - If the left-hand argument is zero, simplify to just that argument
331    /// - If the left-hand argument is non-zero, simplify to the other argument
332    /// ```
333    /// # let mut ctx = fidget_core::context::Context::new();
334    /// let x = ctx.x();
335    /// let y = ctx.y();
336    /// let op = ctx.and(x, y).unwrap();
337    /// let v = ctx.eval_xyz(op, 1.0, 0.0, 0.0).unwrap();
338    /// assert_eq!(v, 0.0);
339    /// let v = ctx.eval_xyz(op, 1.0, 1.0, 0.0).unwrap();
340    /// assert_eq!(v, 1.0);
341    /// let v = ctx.eval_xyz(op, 1.0, 2.0, 0.0).unwrap();
342    /// assert_eq!(v, 2.0);
343    /// ```
344    pub fn and<A: IntoNode, B: IntoNode>(
345        &mut self,
346        a: A,
347        b: B,
348    ) -> Result<Node, BadNode> {
349        let a = a.into_node(self)?;
350        let b = b.into_node(self)?;
351
352        let op_a = *self.get_op(a).ok_or(BadNode)?;
353        if let Op::Const(v) = op_a {
354            if v.0 == 0.0 { Ok(a) } else { Ok(b) }
355        } else {
356            self.op_binary(a, b, BinaryOpcode::And)
357        }
358    }
359
360    /// Builds an `or` node
361    ///
362    /// If the left-hand argument is non-zero, it is returned.  Otherwise, the
363    /// right-hand argument is returned.
364    ///
365    /// This node can be simplified using a tracing evaluator.
366    /// ```
367    /// # let mut ctx = fidget_core::context::Context::new();
368    /// let x = ctx.x();
369    /// let y = ctx.y();
370    /// let op = ctx.or(x, y).unwrap();
371    /// let v = ctx.eval_xyz(op, 1.0, 0.0, 0.0).unwrap();
372    /// assert_eq!(v, 1.0);
373    /// let v = ctx.eval_xyz(op, 0.0, 0.0, 0.0).unwrap();
374    /// assert_eq!(v, 0.0);
375    /// let v = ctx.eval_xyz(op, 0.0, 3.0, 0.0).unwrap();
376    /// assert_eq!(v, 3.0);
377    /// ```
378    pub fn or<A: IntoNode, B: IntoNode>(
379        &mut self,
380        a: A,
381        b: B,
382    ) -> Result<Node, BadNode> {
383        let a = a.into_node(self)?;
384        let b = b.into_node(self)?;
385
386        let op_a = *self.get_op(a).ok_or(BadNode)?;
387        let op_b = *self.get_op(b).ok_or(BadNode)?;
388        if let Op::Const(v) = op_a {
389            if v.0 != 0.0 {
390                return Ok(a);
391            } else {
392                return Ok(b);
393            }
394        } else if let Op::Const(v) = op_b
395            && v.0 == 0.0
396        {
397            return Ok(a);
398        }
399        self.op_binary(a, b, BinaryOpcode::Or)
400    }
401
402    /// Builds a logical negation node
403    ///
404    /// The output is 1 if the argument is 0, and 0 otherwise.
405    pub fn not<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
406        let a = a.into_node(self)?;
407        self.op_unary(a, UnaryOpcode::Not)
408    }
409
410    /// Builds a unary negation node
411    /// ```
412    /// # let mut ctx = fidget_core::context::Context::new();
413    /// let x = ctx.x();
414    /// let op = ctx.neg(x).unwrap();
415    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
416    /// assert_eq!(v, -2.0);
417    /// ```
418    pub fn neg<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
419        let a = a.into_node(self)?;
420        self.op_unary(a, UnaryOpcode::Neg)
421    }
422
423    /// Builds a reciprocal node
424    /// ```
425    /// # let mut ctx = fidget_core::context::Context::new();
426    /// let x = ctx.x();
427    /// let op = ctx.recip(x).unwrap();
428    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
429    /// assert_eq!(v, 0.5);
430    /// ```
431    pub fn recip<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
432        let a = a.into_node(self)?;
433        self.op_unary(a, UnaryOpcode::Recip)
434    }
435
436    /// Builds a node which calculates the absolute value of its input
437    /// ```
438    /// # let mut ctx = fidget_core::context::Context::new();
439    /// let x = ctx.x();
440    /// let op = ctx.abs(x).unwrap();
441    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
442    /// assert_eq!(v, 2.0);
443    /// let v = ctx.eval_xyz(op, -2.0, 0.0, 0.0).unwrap();
444    /// assert_eq!(v, 2.0);
445    /// ```
446    pub fn abs<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
447        let a = a.into_node(self)?;
448        self.op_unary(a, UnaryOpcode::Abs)
449    }
450
451    /// Builds a node which calculates the square root of its input
452    /// ```
453    /// # let mut ctx = fidget_core::context::Context::new();
454    /// let x = ctx.x();
455    /// let op = ctx.sqrt(x).unwrap();
456    /// let v = ctx.eval_xyz(op, 4.0, 0.0, 0.0).unwrap();
457    /// assert_eq!(v, 2.0);
458    /// ```
459    pub fn sqrt<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
460        let a = a.into_node(self)?;
461        self.op_unary(a, UnaryOpcode::Sqrt)
462    }
463
464    /// Builds a node which calculates the sine of its input (in radians)
465    /// ```
466    /// # let mut ctx = fidget_core::context::Context::new();
467    /// let x = ctx.x();
468    /// let op = ctx.sin(x).unwrap();
469    /// let v = ctx.eval_xyz(op, std::f32::consts::PI / 2.0, 0.0, 0.0).unwrap();
470    /// assert_eq!(v, 1.0);
471    /// ```
472    pub fn sin<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
473        let a = a.into_node(self)?;
474        self.op_unary(a, UnaryOpcode::Sin)
475    }
476
477    /// Builds a node which calculates the cosine of its input (in radians)
478    pub fn cos<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
479        let a = a.into_node(self)?;
480        self.op_unary(a, UnaryOpcode::Cos)
481    }
482
483    /// Builds a node which calculates the tangent of its input (in radians)
484    pub fn tan<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
485        let a = a.into_node(self)?;
486        self.op_unary(a, UnaryOpcode::Tan)
487    }
488
489    /// Builds a node which calculates the arcsine of its input (in radians)
490    pub fn asin<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
491        let a = a.into_node(self)?;
492        self.op_unary(a, UnaryOpcode::Asin)
493    }
494
495    /// Builds a node which calculates the arccosine of its input (in radians)
496    pub fn acos<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
497        let a = a.into_node(self)?;
498        self.op_unary(a, UnaryOpcode::Acos)
499    }
500
501    /// Builds a node which calculates the arctangent of its input (in radians)
502    pub fn atan<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
503        let a = a.into_node(self)?;
504        self.op_unary(a, UnaryOpcode::Atan)
505    }
506
507    /// Builds a node which calculates the exponent of its input
508    pub fn exp<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
509        let a = a.into_node(self)?;
510        self.op_unary(a, UnaryOpcode::Exp)
511    }
512
513    /// Builds a node which calculates the natural log of its input
514    pub fn ln<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
515        let a = a.into_node(self)?;
516        self.op_unary(a, UnaryOpcode::Ln)
517    }
518
519    ////////////////////////////////////////////////////////////////////////////
520    // Derived functions
521    /// Builds a node which squares its input
522    /// ```
523    /// # let mut ctx = fidget_core::context::Context::new();
524    /// let x = ctx.x();
525    /// let op = ctx.square(x).unwrap();
526    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
527    /// assert_eq!(v, 4.0);
528    /// ```
529    pub fn square<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
530        let a = a.into_node(self)?;
531        self.op_unary(a, UnaryOpcode::Square)
532    }
533
534    /// Builds a node which takes the floor of its input
535    /// ```
536    /// # let mut ctx = fidget_core::context::Context::new();
537    /// let x = ctx.x();
538    /// let op = ctx.floor(x).unwrap();
539    /// let v = ctx.eval_xyz(op, 1.2, 0.0, 0.0).unwrap();
540    /// assert_eq!(v, 1.0);
541    /// ```
542    pub fn floor<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
543        let a = a.into_node(self)?;
544        self.op_unary(a, UnaryOpcode::Floor)
545    }
546
547    /// Builds a node which takes the ceiling of its input
548    /// ```
549    /// # let mut ctx = fidget_core::context::Context::new();
550    /// let x = ctx.x();
551    /// let op = ctx.ceil(x).unwrap();
552    /// let v = ctx.eval_xyz(op, 1.2, 0.0, 0.0).unwrap();
553    /// assert_eq!(v, 2.0);
554    /// ```
555    pub fn ceil<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
556        let a = a.into_node(self)?;
557        self.op_unary(a, UnaryOpcode::Ceil)
558    }
559
560    /// Builds a node which rounds its input to the nearest integer
561    /// ```
562    /// # let mut ctx = fidget_core::context::Context::new();
563    /// let x = ctx.x();
564    /// let op = ctx.round(x).unwrap();
565    /// let v = ctx.eval_xyz(op, 1.2, 0.0, 0.0).unwrap();
566    /// assert_eq!(v, 1.0);
567    /// let v = ctx.eval_xyz(op, 1.6, 0.0, 0.0).unwrap();
568    /// assert_eq!(v, 2.0);
569    /// let v = ctx.eval_xyz(op, 1.5, 0.0, 0.0).unwrap();
570    /// assert_eq!(v, 2.0); // rounds away from 0.0 if ambiguous
571    /// ```
572    pub fn round<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
573        let a = a.into_node(self)?;
574        self.op_unary(a, UnaryOpcode::Round)
575    }
576
577    /// Builds a node which performs subtraction.
578    /// ```
579    /// # let mut ctx = fidget_core::context::Context::new();
580    /// let x = ctx.x();
581    /// let y = ctx.y();
582    /// let op = ctx.sub(x, y).unwrap();
583    /// let v = ctx.eval_xyz(op, 3.0, 2.0, 0.0).unwrap();
584    /// assert_eq!(v, 1.0);
585    /// ```
586    pub fn sub<A: IntoNode, B: IntoNode>(
587        &mut self,
588        a: A,
589        b: B,
590    ) -> Result<Node, BadNode> {
591        let a = a.into_node(self)?;
592        let b = b.into_node(self)?;
593
594        match (self.get_const(a), self.get_const(b)) {
595            (Ok(0.0), _) => self.neg(b),
596            (_, Ok(0.0)) => Ok(a),
597            _ => self.op_binary(a, b, BinaryOpcode::Sub),
598        }
599    }
600
601    /// Builds a node which performs division.
602    /// ```
603    /// # let mut ctx = fidget_core::context::Context::new();
604    /// let x = ctx.x();
605    /// let y = ctx.y();
606    /// let op = ctx.div(x, y).unwrap();
607    /// let v = ctx.eval_xyz(op, 3.0, 2.0, 0.0).unwrap();
608    /// assert_eq!(v, 1.5);
609    /// ```
610    pub fn div<A: IntoNode, B: IntoNode>(
611        &mut self,
612        a: A,
613        b: B,
614    ) -> Result<Node, BadNode> {
615        let a = a.into_node(self)?;
616        let b = b.into_node(self)?;
617
618        match (self.get_const(a), self.get_const(b)) {
619            (Ok(0.0), _) => Ok(a),
620            (_, Ok(1.0)) => Ok(a),
621            _ => self.op_binary(a, b, BinaryOpcode::Div),
622        }
623    }
624
625    /// Builds a node which computes `atan2(y, x)`
626    /// ```
627    /// # let mut ctx = fidget_core::context::Context::new();
628    /// let x = ctx.x();
629    /// let y = ctx.y();
630    /// let op = ctx.atan2(y, x).unwrap();
631    /// let v = ctx.eval_xyz(op, 0.0, 1.0, 0.0).unwrap();
632    /// assert_eq!(v, std::f32::consts::FRAC_PI_2);
633    /// ```
634    pub fn atan2<A: IntoNode, B: IntoNode>(
635        &mut self,
636        y: A,
637        x: B,
638    ) -> Result<Node, BadNode> {
639        let y = y.into_node(self)?;
640        let x = x.into_node(self)?;
641
642        self.op_binary(y, x, BinaryOpcode::Atan)
643    }
644
645    /// Builds a node that compares two values
646    ///
647    /// The result is -1 if `a < b`, +1 if `a > b`, 0 if `a == b`, and `NaN` if
648    /// either side is `NaN`.
649    /// ```
650    /// # let mut ctx = fidget_core::context::Context::new();
651    /// let x = ctx.x();
652    /// let op = ctx.compare(x, 1.0).unwrap();
653    /// let v = ctx.eval_xyz(op, 0.0, 0.0, 0.0).unwrap();
654    /// assert_eq!(v, -1.0);
655    /// let v = ctx.eval_xyz(op, 2.0, 0.0, 0.0).unwrap();
656    /// assert_eq!(v, 1.0);
657    /// let v = ctx.eval_xyz(op, 1.0, 0.0, 0.0).unwrap();
658    /// assert_eq!(v, 0.0);
659    /// ```
660    pub fn compare<A: IntoNode, B: IntoNode>(
661        &mut self,
662        a: A,
663        b: B,
664    ) -> Result<Node, BadNode> {
665        let a = a.into_node(self)?;
666        let b = b.into_node(self)?;
667        self.op_binary(a, b, BinaryOpcode::Compare)
668    }
669
670    /// Builds a node which performs pseudo-random mixing of two values
671    pub fn mix<A: IntoNode, B: IntoNode>(
672        &mut self,
673        a: A,
674        b: B,
675    ) -> Result<Node, BadNode> {
676        let a = a.into_node(self)?;
677        let b = b.into_node(self)?;
678        self.op_binary(a, b, BinaryOpcode::Mix)
679    }
680
681    /// Builds a node that returns a pseudo-random value
682    pub fn rand<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
683        let a = a.into_node(self)?;
684        self.op_unary(a, UnaryOpcode::Rand)
685    }
686
687    /// Builds a node that is 1 if `lhs < rhs` and 0 otherwise
688    ///
689    /// ```
690    /// # let mut ctx = fidget_core::context::Context::new();
691    /// let x = ctx.x();
692    /// let y = ctx.y();
693    /// let op = ctx.less_than(x, y).unwrap();
694    /// let v = ctx.eval_xyz(op, 0.0, 1.0, 0.0).unwrap();
695    /// assert_eq!(v, 1.0);
696    /// let v = ctx.eval_xyz(op, 1.0, 1.0, 0.0).unwrap();
697    /// assert_eq!(v, 0.0);
698    /// let v = ctx.eval_xyz(op, 2.0, 1.0, 0.0).unwrap();
699    /// assert_eq!(v, 0.0);
700    /// ```
701    pub fn less_than<A: IntoNode, B: IntoNode>(
702        &mut self,
703        lhs: A,
704        rhs: B,
705    ) -> Result<Node, BadNode> {
706        let lhs = lhs.into_node(self)?;
707        let rhs = rhs.into_node(self)?;
708        let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
709        self.max(cmp, 0.0)
710    }
711
712    /// Builds a node that is 1 if `lhs <= rhs` and 0 otherwise
713    ///
714    /// ```
715    /// # let mut ctx = fidget_core::context::Context::new();
716    /// let x = ctx.x();
717    /// let y = ctx.y();
718    /// let op = ctx.less_than_or_equal(x, y).unwrap();
719    /// let v = ctx.eval_xyz(op, 0.0, 1.0, 0.0).unwrap();
720    /// assert_eq!(v, 1.0);
721    /// let v = ctx.eval_xyz(op, 1.0, 1.0, 0.0).unwrap();
722    /// assert_eq!(v, 1.0);
723    /// let v = ctx.eval_xyz(op, 2.0, 1.0, 0.0).unwrap();
724    /// assert_eq!(v, 0.0);
725    /// ```
726    pub fn less_than_or_equal<A: IntoNode, B: IntoNode>(
727        &mut self,
728        lhs: A,
729        rhs: B,
730    ) -> Result<Node, BadNode> {
731        let lhs = lhs.into_node(self)?;
732        let rhs = rhs.into_node(self)?;
733        let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
734        let shift = self.add(cmp, 1.0)?;
735        self.min(shift, 1.0)
736    }
737
738    /// Builds a node that takes the modulo (least non-negative remainder)
739    pub fn modulo<A: IntoNode, B: IntoNode>(
740        &mut self,
741        a: A,
742        b: B,
743    ) -> Result<Node, BadNode> {
744        let a = a.into_node(self)?;
745        let b = b.into_node(self)?;
746        self.op_binary(a, b, BinaryOpcode::Mod)
747    }
748
749    /// Builds a node that returns the first node if the condition is not
750    /// equal to zero, else returns the other node
751    ///
752    /// The result is `a` if `condition != 0`, else the result is `b`.
753    /// ```
754    /// # let mut ctx = fidget_core::context::Context::new();
755    /// let x = ctx.x();
756    /// let y = ctx.y();
757    /// let z = ctx.z();
758    ///
759    /// let if_else = ctx.if_nonzero_else(x, y, z).unwrap();
760    ///
761    /// assert_eq!(ctx.eval_xyz(if_else, 0.0, 2.0, 3.0).unwrap(), 3.0);
762    /// assert_eq!(ctx.eval_xyz(if_else, 1.0, 2.0, 3.0).unwrap(), 2.0);
763    /// assert_eq!(ctx.eval_xyz(if_else, 0.0, f32::NAN, 3.0).unwrap(), 3.0);
764    /// assert_eq!(ctx.eval_xyz(if_else, 1.0, 2.0, f32::NAN).unwrap(), 2.0);
765    /// ```
766    pub fn if_nonzero_else<Condition: IntoNode, A: IntoNode, B: IntoNode>(
767        &mut self,
768        condition: Condition,
769        a: A,
770        b: B,
771    ) -> Result<Node, BadNode> {
772        let condition = condition.into_node(self)?;
773        let a = a.into_node(self)?;
774        let b = b.into_node(self)?;
775
776        let lhs = self.and(condition, a)?;
777        let n_condition = self.not(condition)?;
778        let rhs = self.and(n_condition, b)?;
779        self.or(lhs, rhs)
780    }
781
782    ////////////////////////////////////////////////////////////////////////////
783    /// Evaluates the given node with the provided values for X, Y, and Z.
784    ///
785    /// This is extremely inefficient; consider converting the node into a
786    /// [`Shape`](crate::shape::Shape) and using its evaluators instead.
787    ///
788    /// ```
789    /// # let mut ctx = fidget_core::context::Context::new();
790    /// let x = ctx.x();
791    /// let y = ctx.y();
792    /// let z = ctx.z();
793    /// let op = ctx.mul(x, y).unwrap();
794    /// let op = ctx.div(op, z).unwrap();
795    /// let v = ctx.eval_xyz(op, 3.0, 5.0, 2.0).unwrap();
796    /// assert_eq!(v, 7.5); // (3.0 * 5.0) / 2.0
797    /// ```
798    pub fn eval_xyz(
799        &self,
800        root: Node,
801        x: f32,
802        y: f32,
803        z: f32,
804    ) -> Result<f32, EvalError> {
805        let vars = [(Var::X, x), (Var::Y, y), (Var::Z, z)]
806            .into_iter()
807            .collect();
808        self.eval(root, &vars)
809    }
810
811    /// Evaluates the given node with a generic set of variables
812    ///
813    /// This is extremely inefficient; consider converting the node into a
814    /// [`Shape`](crate::shape::Shape) and using its evaluators instead.
815    pub fn eval(
816        &self,
817        root: Node,
818        vars: &HashMap<Var, f32>,
819    ) -> Result<f32, EvalError> {
820        let mut cache = vec![None; self.ops.len()].into();
821        self.eval_inner(root, vars, &mut cache)
822    }
823
824    fn eval_inner(
825        &self,
826        node: Node,
827        vars: &HashMap<Var, f32>,
828        cache: &mut IndexVec<Option<f32>, Node>,
829    ) -> Result<f32, EvalError> {
830        if node.0 >= cache.len() {
831            return Err(EvalError::BadNode(BadNode));
832        }
833        if let Some(v) = cache[node] {
834            return Ok(v);
835        }
836        let mut get = |n: Node| self.eval_inner(n, vars, cache);
837        let v = match self.get_op(node).ok_or(EvalError::BadNode(BadNode))? {
838            Op::Input(v) => *vars.get(v).ok_or(EvalError::MissingVar(*v))?,
839            Op::Const(c) => c.0,
840
841            Op::Binary(op, a, b) => {
842                let a = get(*a)?;
843                let b = get(*b)?;
844                op.eval(a, b)
845            }
846
847            // Unary operations
848            Op::Unary(op, a) => {
849                let a = get(*a)?;
850                op.eval(a)
851            }
852        };
853
854        cache[node] = Some(v);
855        Ok(v)
856    }
857
858    /// Parses a flat text representation of a math tree. For example, the
859    /// circle `(- (+ (square x) (square y)) 1)` can be parsed from
860    /// ```
861    /// # use fidget_core::context::Context;
862    /// let txt = "
863    /// ## This is a comment!
864    /// 0x600000b90000 var-x
865    /// 0x600000b900a0 square 0x600000b90000
866    /// 0x600000b90050 var-y
867    /// 0x600000b900f0 square 0x600000b90050
868    /// 0x600000b90140 add 0x600000b900a0 0x600000b900f0
869    /// 0x600000b90190 sqrt 0x600000b90140
870    /// 0x600000b901e0 const 1
871    /// ";
872    /// let (ctx, _node) = Context::from_text(&mut txt.as_bytes()).unwrap();
873    /// assert_eq!(ctx.len(), 7);
874    /// ```
875    ///
876    /// This representation is loosely defined and only intended for use in
877    /// quick experiments.
878    pub fn from_text<R: Read>(r: R) -> Result<(Self, Node), ParseError> {
879        let reader = BufReader::new(r);
880        let mut ctx = Self::new();
881        let mut seen = BTreeMap::new();
882        let mut last = None;
883
884        for line in reader.lines().map(|line| line.unwrap()) {
885            if line.is_empty() || line.starts_with('#') {
886                continue;
887            }
888            let mut iter = line.split_whitespace();
889            let i: String = iter.next().unwrap().to_owned();
890            let opcode = iter.next().unwrap();
891
892            let mut pop = || {
893                let txt = iter.next().unwrap();
894                seen.get(txt)
895                    .cloned()
896                    .ok_or_else(|| ParseError::UnknownVariable(txt.to_string()))
897            };
898            let node = match opcode {
899                "const" => ctx.constant(iter.next().unwrap().parse().unwrap()),
900                "var-x" => ctx.x(),
901                "var-y" => ctx.y(),
902                "var-z" => ctx.z(),
903                "abs" => ctx.abs(pop()?)?,
904                "neg" => ctx.neg(pop()?)?,
905                "sqrt" => ctx.sqrt(pop()?)?,
906                "square" => ctx.square(pop()?)?,
907                "floor" => ctx.floor(pop()?)?,
908                "ceil" => ctx.ceil(pop()?)?,
909                "round" => ctx.round(pop()?)?,
910                "sin" => ctx.sin(pop()?)?,
911                "cos" => ctx.cos(pop()?)?,
912                "tan" => ctx.tan(pop()?)?,
913                "asin" => ctx.asin(pop()?)?,
914                "acos" => ctx.acos(pop()?)?,
915                "atan" => ctx.atan(pop()?)?,
916                "ln" => ctx.ln(pop()?)?,
917                "not" => ctx.not(pop()?)?,
918                "rand" => ctx.rand(pop()?)?,
919                "exp" => ctx.exp(pop()?)?,
920                "add" => ctx.add(pop()?, pop()?)?,
921                "mul" => ctx.mul(pop()?, pop()?)?,
922                "min" => ctx.min(pop()?, pop()?)?,
923                "max" => ctx.max(pop()?, pop()?)?,
924                "div" => ctx.div(pop()?, pop()?)?,
925                "atan2" => ctx.atan2(pop()?, pop()?)?,
926                "sub" => ctx.sub(pop()?, pop()?)?,
927                "compare" => ctx.compare(pop()?, pop()?)?,
928                "mod" => ctx.modulo(pop()?, pop()?)?,
929                "and" => ctx.and(pop()?, pop()?)?,
930                "or" => ctx.or(pop()?, pop()?)?,
931                "mix" => ctx.mix(pop()?, pop()?)?,
932                op => return Err(ParseError::UnknownOpcode(op.to_owned())),
933            };
934            seen.insert(i, node);
935            last = Some(node);
936        }
937        match last {
938            Some(node) => Ok((ctx, node)),
939            None => Err(ParseError::EmptyFile),
940        }
941    }
942
943    /// Converts the entire context into a GraphViz drawing
944    pub fn dot(&self) -> String {
945        let mut out = "digraph mygraph{\n".to_owned();
946        for node in self.ops.keys() {
947            let op = self.get_op(node).unwrap();
948            out += &self.dot_node(node);
949            out += &op.dot_edges(node);
950        }
951        out += "}\n";
952        out
953    }
954
955    /// Converts the given node into a GraphViz node
956    ///
957    /// (this is a local function instead of a function on `Op` because it
958    ///  requires looking up variables by name)
959    fn dot_node(&self, i: Node) -> String {
960        let mut out = format!(r#"n{} [label = ""#, i.get());
961        let op = self.get_op(i).unwrap();
962        match op {
963            Op::Const(c) => write!(out, "{c}").unwrap(),
964            Op::Input(v) => {
965                out += &v.to_string();
966            }
967            Op::Binary(op, ..) => match op {
968                BinaryOpcode::Add => out += "add",
969                BinaryOpcode::Sub => out += "sub",
970                BinaryOpcode::Mul => out += "mul",
971                BinaryOpcode::Div => out += "div",
972                BinaryOpcode::Atan => out += "atan2",
973                BinaryOpcode::Min => out += "min",
974                BinaryOpcode::Max => out += "max",
975                BinaryOpcode::Compare => out += "compare",
976                BinaryOpcode::Mod => out += "mod",
977                BinaryOpcode::And => out += "and",
978                BinaryOpcode::Or => out += "or",
979                BinaryOpcode::Mix => out += "mix",
980            },
981            Op::Unary(op, ..) => match op {
982                UnaryOpcode::Neg => out += "neg",
983                UnaryOpcode::Abs => out += "abs",
984                UnaryOpcode::Recip => out += "recip",
985                UnaryOpcode::Sqrt => out += "sqrt",
986                UnaryOpcode::Square => out += "square",
987                UnaryOpcode::Floor => out += "floor",
988                UnaryOpcode::Ceil => out += "ceil",
989                UnaryOpcode::Round => out += "round",
990                UnaryOpcode::Sin => out += "sin",
991                UnaryOpcode::Cos => out += "cos",
992                UnaryOpcode::Tan => out += "tan",
993                UnaryOpcode::Asin => out += "asin",
994                UnaryOpcode::Acos => out += "acos",
995                UnaryOpcode::Atan => out += "atan",
996                UnaryOpcode::Exp => out += "exp",
997                UnaryOpcode::Ln => out += "ln",
998                UnaryOpcode::Not => out += "not",
999                UnaryOpcode::Rand => out += "rand",
1000            },
1001        };
1002        write!(
1003            out,
1004            r#"" color="{0}1" shape="{1}" fontcolor="{0}4"]"#,
1005            op.dot_node_color(),
1006            op.dot_node_shape()
1007        )
1008        .unwrap();
1009        out
1010    }
1011
1012    /// Looks up an operation by `Node` handle
1013    pub fn get_op(&self, node: Node) -> Option<&Op> {
1014        self.ops.get_by_index(node)
1015    }
1016
1017    /// Imports the given tree, deduplicating and returning the root
1018    pub fn import(&mut self, tree: &Tree) -> Node {
1019        // A naive remapping implementation would use recursion.  A naive
1020        // remapping implementation would blow up the stack given any
1021        // significant tree size.
1022        //
1023        // Instead, we maintain our own pseudo-stack here in a pair of Vecs (one
1024        // stack for actions, and a second stack for return values).
1025        enum Action<'a> {
1026            /// Pushes `Up(op)` followed by `Down(c)` for each child
1027            Down(&'a Arc<TreeOp>),
1028            /// Consumes imported trees from the stack and pushes a new tree
1029            Up(&'a Arc<TreeOp>),
1030            /// Pops the latest axis frame
1031            Pop,
1032            /// Pops the latest affine frame
1033            PopAffine,
1034        }
1035        let mut axes = vec![(self.x(), self.y(), self.z())];
1036        let mut todo = vec![Action::Down(tree.arc())];
1037        let mut stack = vec![];
1038        let mut affine: Vec<Matrix4<f32>> = vec![];
1039
1040        // Cache of TreeOp -> Node mapping under a particular frame (axes)
1041        //
1042        // This isn't required for correctness, but can be a speed optimization
1043        // (because it means we don't have to walk the same tree twice).
1044        let mut seen = HashMap::new();
1045
1046        while let Some(t) = todo.pop() {
1047            match t {
1048                Action::Down(t) => {
1049                    // If we've already seen this TreeOp with these axes, then
1050                    // we can return the previous Node.
1051                    if matches!(
1052                        t.as_ref(),
1053                        TreeOp::Unary(..) | TreeOp::Binary(..)
1054                    ) && let Some(p) =
1055                        seen.get(&(*axes.last().unwrap(), Arc::as_ptr(t)))
1056                    {
1057                        stack.push(*p);
1058                        continue;
1059                    }
1060                    match t.as_ref() {
1061                        TreeOp::Const(c) => {
1062                            stack.push(self.constant(*c));
1063                        }
1064                        TreeOp::Input(s) => {
1065                            let axes = axes.last().unwrap();
1066                            stack.push(match *s {
1067                                Var::X => axes.0,
1068                                Var::Y => axes.1,
1069                                Var::Z => axes.2,
1070                                v @ Var::V(..) => self.var(v),
1071                            });
1072                        }
1073                        TreeOp::Unary(_op, arg) => {
1074                            todo.push(Action::Up(t));
1075                            todo.push(Action::Down(arg));
1076                        }
1077                        TreeOp::Binary(_op, lhs, rhs) => {
1078                            todo.push(Action::Up(t));
1079                            todo.push(Action::Down(lhs));
1080                            todo.push(Action::Down(rhs));
1081                        }
1082                        TreeOp::RemapAxes { target: _, x, y, z } => {
1083                            // Action::Up(t) does the remapping and target eval
1084                            todo.push(Action::Up(t));
1085                            todo.push(Action::Down(x));
1086                            todo.push(Action::Down(y));
1087                            todo.push(Action::Down(z));
1088                        }
1089                        TreeOp::RemapAffine { target, mat } => {
1090                            let prev = affine
1091                                .last()
1092                                .cloned()
1093                                .unwrap_or(Matrix4::identity());
1094                            let mat = prev * mat.to_homogeneous();
1095
1096                            // Push either an affine frame or an axis frame,
1097                            // depending on whether the target is also affine
1098                            if matches!(&**target, TreeOp::RemapAffine { .. }) {
1099                                affine.push(mat);
1100                                todo.push(Action::PopAffine);
1101                            } else {
1102                                let (x, y, z) = axes.last().unwrap();
1103                                let mut out = [None; 3];
1104                                for i in 0..3 {
1105                                    let a = self.mul(mat[(i, 0)], *x).unwrap();
1106                                    let b = self.mul(mat[(i, 1)], *y).unwrap();
1107                                    let c = self.mul(mat[(i, 2)], *z).unwrap();
1108                                    let d = self.constant(mat[(i, 3)]);
1109                                    let ab = self.add(a, b).unwrap();
1110                                    let cd = self.add(c, d).unwrap();
1111                                    out[i] = Some(self.add(ab, cd).unwrap());
1112                                }
1113                                let [x, y, z] = out.map(Option::unwrap);
1114                                axes.push((x, y, z));
1115                                todo.push(Action::Pop);
1116                            }
1117                            todo.push(Action::Down(target));
1118                        }
1119                    }
1120                }
1121                Action::Up(t) => {
1122                    match t.as_ref() {
1123                        TreeOp::Const(..)
1124                        | TreeOp::Input(..)
1125                        | TreeOp::RemapAffine { .. } => unreachable!(),
1126                        TreeOp::Unary(op, ..) => {
1127                            let arg = stack.pop().unwrap();
1128                            let out = self.op_unary(arg, *op).unwrap();
1129                            stack.push(out);
1130                        }
1131                        TreeOp::Binary(op, ..) => {
1132                            let lhs = stack.pop().unwrap();
1133                            let rhs = stack.pop().unwrap();
1134                            // Call individual builders to apply optimizations
1135                            let out = match op {
1136                                BinaryOpcode::Add => self.add(lhs, rhs),
1137                                BinaryOpcode::Sub => self.sub(lhs, rhs),
1138                                BinaryOpcode::Mul => self.mul(lhs, rhs),
1139                                BinaryOpcode::Div => self.div(lhs, rhs),
1140                                BinaryOpcode::Atan => self.atan2(lhs, rhs),
1141                                BinaryOpcode::Min => self.min(lhs, rhs),
1142                                BinaryOpcode::Max => self.max(lhs, rhs),
1143                                BinaryOpcode::Compare => self.compare(lhs, rhs),
1144                                BinaryOpcode::Mix => self.mix(lhs, rhs),
1145                                BinaryOpcode::Mod => self.modulo(lhs, rhs),
1146                                BinaryOpcode::And => self.and(lhs, rhs),
1147                                BinaryOpcode::Or => self.or(lhs, rhs),
1148                            }
1149                            .unwrap();
1150                            if Arc::strong_count(t) > 1 {
1151                                seen.insert(
1152                                    (*axes.last().unwrap(), Arc::as_ptr(t)),
1153                                    out,
1154                                );
1155                            }
1156                            stack.push(out);
1157                        }
1158                        TreeOp::RemapAxes { target, .. } => {
1159                            let x = stack.pop().unwrap();
1160                            let y = stack.pop().unwrap();
1161                            let z = stack.pop().unwrap();
1162                            axes.push((x, y, z));
1163                            todo.push(Action::Pop);
1164                            todo.push(Action::Down(target));
1165                        }
1166                    }
1167                    // Update the cache with the new tree, if relevant
1168                    //
1169                    // The `strong_count` check is a rough heuristic to avoid
1170                    // caching if there's only a single owner of the tree.  This
1171                    // isn't perfect, but it doesn't need to be for correctness.
1172                    if matches!(
1173                        t.as_ref(),
1174                        TreeOp::Unary(..) | TreeOp::Binary(..)
1175                    ) && Arc::strong_count(t) > 1
1176                    {
1177                        seen.insert(
1178                            (*axes.last().unwrap(), Arc::as_ptr(t)),
1179                            *stack.last().unwrap(),
1180                        );
1181                    }
1182                }
1183                Action::Pop => {
1184                    axes.pop().unwrap();
1185                }
1186                Action::PopAffine => {
1187                    affine.pop().unwrap();
1188                }
1189            }
1190        }
1191        assert_eq!(stack.len(), 1);
1192        stack.pop().unwrap()
1193    }
1194
1195    /// Converts from a context-specific node into a standalone [`Tree`]
1196    pub fn export(&self, n: Node) -> Result<Tree, BadNode> {
1197        if self.get_op(n).is_none() {
1198            return Err(BadNode);
1199        }
1200
1201        // Do recursion on the heap to avoid stack overflows for deep trees
1202        enum Action {
1203            /// Pushes `Up(n)` followed by `Down(n)` for each child
1204            Down(Node),
1205            /// Consumes trees from the stack and pushes a new tree
1206            Up(Node, Op),
1207        }
1208        let mut todo = vec![Action::Down(n)];
1209        let mut stack = vec![];
1210
1211        // Cache of Node -> Tree mapping, for Tree deduplication
1212        let mut seen: HashMap<Node, Tree> = HashMap::new();
1213
1214        while let Some(t) = todo.pop() {
1215            match t {
1216                Action::Down(n) => {
1217                    // If we've already seen this TreeOp with these axes, then
1218                    // we can return the previous Node.
1219                    if let Some(p) = seen.get(&n) {
1220                        stack.push(p.clone());
1221                        continue;
1222                    }
1223                    let op = self.get_op(n).unwrap();
1224                    match op {
1225                        Op::Const(c) => {
1226                            let t = Tree::from(c.0);
1227                            seen.insert(n, t.clone());
1228                            stack.push(t);
1229                        }
1230                        Op::Input(v) => {
1231                            let t = Tree::from(*v);
1232                            seen.insert(n, t.clone());
1233                            stack.push(t);
1234                        }
1235                        Op::Unary(_op, arg) => {
1236                            todo.push(Action::Up(n, *op));
1237                            todo.push(Action::Down(*arg));
1238                        }
1239                        Op::Binary(_op, lhs, rhs) => {
1240                            todo.push(Action::Up(n, *op));
1241                            todo.push(Action::Down(*lhs));
1242                            todo.push(Action::Down(*rhs));
1243                        }
1244                    }
1245                }
1246                Action::Up(n, op) => match op {
1247                    Op::Const(..) | Op::Input(..) => unreachable!(),
1248                    Op::Unary(op, ..) => {
1249                        let arg = stack.pop().unwrap();
1250                        let out =
1251                            Tree::from(TreeOp::Unary(op, arg.arc().clone()));
1252                        seen.insert(n, out.clone());
1253                        stack.push(out);
1254                    }
1255                    Op::Binary(op, ..) => {
1256                        let lhs = stack.pop().unwrap();
1257                        let rhs = stack.pop().unwrap();
1258                        let out = Tree::from(TreeOp::Binary(
1259                            op,
1260                            lhs.arc().clone(),
1261                            rhs.arc().clone(),
1262                        ));
1263                        seen.insert(n, out.clone());
1264                        stack.push(out);
1265                    }
1266                },
1267            }
1268        }
1269        assert_eq!(stack.len(), 1);
1270        Ok(stack.pop().unwrap())
1271    }
1272
1273    /// Takes the symbolic derivative of a node with respect to a variable
1274    pub fn deriv(&mut self, n: Node, v: Var) -> Result<Node, BadNode> {
1275        if self.get_op(n).is_none() {
1276            return Err(BadNode);
1277        }
1278
1279        // Do recursion on the heap to avoid stack overflows for deep trees
1280        enum Action {
1281            /// Pushes `Up(n)` followed by `Down(n)` for each child
1282            Down(Node),
1283            /// Consumes trees from the stack and pushes a new tree
1284            Up(Node, Op),
1285        }
1286        let mut todo = vec![Action::Down(n)];
1287        let mut stack = vec![];
1288        let zero = self.constant(0.0);
1289
1290        // Cache of Node -> Node mapping, for deduplication
1291        let mut seen: HashMap<Node, Node> = HashMap::new();
1292
1293        while let Some(t) = todo.pop() {
1294            match t {
1295                Action::Down(n) => {
1296                    // If we've already seen this TreeOp with these axes, then
1297                    // we can return the previous Node.
1298                    if let Some(p) = seen.get(&n) {
1299                        stack.push(*p);
1300                        continue;
1301                    }
1302                    let op = *self.get_op(n).unwrap();
1303                    match op {
1304                        Op::Const(_c) => {
1305                            seen.insert(n, zero);
1306                            stack.push(zero);
1307                        }
1308                        Op::Input(u) => {
1309                            let z =
1310                                if v == u { self.constant(1.0) } else { zero };
1311                            seen.insert(n, z);
1312                            stack.push(z);
1313                        }
1314                        Op::Unary(_op, arg) => {
1315                            todo.push(Action::Up(n, op));
1316                            todo.push(Action::Down(arg));
1317                        }
1318                        Op::Binary(_op, lhs, rhs) => {
1319                            todo.push(Action::Up(n, op));
1320                            todo.push(Action::Down(lhs));
1321                            todo.push(Action::Down(rhs));
1322                        }
1323                    }
1324                }
1325                Action::Up(n, op) => match op {
1326                    Op::Const(..) | Op::Input(..) => unreachable!(),
1327                    Op::Unary(op, v_arg) => {
1328                        let d_arg = stack.pop().unwrap();
1329                        let out = match op {
1330                            UnaryOpcode::Neg => self.neg(d_arg),
1331                            UnaryOpcode::Abs => {
1332                                let cond = self.less_than(v_arg, zero).unwrap();
1333                                let pos = d_arg;
1334                                let neg = self.neg(d_arg).unwrap();
1335                                self.if_nonzero_else(cond, neg, pos)
1336                            }
1337                            UnaryOpcode::Recip => {
1338                                let a = self.square(v_arg).unwrap();
1339                                let b = self.neg(d_arg).unwrap();
1340                                self.div(b, a)
1341                            }
1342                            UnaryOpcode::Sqrt => {
1343                                let v = self.mul(n, 2.0).unwrap();
1344                                self.div(d_arg, v)
1345                            }
1346                            UnaryOpcode::Square => {
1347                                let v = self.mul(d_arg, v_arg).unwrap();
1348                                self.mul(2.0, v)
1349                            }
1350
1351                            // Discontinuous operations don't have Dirac deltas
1352                            UnaryOpcode::Floor
1353                            | UnaryOpcode::Ceil
1354                            | UnaryOpcode::Rand
1355                            | UnaryOpcode::Round => Ok(zero),
1356
1357                            UnaryOpcode::Sin => {
1358                                let c = self.cos(v_arg).unwrap();
1359                                self.mul(c, d_arg)
1360                            }
1361
1362                            UnaryOpcode::Cos => {
1363                                let s = self.sin(v_arg).unwrap();
1364                                let s = self.neg(s).unwrap();
1365                                self.mul(s, d_arg)
1366                            }
1367
1368                            UnaryOpcode::Tan => {
1369                                let c = self.cos(v_arg).unwrap();
1370                                let c = self.square(c).unwrap();
1371                                self.div(d_arg, c)
1372                            }
1373
1374                            UnaryOpcode::Asin => {
1375                                let v = self.square(v_arg).unwrap();
1376                                let v = self.sub(1.0, v).unwrap();
1377                                let v = self.sqrt(v).unwrap();
1378                                self.div(d_arg, v)
1379                            }
1380                            UnaryOpcode::Acos => {
1381                                let v = self.square(v_arg).unwrap();
1382                                let v = self.sub(1.0, v).unwrap();
1383                                let v = self.sqrt(v).unwrap();
1384                                let v = self.neg(v).unwrap();
1385                                self.div(d_arg, v)
1386                            }
1387                            UnaryOpcode::Atan => {
1388                                let v = self.square(v_arg).unwrap();
1389                                let v = self.add(1.0, v).unwrap();
1390                                self.div(d_arg, v)
1391                            }
1392                            UnaryOpcode::Exp => self.mul(n, d_arg),
1393                            UnaryOpcode::Ln => self.div(d_arg, v_arg),
1394                            UnaryOpcode::Not => Ok(zero),
1395                        }
1396                        .unwrap();
1397                        seen.insert(n, out);
1398                        stack.push(out);
1399                    }
1400                    Op::Binary(op, v_lhs, v_rhs) => {
1401                        let d_lhs = stack.pop().unwrap();
1402                        let d_rhs = stack.pop().unwrap();
1403                        let out = match op {
1404                            BinaryOpcode::Add => self.add(d_lhs, d_rhs),
1405                            BinaryOpcode::Sub => self.sub(d_lhs, d_rhs),
1406                            BinaryOpcode::Mul => {
1407                                let a = self.mul(d_lhs, v_rhs).unwrap();
1408                                let b = self.mul(v_lhs, d_rhs).unwrap();
1409                                self.add(a, b)
1410                            }
1411                            BinaryOpcode::Div => {
1412                                let v = self.square(v_rhs).unwrap();
1413                                let a = self.mul(v_rhs, d_lhs).unwrap();
1414                                let b = self.mul(v_lhs, d_rhs).unwrap();
1415                                let c = self.sub(a, b).unwrap();
1416                                self.div(c, v)
1417                            }
1418                            BinaryOpcode::Atan => {
1419                                let a = self.square(v_lhs).unwrap();
1420                                let b = self.square(v_rhs).unwrap();
1421                                let d = self.add(a, b).unwrap();
1422
1423                                let a = self.mul(v_rhs, d_lhs).unwrap();
1424                                let b = self.mul(v_lhs, d_rhs).unwrap();
1425                                let v = self.sub(a, b).unwrap();
1426                                self.div(v, d)
1427                            }
1428                            BinaryOpcode::Min => {
1429                                let cond =
1430                                    self.less_than(v_lhs, v_rhs).unwrap();
1431                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1432                            }
1433                            BinaryOpcode::Max => {
1434                                let cond =
1435                                    self.less_than(v_rhs, v_lhs).unwrap();
1436                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1437                            }
1438                            BinaryOpcode::Compare => Ok(zero),
1439                            BinaryOpcode::Mod => {
1440                                let e = self.div(v_lhs, v_rhs).unwrap();
1441                                let q = self.floor(e).unwrap();
1442
1443                                // XXX
1444                                // (we don't actually have %, so hack it from
1445                                // `modulo`, which is actually `rem_euclid`)
1446                                // ???
1447                                let m = self.modulo(q, v_rhs).unwrap();
1448                                let cond = self.less_than(q, zero).unwrap();
1449                                let offset = self
1450                                    .if_nonzero_else(cond, v_rhs, zero)
1451                                    .unwrap();
1452                                let m = self.sub(m, offset).unwrap();
1453
1454                                // Torn from the div_euclid implementation
1455                                let outer = self.less_than(m, zero).unwrap();
1456                                let inner =
1457                                    self.less_than(zero, v_rhs).unwrap();
1458                                let qa = self.sub(q, 1.0).unwrap();
1459                                let qb = self.add(q, 1.0).unwrap();
1460                                let inner = self
1461                                    .if_nonzero_else(inner, qa, qb)
1462                                    .unwrap();
1463                                let e = self
1464                                    .if_nonzero_else(outer, inner, q)
1465                                    .unwrap();
1466
1467                                let v = self.mul(d_rhs, e).unwrap();
1468                                self.sub(d_lhs, v)
1469                            }
1470                            BinaryOpcode::And => {
1471                                let cond = self.compare(v_lhs, zero).unwrap();
1472                                self.if_nonzero_else(cond, d_rhs, d_lhs)
1473                            }
1474                            BinaryOpcode::Or => {
1475                                let cond = self.compare(v_lhs, zero).unwrap();
1476                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1477                            }
1478                            BinaryOpcode::Mix => Ok(zero),
1479                        }
1480                        .unwrap();
1481                        seen.insert(n, out);
1482                        stack.push(out);
1483                    }
1484                },
1485            }
1486        }
1487        assert_eq!(stack.len(), 1);
1488        Ok(stack.pop().unwrap())
1489    }
1490}
1491
1492/// Error indicating that a node is missing from the context
1493#[derive(thiserror::Error, Debug)]
1494#[error("node is not present in this `Context`")]
1495pub struct BadNode;
1496
1497/// Error type for [`Context::from_text`]
1498#[derive(thiserror::Error, Debug)]
1499pub enum ParseError {
1500    /// Unknown opcode {0}
1501    #[error("unknown opcode {0}")]
1502    UnknownOpcode(String),
1503
1504    /// Unknown variable {0}
1505    #[error("unknown variable {0}")]
1506    UnknownVariable(String),
1507
1508    /// Node is missing from the context
1509    #[error(transparent)]
1510    BadNode(#[from] BadNode),
1511
1512    /// Empty file
1513    #[error("empty file")]
1514    EmptyFile,
1515}
1516
1517/// Error for getting a constant from a node
1518#[derive(thiserror::Error, Debug)]
1519pub enum ConstError {
1520    /// The given node is not a constant
1521    #[error("node is not a constant")]
1522    NotAConst,
1523
1524    /// Node is missing from the context
1525    #[error(transparent)]
1526    BadNode(#[from] BadNode),
1527}
1528
1529/// Error indicating that the node is not a [`Var`]
1530#[derive(thiserror::Error, Debug)]
1531#[error("node does not have an associated variable")]
1532pub struct NotAVar;
1533
1534/// Error for getting a [`Var`] from a node
1535#[derive(thiserror::Error, Debug)]
1536pub enum VarError {
1537    /// Node is not a [`Var`]
1538    #[error(transparent)]
1539    NotAVar(#[from] NotAVar),
1540
1541    /// Node is missing from the context
1542    #[error(transparent)]
1543    BadNode(#[from] BadNode),
1544}
1545
1546/// Error during tree-walking evaluation
1547#[derive(thiserror::Error, Debug)]
1548pub enum EvalError {
1549    /// Variable is missing in the evaluation map
1550    #[error("variable {0} is missing in the evaluation map")]
1551    MissingVar(Var),
1552
1553    /// Node is missing from the context
1554    #[error(transparent)]
1555    BadNode(#[from] BadNode),
1556}
1557
1558////////////////////////////////////////////////////////////////////////////////
1559/// Helper trait for things that can be converted into a [`Node`] given a
1560/// [`Context`].
1561///
1562/// This trait allows you to write
1563/// ```
1564/// # let mut ctx = fidget_core::context::Context::new();
1565/// let x = ctx.x();
1566/// let sum = ctx.add(x, 1.0).unwrap();
1567/// ```
1568/// instead of the more verbose
1569/// ```
1570/// # let mut ctx = fidget_core::context::Context::new();
1571/// let x = ctx.x();
1572/// let num = ctx.constant(1.0);
1573/// let sum = ctx.add(x, num).unwrap();
1574/// ```
1575pub trait IntoNode {
1576    /// Converts the given values into a node
1577    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode>;
1578}
1579
1580impl IntoNode for Node {
1581    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
1582        ctx.check_node(self)?;
1583        Ok(self)
1584    }
1585}
1586
1587impl IntoNode for f32 {
1588    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
1589        Ok(ctx.constant(self))
1590    }
1591}
1592
1593////////////////////////////////////////////////////////////////////////////////
1594
1595#[cfg(test)]
1596mod test {
1597    use super::*;
1598    use crate::vm::VmData;
1599
1600    // This can't be in a doctest, because it uses a private function
1601    #[test]
1602    fn test_get_op() {
1603        let mut ctx = Context::new();
1604        let x = ctx.x();
1605        let op_x = ctx.get_op(x).unwrap();
1606        assert!(matches!(op_x, Op::Input(_)));
1607    }
1608
1609    #[test]
1610    fn test_ring() {
1611        let mut ctx = Context::new();
1612        let c0 = ctx.constant(0.5);
1613        let x = ctx.x();
1614        let y = ctx.y();
1615        let x2 = ctx.square(x).unwrap();
1616        let y2 = ctx.square(y).unwrap();
1617        let r = ctx.add(x2, y2).unwrap();
1618        let c6 = ctx.sub(r, c0).unwrap();
1619        let c7 = ctx.constant(0.25);
1620        let c8 = ctx.sub(c7, r).unwrap();
1621        let c9 = ctx.max(c8, c6).unwrap();
1622
1623        let tape = VmData::<255>::new(&ctx, &[c9]).unwrap();
1624        assert_eq!(tape.len(), 9);
1625        assert_eq!(tape.vars.len(), 2);
1626    }
1627
1628    #[test]
1629    fn test_dupe() {
1630        let mut ctx = Context::new();
1631        let x = ctx.x();
1632        let x_squared = ctx.mul(x, x).unwrap();
1633
1634        let tape = VmData::<255>::new(&ctx, &[x_squared]).unwrap();
1635        assert_eq!(tape.len(), 3); // x, square, output
1636        assert_eq!(tape.vars.len(), 1);
1637    }
1638
1639    #[test]
1640    fn test_export() {
1641        let mut ctx = Context::new();
1642        let x = ctx.x();
1643        let s = ctx.sin(x).unwrap();
1644        let c = ctx.cos(x).unwrap();
1645        let sum = ctx.add(s, c).unwrap();
1646        let t = ctx.export(sum).unwrap();
1647        if let TreeOp::Binary(BinaryOpcode::Add, lhs, rhs) = &*t {
1648            match (&**lhs, &**rhs) {
1649                (
1650                    TreeOp::Unary(UnaryOpcode::Sin, x1),
1651                    TreeOp::Unary(UnaryOpcode::Cos, x2),
1652                ) => {
1653                    assert_eq!(Arc::as_ptr(x1), Arc::as_ptr(x2));
1654                    let TreeOp::Input(Var::X) = &**x1 else {
1655                        panic!("invalid X: {x1:?}");
1656                    };
1657                }
1658                _ => panic!("invalid lhs / rhs: {lhs:?} {rhs:?}"),
1659            }
1660        } else {
1661            panic!("unexpected opcode {t:?}");
1662        }
1663    }
1664
1665    #[test]
1666    fn import_optimization() {
1667        let t = Tree::x() + 0;
1668        let mut ctx = Context::new();
1669        let root = ctx.import(&t);
1670        assert_eq!(ctx.get_op(root).unwrap(), &Op::Input(Var::X));
1671    }
1672}