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<f64, 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: f64) -> 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::f64::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::f64::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 that is 1 if `lhs < rhs` and 0 otherwise
671    ///
672    /// ```
673    /// # let mut ctx = fidget_core::context::Context::new();
674    /// let x = ctx.x();
675    /// let y = ctx.y();
676    /// let op = ctx.less_than(x, y).unwrap();
677    /// let v = ctx.eval_xyz(op, 0.0, 1.0, 0.0).unwrap();
678    /// assert_eq!(v, 1.0);
679    /// let v = ctx.eval_xyz(op, 1.0, 1.0, 0.0).unwrap();
680    /// assert_eq!(v, 0.0);
681    /// let v = ctx.eval_xyz(op, 2.0, 1.0, 0.0).unwrap();
682    /// assert_eq!(v, 0.0);
683    /// ```
684    pub fn less_than<A: IntoNode, B: IntoNode>(
685        &mut self,
686        lhs: A,
687        rhs: B,
688    ) -> Result<Node, BadNode> {
689        let lhs = lhs.into_node(self)?;
690        let rhs = rhs.into_node(self)?;
691        let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
692        self.max(cmp, 0.0)
693    }
694
695    /// Builds a node that is 1 if `lhs <= rhs` and 0 otherwise
696    ///
697    /// ```
698    /// # let mut ctx = fidget_core::context::Context::new();
699    /// let x = ctx.x();
700    /// let y = ctx.y();
701    /// let op = ctx.less_than_or_equal(x, y).unwrap();
702    /// let v = ctx.eval_xyz(op, 0.0, 1.0, 0.0).unwrap();
703    /// assert_eq!(v, 1.0);
704    /// let v = ctx.eval_xyz(op, 1.0, 1.0, 0.0).unwrap();
705    /// assert_eq!(v, 1.0);
706    /// let v = ctx.eval_xyz(op, 2.0, 1.0, 0.0).unwrap();
707    /// assert_eq!(v, 0.0);
708    /// ```
709    pub fn less_than_or_equal<A: IntoNode, B: IntoNode>(
710        &mut self,
711        lhs: A,
712        rhs: B,
713    ) -> Result<Node, BadNode> {
714        let lhs = lhs.into_node(self)?;
715        let rhs = rhs.into_node(self)?;
716        let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
717        let shift = self.add(cmp, 1.0)?;
718        self.min(shift, 1.0)
719    }
720
721    /// Builds a node that takes the modulo (least non-negative remainder)
722    pub fn modulo<A: IntoNode, B: IntoNode>(
723        &mut self,
724        a: A,
725        b: B,
726    ) -> Result<Node, BadNode> {
727        let a = a.into_node(self)?;
728        let b = b.into_node(self)?;
729        self.op_binary(a, b, BinaryOpcode::Mod)
730    }
731
732    /// Builds a node that returns the first node if the condition is not
733    /// equal to zero, else returns the other node
734    ///
735    /// The result is `a` if `condition != 0`, else the result is `b`.
736    /// ```
737    /// # let mut ctx = fidget_core::context::Context::new();
738    /// let x = ctx.x();
739    /// let y = ctx.y();
740    /// let z = ctx.z();
741    ///
742    /// let if_else = ctx.if_nonzero_else(x, y, z).unwrap();
743    ///
744    /// assert_eq!(ctx.eval_xyz(if_else, 0.0, 2.0, 3.0).unwrap(), 3.0);
745    /// assert_eq!(ctx.eval_xyz(if_else, 1.0, 2.0, 3.0).unwrap(), 2.0);
746    /// assert_eq!(ctx.eval_xyz(if_else, 0.0, f64::NAN, 3.0).unwrap(), 3.0);
747    /// assert_eq!(ctx.eval_xyz(if_else, 1.0, 2.0, f64::NAN).unwrap(), 2.0);
748    /// ```
749    pub fn if_nonzero_else<Condition: IntoNode, A: IntoNode, B: IntoNode>(
750        &mut self,
751        condition: Condition,
752        a: A,
753        b: B,
754    ) -> Result<Node, BadNode> {
755        let condition = condition.into_node(self)?;
756        let a = a.into_node(self)?;
757        let b = b.into_node(self)?;
758
759        let lhs = self.and(condition, a)?;
760        let n_condition = self.not(condition)?;
761        let rhs = self.and(n_condition, b)?;
762        self.or(lhs, rhs)
763    }
764
765    ////////////////////////////////////////////////////////////////////////////
766    /// Evaluates the given node with the provided values for X, Y, and Z.
767    ///
768    /// This is extremely inefficient; consider converting the node into a
769    /// [`Shape`](crate::shape::Shape) and using its evaluators instead.
770    ///
771    /// ```
772    /// # let mut ctx = fidget_core::context::Context::new();
773    /// let x = ctx.x();
774    /// let y = ctx.y();
775    /// let z = ctx.z();
776    /// let op = ctx.mul(x, y).unwrap();
777    /// let op = ctx.div(op, z).unwrap();
778    /// let v = ctx.eval_xyz(op, 3.0, 5.0, 2.0).unwrap();
779    /// assert_eq!(v, 7.5); // (3.0 * 5.0) / 2.0
780    /// ```
781    pub fn eval_xyz(
782        &self,
783        root: Node,
784        x: f64,
785        y: f64,
786        z: f64,
787    ) -> Result<f64, EvalError> {
788        let vars = [(Var::X, x), (Var::Y, y), (Var::Z, z)]
789            .into_iter()
790            .collect();
791        self.eval(root, &vars)
792    }
793
794    /// Evaluates the given node with a generic set of variables
795    ///
796    /// This is extremely inefficient; consider converting the node into a
797    /// [`Shape`](crate::shape::Shape) and using its evaluators instead.
798    pub fn eval(
799        &self,
800        root: Node,
801        vars: &HashMap<Var, f64>,
802    ) -> Result<f64, EvalError> {
803        let mut cache = vec![None; self.ops.len()].into();
804        self.eval_inner(root, vars, &mut cache)
805    }
806
807    fn eval_inner(
808        &self,
809        node: Node,
810        vars: &HashMap<Var, f64>,
811        cache: &mut IndexVec<Option<f64>, Node>,
812    ) -> Result<f64, EvalError> {
813        if node.0 >= cache.len() {
814            return Err(EvalError::BadNode(BadNode));
815        }
816        if let Some(v) = cache[node] {
817            return Ok(v);
818        }
819        let mut get = |n: Node| self.eval_inner(n, vars, cache);
820        let v = match self.get_op(node).ok_or(EvalError::BadNode(BadNode))? {
821            Op::Input(v) => *vars.get(v).ok_or(EvalError::MissingVar(*v))?,
822            Op::Const(c) => c.0,
823
824            Op::Binary(op, a, b) => {
825                let a = get(*a)?;
826                let b = get(*b)?;
827                op.eval(a, b)
828            }
829
830            // Unary operations
831            Op::Unary(op, a) => {
832                let a = get(*a)?;
833                op.eval(a)
834            }
835        };
836
837        cache[node] = Some(v);
838        Ok(v)
839    }
840
841    /// Parses a flat text representation of a math tree. For example, the
842    /// circle `(- (+ (square x) (square y)) 1)` can be parsed from
843    /// ```
844    /// # use fidget_core::context::Context;
845    /// let txt = "
846    /// ## This is a comment!
847    /// 0x600000b90000 var-x
848    /// 0x600000b900a0 square 0x600000b90000
849    /// 0x600000b90050 var-y
850    /// 0x600000b900f0 square 0x600000b90050
851    /// 0x600000b90140 add 0x600000b900a0 0x600000b900f0
852    /// 0x600000b90190 sqrt 0x600000b90140
853    /// 0x600000b901e0 const 1
854    /// ";
855    /// let (ctx, _node) = Context::from_text(&mut txt.as_bytes()).unwrap();
856    /// assert_eq!(ctx.len(), 7);
857    /// ```
858    ///
859    /// This representation is loosely defined and only intended for use in
860    /// quick experiments.
861    pub fn from_text<R: Read>(r: R) -> Result<(Self, Node), ParseError> {
862        let reader = BufReader::new(r);
863        let mut ctx = Self::new();
864        let mut seen = BTreeMap::new();
865        let mut last = None;
866
867        for line in reader.lines().map(|line| line.unwrap()) {
868            if line.is_empty() || line.starts_with('#') {
869                continue;
870            }
871            let mut iter = line.split_whitespace();
872            let i: String = iter.next().unwrap().to_owned();
873            let opcode = iter.next().unwrap();
874
875            let mut pop = || {
876                let txt = iter.next().unwrap();
877                seen.get(txt)
878                    .cloned()
879                    .ok_or_else(|| ParseError::UnknownVariable(txt.to_string()))
880            };
881            let node = match opcode {
882                "const" => ctx.constant(iter.next().unwrap().parse().unwrap()),
883                "var-x" => ctx.x(),
884                "var-y" => ctx.y(),
885                "var-z" => ctx.z(),
886                "abs" => ctx.abs(pop()?)?,
887                "neg" => ctx.neg(pop()?)?,
888                "sqrt" => ctx.sqrt(pop()?)?,
889                "square" => ctx.square(pop()?)?,
890                "floor" => ctx.floor(pop()?)?,
891                "ceil" => ctx.ceil(pop()?)?,
892                "round" => ctx.round(pop()?)?,
893                "sin" => ctx.sin(pop()?)?,
894                "cos" => ctx.cos(pop()?)?,
895                "tan" => ctx.tan(pop()?)?,
896                "asin" => ctx.asin(pop()?)?,
897                "acos" => ctx.acos(pop()?)?,
898                "atan" => ctx.atan(pop()?)?,
899                "ln" => ctx.ln(pop()?)?,
900                "not" => ctx.not(pop()?)?,
901                "exp" => ctx.exp(pop()?)?,
902                "add" => ctx.add(pop()?, pop()?)?,
903                "mul" => ctx.mul(pop()?, pop()?)?,
904                "min" => ctx.min(pop()?, pop()?)?,
905                "max" => ctx.max(pop()?, pop()?)?,
906                "div" => ctx.div(pop()?, pop()?)?,
907                "atan2" => ctx.atan2(pop()?, pop()?)?,
908                "sub" => ctx.sub(pop()?, pop()?)?,
909                "compare" => ctx.compare(pop()?, pop()?)?,
910                "mod" => ctx.modulo(pop()?, pop()?)?,
911                "and" => ctx.and(pop()?, pop()?)?,
912                "or" => ctx.or(pop()?, pop()?)?,
913                op => return Err(ParseError::UnknownOpcode(op.to_owned())),
914            };
915            seen.insert(i, node);
916            last = Some(node);
917        }
918        match last {
919            Some(node) => Ok((ctx, node)),
920            None => Err(ParseError::EmptyFile),
921        }
922    }
923
924    /// Converts the entire context into a GraphViz drawing
925    pub fn dot(&self) -> String {
926        let mut out = "digraph mygraph{\n".to_owned();
927        for node in self.ops.keys() {
928            let op = self.get_op(node).unwrap();
929            out += &self.dot_node(node);
930            out += &op.dot_edges(node);
931        }
932        out += "}\n";
933        out
934    }
935
936    /// Converts the given node into a GraphViz node
937    ///
938    /// (this is a local function instead of a function on `Op` because it
939    ///  requires looking up variables by name)
940    fn dot_node(&self, i: Node) -> String {
941        let mut out = format!(r#"n{} [label = ""#, i.get());
942        let op = self.get_op(i).unwrap();
943        match op {
944            Op::Const(c) => write!(out, "{c}").unwrap(),
945            Op::Input(v) => {
946                out += &v.to_string();
947            }
948            Op::Binary(op, ..) => match op {
949                BinaryOpcode::Add => out += "add",
950                BinaryOpcode::Sub => out += "sub",
951                BinaryOpcode::Mul => out += "mul",
952                BinaryOpcode::Div => out += "div",
953                BinaryOpcode::Atan => out += "atan2",
954                BinaryOpcode::Min => out += "min",
955                BinaryOpcode::Max => out += "max",
956                BinaryOpcode::Compare => out += "compare",
957                BinaryOpcode::Mod => out += "mod",
958                BinaryOpcode::And => out += "and",
959                BinaryOpcode::Or => out += "or",
960            },
961            Op::Unary(op, ..) => match op {
962                UnaryOpcode::Neg => out += "neg",
963                UnaryOpcode::Abs => out += "abs",
964                UnaryOpcode::Recip => out += "recip",
965                UnaryOpcode::Sqrt => out += "sqrt",
966                UnaryOpcode::Square => out += "square",
967                UnaryOpcode::Floor => out += "floor",
968                UnaryOpcode::Ceil => out += "ceil",
969                UnaryOpcode::Round => out += "round",
970                UnaryOpcode::Sin => out += "sin",
971                UnaryOpcode::Cos => out += "cos",
972                UnaryOpcode::Tan => out += "tan",
973                UnaryOpcode::Asin => out += "asin",
974                UnaryOpcode::Acos => out += "acos",
975                UnaryOpcode::Atan => out += "atan",
976                UnaryOpcode::Exp => out += "exp",
977                UnaryOpcode::Ln => out += "ln",
978                UnaryOpcode::Not => out += "not",
979            },
980        };
981        write!(
982            out,
983            r#"" color="{0}1" shape="{1}" fontcolor="{0}4"]"#,
984            op.dot_node_color(),
985            op.dot_node_shape()
986        )
987        .unwrap();
988        out
989    }
990
991    /// Looks up an operation by `Node` handle
992    pub fn get_op(&self, node: Node) -> Option<&Op> {
993        self.ops.get_by_index(node)
994    }
995
996    /// Imports the given tree, deduplicating and returning the root
997    pub fn import(&mut self, tree: &Tree) -> Node {
998        // A naive remapping implementation would use recursion.  A naive
999        // remapping implementation would blow up the stack given any
1000        // significant tree size.
1001        //
1002        // Instead, we maintain our own pseudo-stack here in a pair of Vecs (one
1003        // stack for actions, and a second stack for return values).
1004        enum Action<'a> {
1005            /// Pushes `Up(op)` followed by `Down(c)` for each child
1006            Down(&'a Arc<TreeOp>),
1007            /// Consumes imported trees from the stack and pushes a new tree
1008            Up(&'a Arc<TreeOp>),
1009            /// Pops the latest axis frame
1010            Pop,
1011            /// Pops the latest affine frame
1012            PopAffine,
1013        }
1014        let mut axes = vec![(self.x(), self.y(), self.z())];
1015        let mut todo = vec![Action::Down(tree.arc())];
1016        let mut stack = vec![];
1017        let mut affine: Vec<Matrix4<f64>> = vec![];
1018
1019        // Cache of TreeOp -> Node mapping under a particular frame (axes)
1020        //
1021        // This isn't required for correctness, but can be a speed optimization
1022        // (because it means we don't have to walk the same tree twice).
1023        let mut seen = HashMap::new();
1024
1025        while let Some(t) = todo.pop() {
1026            match t {
1027                Action::Down(t) => {
1028                    // If we've already seen this TreeOp with these axes, then
1029                    // we can return the previous Node.
1030                    if matches!(
1031                        t.as_ref(),
1032                        TreeOp::Unary(..) | TreeOp::Binary(..)
1033                    ) && let Some(p) =
1034                        seen.get(&(*axes.last().unwrap(), Arc::as_ptr(t)))
1035                    {
1036                        stack.push(*p);
1037                        continue;
1038                    }
1039                    match t.as_ref() {
1040                        TreeOp::Const(c) => {
1041                            stack.push(self.constant(*c));
1042                        }
1043                        TreeOp::Input(s) => {
1044                            let axes = axes.last().unwrap();
1045                            stack.push(match *s {
1046                                Var::X => axes.0,
1047                                Var::Y => axes.1,
1048                                Var::Z => axes.2,
1049                                v @ Var::V(..) => self.var(v),
1050                            });
1051                        }
1052                        TreeOp::Unary(_op, arg) => {
1053                            todo.push(Action::Up(t));
1054                            todo.push(Action::Down(arg));
1055                        }
1056                        TreeOp::Binary(_op, lhs, rhs) => {
1057                            todo.push(Action::Up(t));
1058                            todo.push(Action::Down(lhs));
1059                            todo.push(Action::Down(rhs));
1060                        }
1061                        TreeOp::RemapAxes { target: _, x, y, z } => {
1062                            // Action::Up(t) does the remapping and target eval
1063                            todo.push(Action::Up(t));
1064                            todo.push(Action::Down(x));
1065                            todo.push(Action::Down(y));
1066                            todo.push(Action::Down(z));
1067                        }
1068                        TreeOp::RemapAffine { target, mat } => {
1069                            let prev = affine
1070                                .last()
1071                                .cloned()
1072                                .unwrap_or(Matrix4::identity());
1073                            let mat = prev * mat.to_homogeneous();
1074
1075                            // Push either an affine frame or an axis frame,
1076                            // depending on whether the target is also affine
1077                            if matches!(&**target, TreeOp::RemapAffine { .. }) {
1078                                affine.push(mat);
1079                                todo.push(Action::PopAffine);
1080                            } else {
1081                                let (x, y, z) = axes.last().unwrap();
1082                                let mut out = [None; 3];
1083                                for i in 0..3 {
1084                                    let a = self.mul(mat[(i, 0)], *x).unwrap();
1085                                    let b = self.mul(mat[(i, 1)], *y).unwrap();
1086                                    let c = self.mul(mat[(i, 2)], *z).unwrap();
1087                                    let d = self.constant(mat[(i, 3)]);
1088                                    let ab = self.add(a, b).unwrap();
1089                                    let cd = self.add(c, d).unwrap();
1090                                    out[i] = Some(self.add(ab, cd).unwrap());
1091                                }
1092                                let [x, y, z] = out.map(Option::unwrap);
1093                                axes.push((x, y, z));
1094                                todo.push(Action::Pop);
1095                            }
1096                            todo.push(Action::Down(target));
1097                        }
1098                    }
1099                }
1100                Action::Up(t) => {
1101                    match t.as_ref() {
1102                        TreeOp::Const(..)
1103                        | TreeOp::Input(..)
1104                        | TreeOp::RemapAffine { .. } => unreachable!(),
1105                        TreeOp::Unary(op, ..) => {
1106                            let arg = stack.pop().unwrap();
1107                            let out = self.op_unary(arg, *op).unwrap();
1108                            stack.push(out);
1109                        }
1110                        TreeOp::Binary(op, ..) => {
1111                            let lhs = stack.pop().unwrap();
1112                            let rhs = stack.pop().unwrap();
1113                            // Call individual builders to apply optimizations
1114                            let out = match op {
1115                                BinaryOpcode::Add => self.add(lhs, rhs),
1116                                BinaryOpcode::Sub => self.sub(lhs, rhs),
1117                                BinaryOpcode::Mul => self.mul(lhs, rhs),
1118                                BinaryOpcode::Div => self.div(lhs, rhs),
1119                                BinaryOpcode::Atan => self.atan2(lhs, rhs),
1120                                BinaryOpcode::Min => self.min(lhs, rhs),
1121                                BinaryOpcode::Max => self.max(lhs, rhs),
1122                                BinaryOpcode::Compare => self.compare(lhs, rhs),
1123                                BinaryOpcode::Mod => self.modulo(lhs, rhs),
1124                                BinaryOpcode::And => self.and(lhs, rhs),
1125                                BinaryOpcode::Or => self.or(lhs, rhs),
1126                            }
1127                            .unwrap();
1128                            if Arc::strong_count(t) > 1 {
1129                                seen.insert(
1130                                    (*axes.last().unwrap(), Arc::as_ptr(t)),
1131                                    out,
1132                                );
1133                            }
1134                            stack.push(out);
1135                        }
1136                        TreeOp::RemapAxes { target, .. } => {
1137                            let x = stack.pop().unwrap();
1138                            let y = stack.pop().unwrap();
1139                            let z = stack.pop().unwrap();
1140                            axes.push((x, y, z));
1141                            todo.push(Action::Pop);
1142                            todo.push(Action::Down(target));
1143                        }
1144                    }
1145                    // Update the cache with the new tree, if relevant
1146                    //
1147                    // The `strong_count` check is a rough heuristic to avoid
1148                    // caching if there's only a single owner of the tree.  This
1149                    // isn't perfect, but it doesn't need to be for correctness.
1150                    if matches!(
1151                        t.as_ref(),
1152                        TreeOp::Unary(..) | TreeOp::Binary(..)
1153                    ) && Arc::strong_count(t) > 1
1154                    {
1155                        seen.insert(
1156                            (*axes.last().unwrap(), Arc::as_ptr(t)),
1157                            *stack.last().unwrap(),
1158                        );
1159                    }
1160                }
1161                Action::Pop => {
1162                    axes.pop().unwrap();
1163                }
1164                Action::PopAffine => {
1165                    affine.pop().unwrap();
1166                }
1167            }
1168        }
1169        assert_eq!(stack.len(), 1);
1170        stack.pop().unwrap()
1171    }
1172
1173    /// Converts from a context-specific node into a standalone [`Tree`]
1174    pub fn export(&self, n: Node) -> Result<Tree, BadNode> {
1175        if self.get_op(n).is_none() {
1176            return Err(BadNode);
1177        }
1178
1179        // Do recursion on the heap to avoid stack overflows for deep trees
1180        enum Action {
1181            /// Pushes `Up(n)` followed by `Down(n)` for each child
1182            Down(Node),
1183            /// Consumes trees from the stack and pushes a new tree
1184            Up(Node, Op),
1185        }
1186        let mut todo = vec![Action::Down(n)];
1187        let mut stack = vec![];
1188
1189        // Cache of Node -> Tree mapping, for Tree deduplication
1190        let mut seen: HashMap<Node, Tree> = HashMap::new();
1191
1192        while let Some(t) = todo.pop() {
1193            match t {
1194                Action::Down(n) => {
1195                    // If we've already seen this TreeOp with these axes, then
1196                    // we can return the previous Node.
1197                    if let Some(p) = seen.get(&n) {
1198                        stack.push(p.clone());
1199                        continue;
1200                    }
1201                    let op = self.get_op(n).unwrap();
1202                    match op {
1203                        Op::Const(c) => {
1204                            let t = Tree::from(c.0);
1205                            seen.insert(n, t.clone());
1206                            stack.push(t);
1207                        }
1208                        Op::Input(v) => {
1209                            let t = Tree::from(*v);
1210                            seen.insert(n, t.clone());
1211                            stack.push(t);
1212                        }
1213                        Op::Unary(_op, arg) => {
1214                            todo.push(Action::Up(n, *op));
1215                            todo.push(Action::Down(*arg));
1216                        }
1217                        Op::Binary(_op, lhs, rhs) => {
1218                            todo.push(Action::Up(n, *op));
1219                            todo.push(Action::Down(*lhs));
1220                            todo.push(Action::Down(*rhs));
1221                        }
1222                    }
1223                }
1224                Action::Up(n, op) => match op {
1225                    Op::Const(..) | Op::Input(..) => unreachable!(),
1226                    Op::Unary(op, ..) => {
1227                        let arg = stack.pop().unwrap();
1228                        let out =
1229                            Tree::from(TreeOp::Unary(op, arg.arc().clone()));
1230                        seen.insert(n, out.clone());
1231                        stack.push(out);
1232                    }
1233                    Op::Binary(op, ..) => {
1234                        let lhs = stack.pop().unwrap();
1235                        let rhs = stack.pop().unwrap();
1236                        let out = Tree::from(TreeOp::Binary(
1237                            op,
1238                            lhs.arc().clone(),
1239                            rhs.arc().clone(),
1240                        ));
1241                        seen.insert(n, out.clone());
1242                        stack.push(out);
1243                    }
1244                },
1245            }
1246        }
1247        assert_eq!(stack.len(), 1);
1248        Ok(stack.pop().unwrap())
1249    }
1250
1251    /// Takes the symbolic derivative of a node with respect to a variable
1252    pub fn deriv(&mut self, n: Node, v: Var) -> Result<Node, BadNode> {
1253        if self.get_op(n).is_none() {
1254            return Err(BadNode);
1255        }
1256
1257        // Do recursion on the heap to avoid stack overflows for deep trees
1258        enum Action {
1259            /// Pushes `Up(n)` followed by `Down(n)` for each child
1260            Down(Node),
1261            /// Consumes trees from the stack and pushes a new tree
1262            Up(Node, Op),
1263        }
1264        let mut todo = vec![Action::Down(n)];
1265        let mut stack = vec![];
1266        let zero = self.constant(0.0);
1267
1268        // Cache of Node -> Node mapping, for deduplication
1269        let mut seen: HashMap<Node, Node> = HashMap::new();
1270
1271        while let Some(t) = todo.pop() {
1272            match t {
1273                Action::Down(n) => {
1274                    // If we've already seen this TreeOp with these axes, then
1275                    // we can return the previous Node.
1276                    if let Some(p) = seen.get(&n) {
1277                        stack.push(*p);
1278                        continue;
1279                    }
1280                    let op = *self.get_op(n).unwrap();
1281                    match op {
1282                        Op::Const(_c) => {
1283                            seen.insert(n, zero);
1284                            stack.push(zero);
1285                        }
1286                        Op::Input(u) => {
1287                            let z =
1288                                if v == u { self.constant(1.0) } else { zero };
1289                            seen.insert(n, z);
1290                            stack.push(z);
1291                        }
1292                        Op::Unary(_op, arg) => {
1293                            todo.push(Action::Up(n, op));
1294                            todo.push(Action::Down(arg));
1295                        }
1296                        Op::Binary(_op, lhs, rhs) => {
1297                            todo.push(Action::Up(n, op));
1298                            todo.push(Action::Down(lhs));
1299                            todo.push(Action::Down(rhs));
1300                        }
1301                    }
1302                }
1303                Action::Up(n, op) => match op {
1304                    Op::Const(..) | Op::Input(..) => unreachable!(),
1305                    Op::Unary(op, v_arg) => {
1306                        let d_arg = stack.pop().unwrap();
1307                        let out = match op {
1308                            UnaryOpcode::Neg => self.neg(d_arg),
1309                            UnaryOpcode::Abs => {
1310                                let cond = self.less_than(v_arg, zero).unwrap();
1311                                let pos = d_arg;
1312                                let neg = self.neg(d_arg).unwrap();
1313                                self.if_nonzero_else(cond, neg, pos)
1314                            }
1315                            UnaryOpcode::Recip => {
1316                                let a = self.square(v_arg).unwrap();
1317                                let b = self.neg(d_arg).unwrap();
1318                                self.div(b, a)
1319                            }
1320                            UnaryOpcode::Sqrt => {
1321                                let v = self.mul(n, 2.0).unwrap();
1322                                self.div(d_arg, v)
1323                            }
1324                            UnaryOpcode::Square => {
1325                                let v = self.mul(d_arg, v_arg).unwrap();
1326                                self.mul(2.0, v)
1327                            }
1328                            // Discontinuous constants don't have Dirac deltas
1329                            UnaryOpcode::Floor
1330                            | UnaryOpcode::Ceil
1331                            | UnaryOpcode::Round => Ok(zero),
1332
1333                            UnaryOpcode::Sin => {
1334                                let c = self.cos(v_arg).unwrap();
1335                                self.mul(c, d_arg)
1336                            }
1337
1338                            UnaryOpcode::Cos => {
1339                                let s = self.sin(v_arg).unwrap();
1340                                let s = self.neg(s).unwrap();
1341                                self.mul(s, d_arg)
1342                            }
1343
1344                            UnaryOpcode::Tan => {
1345                                let c = self.cos(v_arg).unwrap();
1346                                let c = self.square(c).unwrap();
1347                                self.div(d_arg, c)
1348                            }
1349
1350                            UnaryOpcode::Asin => {
1351                                let v = self.square(v_arg).unwrap();
1352                                let v = self.sub(1.0, v).unwrap();
1353                                let v = self.sqrt(v).unwrap();
1354                                self.div(d_arg, v)
1355                            }
1356                            UnaryOpcode::Acos => {
1357                                let v = self.square(v_arg).unwrap();
1358                                let v = self.sub(1.0, v).unwrap();
1359                                let v = self.sqrt(v).unwrap();
1360                                let v = self.neg(v).unwrap();
1361                                self.div(d_arg, v)
1362                            }
1363                            UnaryOpcode::Atan => {
1364                                let v = self.square(v_arg).unwrap();
1365                                let v = self.add(1.0, v).unwrap();
1366                                self.div(d_arg, v)
1367                            }
1368                            UnaryOpcode::Exp => self.mul(n, d_arg),
1369                            UnaryOpcode::Ln => self.div(d_arg, v_arg),
1370                            UnaryOpcode::Not => Ok(zero),
1371                        }
1372                        .unwrap();
1373                        seen.insert(n, out);
1374                        stack.push(out);
1375                    }
1376                    Op::Binary(op, v_lhs, v_rhs) => {
1377                        let d_lhs = stack.pop().unwrap();
1378                        let d_rhs = stack.pop().unwrap();
1379                        let out = match op {
1380                            BinaryOpcode::Add => self.add(d_lhs, d_rhs),
1381                            BinaryOpcode::Sub => self.sub(d_lhs, d_rhs),
1382                            BinaryOpcode::Mul => {
1383                                let a = self.mul(d_lhs, v_rhs).unwrap();
1384                                let b = self.mul(v_lhs, d_rhs).unwrap();
1385                                self.add(a, b)
1386                            }
1387                            BinaryOpcode::Div => {
1388                                let v = self.square(v_rhs).unwrap();
1389                                let a = self.mul(v_rhs, d_lhs).unwrap();
1390                                let b = self.mul(v_lhs, d_rhs).unwrap();
1391                                let c = self.sub(a, b).unwrap();
1392                                self.div(c, v)
1393                            }
1394                            BinaryOpcode::Atan => {
1395                                let a = self.square(v_lhs).unwrap();
1396                                let b = self.square(v_rhs).unwrap();
1397                                let d = self.add(a, b).unwrap();
1398
1399                                let a = self.mul(v_rhs, d_lhs).unwrap();
1400                                let b = self.mul(v_lhs, d_rhs).unwrap();
1401                                let v = self.sub(a, b).unwrap();
1402                                self.div(v, d)
1403                            }
1404                            BinaryOpcode::Min => {
1405                                let cond =
1406                                    self.less_than(v_lhs, v_rhs).unwrap();
1407                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1408                            }
1409                            BinaryOpcode::Max => {
1410                                let cond =
1411                                    self.less_than(v_rhs, v_lhs).unwrap();
1412                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1413                            }
1414                            BinaryOpcode::Compare => Ok(zero),
1415                            BinaryOpcode::Mod => {
1416                                let e = self.div(v_lhs, v_rhs).unwrap();
1417                                let q = self.floor(e).unwrap();
1418
1419                                // XXX
1420                                // (we don't actually have %, so hack it from
1421                                // `modulo`, which is actually `rem_euclid`)
1422                                // ???
1423                                let m = self.modulo(q, v_rhs).unwrap();
1424                                let cond = self.less_than(q, zero).unwrap();
1425                                let offset = self
1426                                    .if_nonzero_else(cond, v_rhs, zero)
1427                                    .unwrap();
1428                                let m = self.sub(m, offset).unwrap();
1429
1430                                // Torn from the div_euclid implementation
1431                                let outer = self.less_than(m, zero).unwrap();
1432                                let inner =
1433                                    self.less_than(zero, v_rhs).unwrap();
1434                                let qa = self.sub(q, 1.0).unwrap();
1435                                let qb = self.add(q, 1.0).unwrap();
1436                                let inner = self
1437                                    .if_nonzero_else(inner, qa, qb)
1438                                    .unwrap();
1439                                let e = self
1440                                    .if_nonzero_else(outer, inner, q)
1441                                    .unwrap();
1442
1443                                let v = self.mul(d_rhs, e).unwrap();
1444                                self.sub(d_lhs, v)
1445                            }
1446                            BinaryOpcode::And => {
1447                                let cond = self.compare(v_lhs, zero).unwrap();
1448                                self.if_nonzero_else(cond, d_rhs, d_lhs)
1449                            }
1450                            BinaryOpcode::Or => {
1451                                let cond = self.compare(v_lhs, zero).unwrap();
1452                                self.if_nonzero_else(cond, d_lhs, d_rhs)
1453                            }
1454                        }
1455                        .unwrap();
1456                        seen.insert(n, out);
1457                        stack.push(out);
1458                    }
1459                },
1460            }
1461        }
1462        assert_eq!(stack.len(), 1);
1463        Ok(stack.pop().unwrap())
1464    }
1465}
1466
1467/// Error indicating that a node is missing from the context
1468#[derive(thiserror::Error, Debug)]
1469#[error("node is not present in this `Context`")]
1470pub struct BadNode;
1471
1472/// Error type for [`Context::from_text`]
1473#[derive(thiserror::Error, Debug)]
1474pub enum ParseError {
1475    /// Unknown opcode {0}
1476    #[error("unknown opcode {0}")]
1477    UnknownOpcode(String),
1478
1479    /// Unknown variable {0}
1480    #[error("unknown variable {0}")]
1481    UnknownVariable(String),
1482
1483    /// Node is missing from the context
1484    #[error(transparent)]
1485    BadNode(#[from] BadNode),
1486
1487    /// Empty file
1488    #[error("empty file")]
1489    EmptyFile,
1490}
1491
1492/// Error for getting a constant from a node
1493#[derive(thiserror::Error, Debug)]
1494pub enum ConstError {
1495    /// The given node is not a constant
1496    #[error("node is not a constant")]
1497    NotAConst,
1498
1499    /// Node is missing from the context
1500    #[error(transparent)]
1501    BadNode(#[from] BadNode),
1502}
1503
1504/// Error indicating that the node is not a [`Var`]
1505#[derive(thiserror::Error, Debug)]
1506#[error("node does not have an associated variable")]
1507pub struct NotAVar;
1508
1509/// Error for getting a [`Var`] from a node
1510#[derive(thiserror::Error, Debug)]
1511pub enum VarError {
1512    /// Node is not a [`Var`]
1513    #[error(transparent)]
1514    NotAVar(#[from] NotAVar),
1515
1516    /// Node is missing from the context
1517    #[error(transparent)]
1518    BadNode(#[from] BadNode),
1519}
1520
1521/// Error during tree-walking evaluation
1522#[derive(thiserror::Error, Debug)]
1523pub enum EvalError {
1524    /// Variable is missing in the evaluation map
1525    #[error("variable {0} is missing in the evaluation map")]
1526    MissingVar(Var),
1527
1528    /// Node is missing from the context
1529    #[error(transparent)]
1530    BadNode(#[from] BadNode),
1531}
1532
1533////////////////////////////////////////////////////////////////////////////////
1534/// Helper trait for things that can be converted into a [`Node`] given a
1535/// [`Context`].
1536///
1537/// This trait allows you to write
1538/// ```
1539/// # let mut ctx = fidget_core::context::Context::new();
1540/// let x = ctx.x();
1541/// let sum = ctx.add(x, 1.0).unwrap();
1542/// ```
1543/// instead of the more verbose
1544/// ```
1545/// # let mut ctx = fidget_core::context::Context::new();
1546/// let x = ctx.x();
1547/// let num = ctx.constant(1.0);
1548/// let sum = ctx.add(x, num).unwrap();
1549/// ```
1550pub trait IntoNode {
1551    /// Converts the given values into a node
1552    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode>;
1553}
1554
1555impl IntoNode for Node {
1556    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
1557        ctx.check_node(self)?;
1558        Ok(self)
1559    }
1560}
1561
1562impl IntoNode for f32 {
1563    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
1564        Ok(ctx.constant(self as f64))
1565    }
1566}
1567
1568impl IntoNode for f64 {
1569    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
1570        Ok(ctx.constant(self))
1571    }
1572}
1573
1574////////////////////////////////////////////////////////////////////////////////
1575
1576#[cfg(test)]
1577mod test {
1578    use super::*;
1579    use crate::vm::VmData;
1580
1581    // This can't be in a doctest, because it uses a private function
1582    #[test]
1583    fn test_get_op() {
1584        let mut ctx = Context::new();
1585        let x = ctx.x();
1586        let op_x = ctx.get_op(x).unwrap();
1587        assert!(matches!(op_x, Op::Input(_)));
1588    }
1589
1590    #[test]
1591    fn test_ring() {
1592        let mut ctx = Context::new();
1593        let c0 = ctx.constant(0.5);
1594        let x = ctx.x();
1595        let y = ctx.y();
1596        let x2 = ctx.square(x).unwrap();
1597        let y2 = ctx.square(y).unwrap();
1598        let r = ctx.add(x2, y2).unwrap();
1599        let c6 = ctx.sub(r, c0).unwrap();
1600        let c7 = ctx.constant(0.25);
1601        let c8 = ctx.sub(c7, r).unwrap();
1602        let c9 = ctx.max(c8, c6).unwrap();
1603
1604        let tape = VmData::<255>::new(&ctx, &[c9]).unwrap();
1605        assert_eq!(tape.len(), 9);
1606        assert_eq!(tape.vars.len(), 2);
1607    }
1608
1609    #[test]
1610    fn test_dupe() {
1611        let mut ctx = Context::new();
1612        let x = ctx.x();
1613        let x_squared = ctx.mul(x, x).unwrap();
1614
1615        let tape = VmData::<255>::new(&ctx, &[x_squared]).unwrap();
1616        assert_eq!(tape.len(), 3); // x, square, output
1617        assert_eq!(tape.vars.len(), 1);
1618    }
1619
1620    #[test]
1621    fn test_export() {
1622        let mut ctx = Context::new();
1623        let x = ctx.x();
1624        let s = ctx.sin(x).unwrap();
1625        let c = ctx.cos(x).unwrap();
1626        let sum = ctx.add(s, c).unwrap();
1627        let t = ctx.export(sum).unwrap();
1628        if let TreeOp::Binary(BinaryOpcode::Add, lhs, rhs) = &*t {
1629            match (&**lhs, &**rhs) {
1630                (
1631                    TreeOp::Unary(UnaryOpcode::Sin, x1),
1632                    TreeOp::Unary(UnaryOpcode::Cos, x2),
1633                ) => {
1634                    assert_eq!(Arc::as_ptr(x1), Arc::as_ptr(x2));
1635                    let TreeOp::Input(Var::X) = &**x1 else {
1636                        panic!("invalid X: {x1:?}");
1637                    };
1638                }
1639                _ => panic!("invalid lhs / rhs: {lhs:?} {rhs:?}"),
1640            }
1641        } else {
1642            panic!("unexpected opcode {t:?}");
1643        }
1644    }
1645
1646    #[test]
1647    fn import_optimization() {
1648        let t = Tree::x() + 0;
1649        let mut ctx = Context::new();
1650        let root = ctx.import(&t);
1651        assert_eq!(ctx.get_op(root).unwrap(), &Op::Input(Var::X));
1652    }
1653}