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}