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}