mod indexed;
mod op;
mod tree;
use indexed::{Index, IndexMap, IndexVec, define_index};
pub use op::{BinaryOpcode, Op, UnaryOpcode};
pub use tree::{Tree, TreeOp};
use crate::var::Var;
use std::collections::{BTreeMap, HashMap};
use std::fmt::Write;
use std::io::{BufRead, BufReader, Read};
use std::sync::Arc;
use nalgebra::Matrix4;
use ordered_float::OrderedFloat;
define_index!(Node, "An index in the `Context::ops` map");
#[derive(Debug, Default)]
pub struct Context {
ops: IndexMap<Op, Node>,
}
impl Context {
pub fn new() -> Self {
Self::default()
}
pub fn clear(&mut self) {
self.ops.clear();
}
pub fn len(&self) -> usize {
self.ops.len()
}
pub fn is_empty(&self) -> bool {
self.ops.is_empty()
}
fn check_node(&self, node: Node) -> Result<(), BadNode> {
self.get_op(node).ok_or(BadNode).map(|_| ())
}
pub fn get_const(&self, n: Node) -> Result<f64, ConstError> {
match self.get_op(n) {
Some(Op::Const(c)) => Ok(c.0),
Some(_) => Err(ConstError::NotAConst),
None => Err(ConstError::BadNode(BadNode)),
}
}
pub fn get_var(&self, n: Node) -> Result<Var, VarError> {
match self.get_op(n) {
Some(Op::Input(v)) => Ok(*v),
Some(..) => Err(VarError::NotAVar(NotAVar)),
None => Err(VarError::BadNode(BadNode)),
}
}
pub fn x(&mut self) -> Node {
self.var(Var::X)
}
pub fn y(&mut self) -> Node {
self.var(Var::Y)
}
pub fn z(&mut self) -> Node {
self.var(Var::Z)
}
pub fn var(&mut self, v: Var) -> Node {
self.ops.insert(Op::Input(v))
}
pub fn axes(&mut self) -> [Node; 3] {
[self.x(), self.y(), self.z()]
}
pub fn constant(&mut self, f: f64) -> Node {
self.ops.insert(Op::Const(OrderedFloat(f)))
}
fn op_unary(&mut self, a: Node, op: UnaryOpcode) -> Result<Node, BadNode> {
let op_a = *self.get_op(a).ok_or(BadNode)?;
let out = if let Op::Const(a) = op_a {
self.constant(op.eval(a.0))
} else {
self.ops.insert(Op::Unary(op, a))
};
Ok(out)
}
fn op_binary(
&mut self,
a: Node,
b: Node,
op: BinaryOpcode,
) -> Result<Node, BadNode> {
let op_a = *self.get_op(a).ok_or(BadNode)?;
let op_b = *self.get_op(b).ok_or(BadNode)?;
let out = if let (Op::Const(a), Op::Const(b)) = (op_a, op_b) {
self.constant(op.eval(a.0, b.0))
} else {
self.ops.insert(Op::Binary(op, a, b))
};
Ok(out)
}
fn op_binary_commutative(
&mut self,
a: Node,
b: Node,
op: BinaryOpcode,
) -> Result<Node, BadNode> {
self.op_binary(a.min(b), a.max(b), op)
}
pub fn add<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a: Node = a.into_node(self)?;
let b: Node = b.into_node(self)?;
if a == b {
let two = self.constant(2.0);
self.mul(a, two)
} else {
match (self.get_const(a), self.get_const(b)) {
(Ok(0.0), _) => Ok(b),
(_, Ok(0.0)) => Ok(a),
_ => self.op_binary_commutative(a, b, BinaryOpcode::Add),
}
}
}
pub fn mul<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
if a == b {
self.square(a)
} else {
match (self.get_const(a), self.get_const(b)) {
(Ok(1.0), _) => Ok(b),
(_, Ok(1.0)) => Ok(a),
(Ok(0.0), _) => Ok(a),
(_, Ok(0.0)) => Ok(b),
_ => self.op_binary_commutative(a, b, BinaryOpcode::Mul),
}
}
}
pub fn min<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
if a == b {
Ok(a)
} else {
self.op_binary_commutative(a, b, BinaryOpcode::Min)
}
}
pub fn max<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
if a == b {
Ok(a)
} else {
self.op_binary_commutative(a, b, BinaryOpcode::Max)
}
}
pub fn and<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
let op_a = *self.get_op(a).ok_or(BadNode)?;
if let Op::Const(v) = op_a {
if v.0 == 0.0 { Ok(a) } else { Ok(b) }
} else {
self.op_binary(a, b, BinaryOpcode::And)
}
}
pub fn or<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
let op_a = *self.get_op(a).ok_or(BadNode)?;
let op_b = *self.get_op(b).ok_or(BadNode)?;
if let Op::Const(v) = op_a {
if v.0 != 0.0 {
return Ok(a);
} else {
return Ok(b);
}
} else if let Op::Const(v) = op_b
&& v.0 == 0.0
{
return Ok(a);
}
self.op_binary(a, b, BinaryOpcode::Or)
}
pub fn not<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Not)
}
pub fn neg<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Neg)
}
pub fn recip<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Recip)
}
pub fn abs<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Abs)
}
pub fn sqrt<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Sqrt)
}
pub fn sin<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Sin)
}
pub fn cos<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Cos)
}
pub fn tan<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Tan)
}
pub fn asin<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Asin)
}
pub fn acos<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Acos)
}
pub fn atan<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Atan)
}
pub fn exp<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Exp)
}
pub fn ln<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Ln)
}
pub fn square<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Square)
}
pub fn floor<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Floor)
}
pub fn ceil<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Ceil)
}
pub fn round<A: IntoNode>(&mut self, a: A) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
self.op_unary(a, UnaryOpcode::Round)
}
pub fn sub<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
match (self.get_const(a), self.get_const(b)) {
(Ok(0.0), _) => self.neg(b),
(_, Ok(0.0)) => Ok(a),
_ => self.op_binary(a, b, BinaryOpcode::Sub),
}
}
pub fn div<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
match (self.get_const(a), self.get_const(b)) {
(Ok(0.0), _) => Ok(a),
(_, Ok(1.0)) => Ok(a),
_ => self.op_binary(a, b, BinaryOpcode::Div),
}
}
pub fn atan2<A: IntoNode, B: IntoNode>(
&mut self,
y: A,
x: B,
) -> Result<Node, BadNode> {
let y = y.into_node(self)?;
let x = x.into_node(self)?;
self.op_binary(y, x, BinaryOpcode::Atan)
}
pub fn compare<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
self.op_binary(a, b, BinaryOpcode::Compare)
}
pub fn less_than<A: IntoNode, B: IntoNode>(
&mut self,
lhs: A,
rhs: B,
) -> Result<Node, BadNode> {
let lhs = lhs.into_node(self)?;
let rhs = rhs.into_node(self)?;
let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
self.max(cmp, 0.0)
}
pub fn less_than_or_equal<A: IntoNode, B: IntoNode>(
&mut self,
lhs: A,
rhs: B,
) -> Result<Node, BadNode> {
let lhs = lhs.into_node(self)?;
let rhs = rhs.into_node(self)?;
let cmp = self.op_binary(rhs, lhs, BinaryOpcode::Compare)?;
let shift = self.add(cmp, 1.0)?;
self.min(shift, 1.0)
}
pub fn modulo<A: IntoNode, B: IntoNode>(
&mut self,
a: A,
b: B,
) -> Result<Node, BadNode> {
let a = a.into_node(self)?;
let b = b.into_node(self)?;
self.op_binary(a, b, BinaryOpcode::Mod)
}
pub fn if_nonzero_else<Condition: IntoNode, A: IntoNode, B: IntoNode>(
&mut self,
condition: Condition,
a: A,
b: B,
) -> Result<Node, BadNode> {
let condition = condition.into_node(self)?;
let a = a.into_node(self)?;
let b = b.into_node(self)?;
let lhs = self.and(condition, a)?;
let n_condition = self.not(condition)?;
let rhs = self.and(n_condition, b)?;
self.or(lhs, rhs)
}
pub fn eval_xyz(
&self,
root: Node,
x: f64,
y: f64,
z: f64,
) -> Result<f64, EvalError> {
let vars = [(Var::X, x), (Var::Y, y), (Var::Z, z)]
.into_iter()
.collect();
self.eval(root, &vars)
}
pub fn eval(
&self,
root: Node,
vars: &HashMap<Var, f64>,
) -> Result<f64, EvalError> {
let mut cache = vec![None; self.ops.len()].into();
self.eval_inner(root, vars, &mut cache)
}
fn eval_inner(
&self,
node: Node,
vars: &HashMap<Var, f64>,
cache: &mut IndexVec<Option<f64>, Node>,
) -> Result<f64, EvalError> {
if node.0 >= cache.len() {
return Err(EvalError::BadNode(BadNode));
}
if let Some(v) = cache[node] {
return Ok(v);
}
let mut get = |n: Node| self.eval_inner(n, vars, cache);
let v = match self.get_op(node).ok_or(EvalError::BadNode(BadNode))? {
Op::Input(v) => *vars.get(v).ok_or(EvalError::MissingVar(*v))?,
Op::Const(c) => c.0,
Op::Binary(op, a, b) => {
let a = get(*a)?;
let b = get(*b)?;
op.eval(a, b)
}
Op::Unary(op, a) => {
let a = get(*a)?;
op.eval(a)
}
};
cache[node] = Some(v);
Ok(v)
}
pub fn from_text<R: Read>(r: R) -> Result<(Self, Node), ParseError> {
let reader = BufReader::new(r);
let mut ctx = Self::new();
let mut seen = BTreeMap::new();
let mut last = None;
for line in reader.lines().map(|line| line.unwrap()) {
if line.is_empty() || line.starts_with('#') {
continue;
}
let mut iter = line.split_whitespace();
let i: String = iter.next().unwrap().to_owned();
let opcode = iter.next().unwrap();
let mut pop = || {
let txt = iter.next().unwrap();
seen.get(txt)
.cloned()
.ok_or_else(|| ParseError::UnknownVariable(txt.to_string()))
};
let node = match opcode {
"const" => ctx.constant(iter.next().unwrap().parse().unwrap()),
"var-x" => ctx.x(),
"var-y" => ctx.y(),
"var-z" => ctx.z(),
"abs" => ctx.abs(pop()?)?,
"neg" => ctx.neg(pop()?)?,
"sqrt" => ctx.sqrt(pop()?)?,
"square" => ctx.square(pop()?)?,
"floor" => ctx.floor(pop()?)?,
"ceil" => ctx.ceil(pop()?)?,
"round" => ctx.round(pop()?)?,
"sin" => ctx.sin(pop()?)?,
"cos" => ctx.cos(pop()?)?,
"tan" => ctx.tan(pop()?)?,
"asin" => ctx.asin(pop()?)?,
"acos" => ctx.acos(pop()?)?,
"atan" => ctx.atan(pop()?)?,
"ln" => ctx.ln(pop()?)?,
"not" => ctx.not(pop()?)?,
"exp" => ctx.exp(pop()?)?,
"add" => ctx.add(pop()?, pop()?)?,
"mul" => ctx.mul(pop()?, pop()?)?,
"min" => ctx.min(pop()?, pop()?)?,
"max" => ctx.max(pop()?, pop()?)?,
"div" => ctx.div(pop()?, pop()?)?,
"atan2" => ctx.atan2(pop()?, pop()?)?,
"sub" => ctx.sub(pop()?, pop()?)?,
"compare" => ctx.compare(pop()?, pop()?)?,
"mod" => ctx.modulo(pop()?, pop()?)?,
"and" => ctx.and(pop()?, pop()?)?,
"or" => ctx.or(pop()?, pop()?)?,
op => return Err(ParseError::UnknownOpcode(op.to_owned())),
};
seen.insert(i, node);
last = Some(node);
}
match last {
Some(node) => Ok((ctx, node)),
None => Err(ParseError::EmptyFile),
}
}
pub fn dot(&self) -> String {
let mut out = "digraph mygraph{\n".to_owned();
for node in self.ops.keys() {
let op = self.get_op(node).unwrap();
out += &self.dot_node(node);
out += &op.dot_edges(node);
}
out += "}\n";
out
}
fn dot_node(&self, i: Node) -> String {
let mut out = format!(r#"n{} [label = ""#, i.get());
let op = self.get_op(i).unwrap();
match op {
Op::Const(c) => write!(out, "{c}").unwrap(),
Op::Input(v) => {
out += &v.to_string();
}
Op::Binary(op, ..) => match op {
BinaryOpcode::Add => out += "add",
BinaryOpcode::Sub => out += "sub",
BinaryOpcode::Mul => out += "mul",
BinaryOpcode::Div => out += "div",
BinaryOpcode::Atan => out += "atan2",
BinaryOpcode::Min => out += "min",
BinaryOpcode::Max => out += "max",
BinaryOpcode::Compare => out += "compare",
BinaryOpcode::Mod => out += "mod",
BinaryOpcode::And => out += "and",
BinaryOpcode::Or => out += "or",
},
Op::Unary(op, ..) => match op {
UnaryOpcode::Neg => out += "neg",
UnaryOpcode::Abs => out += "abs",
UnaryOpcode::Recip => out += "recip",
UnaryOpcode::Sqrt => out += "sqrt",
UnaryOpcode::Square => out += "square",
UnaryOpcode::Floor => out += "floor",
UnaryOpcode::Ceil => out += "ceil",
UnaryOpcode::Round => out += "round",
UnaryOpcode::Sin => out += "sin",
UnaryOpcode::Cos => out += "cos",
UnaryOpcode::Tan => out += "tan",
UnaryOpcode::Asin => out += "asin",
UnaryOpcode::Acos => out += "acos",
UnaryOpcode::Atan => out += "atan",
UnaryOpcode::Exp => out += "exp",
UnaryOpcode::Ln => out += "ln",
UnaryOpcode::Not => out += "not",
},
};
write!(
out,
r#"" color="{0}1" shape="{1}" fontcolor="{0}4"]"#,
op.dot_node_color(),
op.dot_node_shape()
)
.unwrap();
out
}
pub fn get_op(&self, node: Node) -> Option<&Op> {
self.ops.get_by_index(node)
}
pub fn import(&mut self, tree: &Tree) -> Node {
enum Action<'a> {
Down(&'a Arc<TreeOp>),
Up(&'a Arc<TreeOp>),
Pop,
PopAffine,
}
let mut axes = vec![(self.x(), self.y(), self.z())];
let mut todo = vec![Action::Down(tree.arc())];
let mut stack = vec![];
let mut affine: Vec<Matrix4<f64>> = vec![];
let mut seen = HashMap::new();
while let Some(t) = todo.pop() {
match t {
Action::Down(t) => {
if matches!(
t.as_ref(),
TreeOp::Unary(..) | TreeOp::Binary(..)
) && let Some(p) =
seen.get(&(*axes.last().unwrap(), Arc::as_ptr(t)))
{
stack.push(*p);
continue;
}
match t.as_ref() {
TreeOp::Const(c) => {
stack.push(self.constant(*c));
}
TreeOp::Input(s) => {
let axes = axes.last().unwrap();
stack.push(match *s {
Var::X => axes.0,
Var::Y => axes.1,
Var::Z => axes.2,
v @ Var::V(..) => self.var(v),
});
}
TreeOp::Unary(_op, arg) => {
todo.push(Action::Up(t));
todo.push(Action::Down(arg));
}
TreeOp::Binary(_op, lhs, rhs) => {
todo.push(Action::Up(t));
todo.push(Action::Down(lhs));
todo.push(Action::Down(rhs));
}
TreeOp::RemapAxes { target: _, x, y, z } => {
todo.push(Action::Up(t));
todo.push(Action::Down(x));
todo.push(Action::Down(y));
todo.push(Action::Down(z));
}
TreeOp::RemapAffine { target, mat } => {
let prev = affine
.last()
.cloned()
.unwrap_or(Matrix4::identity());
let mat = prev * mat.to_homogeneous();
if matches!(&**target, TreeOp::RemapAffine { .. }) {
affine.push(mat);
todo.push(Action::PopAffine);
} else {
let (x, y, z) = axes.last().unwrap();
let mut out = [None; 3];
for i in 0..3 {
let a = self.mul(mat[(i, 0)], *x).unwrap();
let b = self.mul(mat[(i, 1)], *y).unwrap();
let c = self.mul(mat[(i, 2)], *z).unwrap();
let d = self.constant(mat[(i, 3)]);
let ab = self.add(a, b).unwrap();
let cd = self.add(c, d).unwrap();
out[i] = Some(self.add(ab, cd).unwrap());
}
let [x, y, z] = out.map(Option::unwrap);
axes.push((x, y, z));
todo.push(Action::Pop);
}
todo.push(Action::Down(target));
}
}
}
Action::Up(t) => {
match t.as_ref() {
TreeOp::Const(..)
| TreeOp::Input(..)
| TreeOp::RemapAffine { .. } => unreachable!(),
TreeOp::Unary(op, ..) => {
let arg = stack.pop().unwrap();
let out = self.op_unary(arg, *op).unwrap();
stack.push(out);
}
TreeOp::Binary(op, ..) => {
let lhs = stack.pop().unwrap();
let rhs = stack.pop().unwrap();
let out = match op {
BinaryOpcode::Add => self.add(lhs, rhs),
BinaryOpcode::Sub => self.sub(lhs, rhs),
BinaryOpcode::Mul => self.mul(lhs, rhs),
BinaryOpcode::Div => self.div(lhs, rhs),
BinaryOpcode::Atan => self.atan2(lhs, rhs),
BinaryOpcode::Min => self.min(lhs, rhs),
BinaryOpcode::Max => self.max(lhs, rhs),
BinaryOpcode::Compare => self.compare(lhs, rhs),
BinaryOpcode::Mod => self.modulo(lhs, rhs),
BinaryOpcode::And => self.and(lhs, rhs),
BinaryOpcode::Or => self.or(lhs, rhs),
}
.unwrap();
if Arc::strong_count(t) > 1 {
seen.insert(
(*axes.last().unwrap(), Arc::as_ptr(t)),
out,
);
}
stack.push(out);
}
TreeOp::RemapAxes { target, .. } => {
let x = stack.pop().unwrap();
let y = stack.pop().unwrap();
let z = stack.pop().unwrap();
axes.push((x, y, z));
todo.push(Action::Pop);
todo.push(Action::Down(target));
}
}
if matches!(
t.as_ref(),
TreeOp::Unary(..) | TreeOp::Binary(..)
) && Arc::strong_count(t) > 1
{
seen.insert(
(*axes.last().unwrap(), Arc::as_ptr(t)),
*stack.last().unwrap(),
);
}
}
Action::Pop => {
axes.pop().unwrap();
}
Action::PopAffine => {
affine.pop().unwrap();
}
}
}
assert_eq!(stack.len(), 1);
stack.pop().unwrap()
}
pub fn export(&self, n: Node) -> Result<Tree, BadNode> {
if self.get_op(n).is_none() {
return Err(BadNode);
}
enum Action {
Down(Node),
Up(Node, Op),
}
let mut todo = vec![Action::Down(n)];
let mut stack = vec![];
let mut seen: HashMap<Node, Tree> = HashMap::new();
while let Some(t) = todo.pop() {
match t {
Action::Down(n) => {
if let Some(p) = seen.get(&n) {
stack.push(p.clone());
continue;
}
let op = self.get_op(n).unwrap();
match op {
Op::Const(c) => {
let t = Tree::from(c.0);
seen.insert(n, t.clone());
stack.push(t);
}
Op::Input(v) => {
let t = Tree::from(*v);
seen.insert(n, t.clone());
stack.push(t);
}
Op::Unary(_op, arg) => {
todo.push(Action::Up(n, *op));
todo.push(Action::Down(*arg));
}
Op::Binary(_op, lhs, rhs) => {
todo.push(Action::Up(n, *op));
todo.push(Action::Down(*lhs));
todo.push(Action::Down(*rhs));
}
}
}
Action::Up(n, op) => match op {
Op::Const(..) | Op::Input(..) => unreachable!(),
Op::Unary(op, ..) => {
let arg = stack.pop().unwrap();
let out =
Tree::from(TreeOp::Unary(op, arg.arc().clone()));
seen.insert(n, out.clone());
stack.push(out);
}
Op::Binary(op, ..) => {
let lhs = stack.pop().unwrap();
let rhs = stack.pop().unwrap();
let out = Tree::from(TreeOp::Binary(
op,
lhs.arc().clone(),
rhs.arc().clone(),
));
seen.insert(n, out.clone());
stack.push(out);
}
},
}
}
assert_eq!(stack.len(), 1);
Ok(stack.pop().unwrap())
}
pub fn deriv(&mut self, n: Node, v: Var) -> Result<Node, BadNode> {
if self.get_op(n).is_none() {
return Err(BadNode);
}
enum Action {
Down(Node),
Up(Node, Op),
}
let mut todo = vec![Action::Down(n)];
let mut stack = vec![];
let zero = self.constant(0.0);
let mut seen: HashMap<Node, Node> = HashMap::new();
while let Some(t) = todo.pop() {
match t {
Action::Down(n) => {
if let Some(p) = seen.get(&n) {
stack.push(*p);
continue;
}
let op = *self.get_op(n).unwrap();
match op {
Op::Const(_c) => {
seen.insert(n, zero);
stack.push(zero);
}
Op::Input(u) => {
let z =
if v == u { self.constant(1.0) } else { zero };
seen.insert(n, z);
stack.push(z);
}
Op::Unary(_op, arg) => {
todo.push(Action::Up(n, op));
todo.push(Action::Down(arg));
}
Op::Binary(_op, lhs, rhs) => {
todo.push(Action::Up(n, op));
todo.push(Action::Down(lhs));
todo.push(Action::Down(rhs));
}
}
}
Action::Up(n, op) => match op {
Op::Const(..) | Op::Input(..) => unreachable!(),
Op::Unary(op, v_arg) => {
let d_arg = stack.pop().unwrap();
let out = match op {
UnaryOpcode::Neg => self.neg(d_arg),
UnaryOpcode::Abs => {
let cond = self.less_than(v_arg, zero).unwrap();
let pos = d_arg;
let neg = self.neg(d_arg).unwrap();
self.if_nonzero_else(cond, neg, pos)
}
UnaryOpcode::Recip => {
let a = self.square(v_arg).unwrap();
let b = self.neg(d_arg).unwrap();
self.div(b, a)
}
UnaryOpcode::Sqrt => {
let v = self.mul(n, 2.0).unwrap();
self.div(d_arg, v)
}
UnaryOpcode::Square => {
let v = self.mul(d_arg, v_arg).unwrap();
self.mul(2.0, v)
}
UnaryOpcode::Floor
| UnaryOpcode::Ceil
| UnaryOpcode::Round => Ok(zero),
UnaryOpcode::Sin => {
let c = self.cos(v_arg).unwrap();
self.mul(c, d_arg)
}
UnaryOpcode::Cos => {
let s = self.sin(v_arg).unwrap();
let s = self.neg(s).unwrap();
self.mul(s, d_arg)
}
UnaryOpcode::Tan => {
let c = self.cos(v_arg).unwrap();
let c = self.square(c).unwrap();
self.div(d_arg, c)
}
UnaryOpcode::Asin => {
let v = self.square(v_arg).unwrap();
let v = self.sub(1.0, v).unwrap();
let v = self.sqrt(v).unwrap();
self.div(d_arg, v)
}
UnaryOpcode::Acos => {
let v = self.square(v_arg).unwrap();
let v = self.sub(1.0, v).unwrap();
let v = self.sqrt(v).unwrap();
let v = self.neg(v).unwrap();
self.div(d_arg, v)
}
UnaryOpcode::Atan => {
let v = self.square(v_arg).unwrap();
let v = self.add(1.0, v).unwrap();
self.div(d_arg, v)
}
UnaryOpcode::Exp => self.mul(n, d_arg),
UnaryOpcode::Ln => self.div(d_arg, v_arg),
UnaryOpcode::Not => Ok(zero),
}
.unwrap();
seen.insert(n, out);
stack.push(out);
}
Op::Binary(op, v_lhs, v_rhs) => {
let d_lhs = stack.pop().unwrap();
let d_rhs = stack.pop().unwrap();
let out = match op {
BinaryOpcode::Add => self.add(d_lhs, d_rhs),
BinaryOpcode::Sub => self.sub(d_lhs, d_rhs),
BinaryOpcode::Mul => {
let a = self.mul(d_lhs, v_rhs).unwrap();
let b = self.mul(v_lhs, d_rhs).unwrap();
self.add(a, b)
}
BinaryOpcode::Div => {
let v = self.square(v_rhs).unwrap();
let a = self.mul(v_rhs, d_lhs).unwrap();
let b = self.mul(v_lhs, d_rhs).unwrap();
let c = self.sub(a, b).unwrap();
self.div(c, v)
}
BinaryOpcode::Atan => {
let a = self.square(v_lhs).unwrap();
let b = self.square(v_rhs).unwrap();
let d = self.add(a, b).unwrap();
let a = self.mul(v_rhs, d_lhs).unwrap();
let b = self.mul(v_lhs, d_rhs).unwrap();
let v = self.sub(a, b).unwrap();
self.div(v, d)
}
BinaryOpcode::Min => {
let cond =
self.less_than(v_lhs, v_rhs).unwrap();
self.if_nonzero_else(cond, d_lhs, d_rhs)
}
BinaryOpcode::Max => {
let cond =
self.less_than(v_rhs, v_lhs).unwrap();
self.if_nonzero_else(cond, d_lhs, d_rhs)
}
BinaryOpcode::Compare => Ok(zero),
BinaryOpcode::Mod => {
let e = self.div(v_lhs, v_rhs).unwrap();
let q = self.floor(e).unwrap();
let m = self.modulo(q, v_rhs).unwrap();
let cond = self.less_than(q, zero).unwrap();
let offset = self
.if_nonzero_else(cond, v_rhs, zero)
.unwrap();
let m = self.sub(m, offset).unwrap();
let outer = self.less_than(m, zero).unwrap();
let inner =
self.less_than(zero, v_rhs).unwrap();
let qa = self.sub(q, 1.0).unwrap();
let qb = self.add(q, 1.0).unwrap();
let inner = self
.if_nonzero_else(inner, qa, qb)
.unwrap();
let e = self
.if_nonzero_else(outer, inner, q)
.unwrap();
let v = self.mul(d_rhs, e).unwrap();
self.sub(d_lhs, v)
}
BinaryOpcode::And => {
let cond = self.compare(v_lhs, zero).unwrap();
self.if_nonzero_else(cond, d_rhs, d_lhs)
}
BinaryOpcode::Or => {
let cond = self.compare(v_lhs, zero).unwrap();
self.if_nonzero_else(cond, d_lhs, d_rhs)
}
}
.unwrap();
seen.insert(n, out);
stack.push(out);
}
},
}
}
assert_eq!(stack.len(), 1);
Ok(stack.pop().unwrap())
}
}
#[derive(thiserror::Error, Debug)]
#[error("node is not present in this `Context`")]
pub struct BadNode;
#[derive(thiserror::Error, Debug)]
pub enum ParseError {
#[error("unknown opcode {0}")]
UnknownOpcode(String),
#[error("unknown variable {0}")]
UnknownVariable(String),
#[error(transparent)]
BadNode(#[from] BadNode),
#[error("empty file")]
EmptyFile,
}
#[derive(thiserror::Error, Debug)]
pub enum ConstError {
#[error("node is not a constant")]
NotAConst,
#[error(transparent)]
BadNode(#[from] BadNode),
}
#[derive(thiserror::Error, Debug)]
#[error("node does not have an associated variable")]
pub struct NotAVar;
#[derive(thiserror::Error, Debug)]
pub enum VarError {
#[error(transparent)]
NotAVar(#[from] NotAVar),
#[error(transparent)]
BadNode(#[from] BadNode),
}
#[derive(thiserror::Error, Debug)]
pub enum EvalError {
#[error("variable {0} is missing in the evaluation map")]
MissingVar(Var),
#[error(transparent)]
BadNode(#[from] BadNode),
}
pub trait IntoNode {
fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode>;
}
impl IntoNode for Node {
fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
ctx.check_node(self)?;
Ok(self)
}
}
impl IntoNode for f32 {
fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
Ok(ctx.constant(self as f64))
}
}
impl IntoNode for f64 {
fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
Ok(ctx.constant(self))
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::vm::VmData;
#[test]
fn test_get_op() {
let mut ctx = Context::new();
let x = ctx.x();
let op_x = ctx.get_op(x).unwrap();
assert!(matches!(op_x, Op::Input(_)));
}
#[test]
fn test_ring() {
let mut ctx = Context::new();
let c0 = ctx.constant(0.5);
let x = ctx.x();
let y = ctx.y();
let x2 = ctx.square(x).unwrap();
let y2 = ctx.square(y).unwrap();
let r = ctx.add(x2, y2).unwrap();
let c6 = ctx.sub(r, c0).unwrap();
let c7 = ctx.constant(0.25);
let c8 = ctx.sub(c7, r).unwrap();
let c9 = ctx.max(c8, c6).unwrap();
let tape = VmData::<255>::new(&ctx, &[c9]).unwrap();
assert_eq!(tape.len(), 9);
assert_eq!(tape.vars.len(), 2);
}
#[test]
fn test_dupe() {
let mut ctx = Context::new();
let x = ctx.x();
let x_squared = ctx.mul(x, x).unwrap();
let tape = VmData::<255>::new(&ctx, &[x_squared]).unwrap();
assert_eq!(tape.len(), 3); assert_eq!(tape.vars.len(), 1);
}
#[test]
fn test_export() {
let mut ctx = Context::new();
let x = ctx.x();
let s = ctx.sin(x).unwrap();
let c = ctx.cos(x).unwrap();
let sum = ctx.add(s, c).unwrap();
let t = ctx.export(sum).unwrap();
if let TreeOp::Binary(BinaryOpcode::Add, lhs, rhs) = &*t {
match (&**lhs, &**rhs) {
(
TreeOp::Unary(UnaryOpcode::Sin, x1),
TreeOp::Unary(UnaryOpcode::Cos, x2),
) => {
assert_eq!(Arc::as_ptr(x1), Arc::as_ptr(x2));
let TreeOp::Input(Var::X) = &**x1 else {
panic!("invalid X: {x1:?}");
};
}
_ => panic!("invalid lhs / rhs: {lhs:?} {rhs:?}"),
}
} else {
panic!("unexpected opcode {t:?}");
}
}
#[test]
fn import_optimization() {
let t = Tree::x() + 0;
let mut ctx = Context::new();
let root = ctx.import(&t);
assert_eq!(ctx.get_op(root).unwrap(), &Op::Input(Var::X));
}
}