1use smallvec::SmallVec;
2
3#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
4pub struct ExprId(pub u32);
5
6impl ExprId {
7 #[inline]
8 pub fn index(self) -> usize {
9 self.0 as usize
10 }
11}
12
13#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
14pub struct VarId(pub u32);
15
16impl VarId {
17 #[inline]
18 pub fn index(self) -> usize {
19 self.0 as usize
20 }
21}
22
23#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
24pub struct ParamId(pub u32);
25
26impl ParamId {
27 #[inline]
28 pub fn index(self) -> usize {
29 self.0 as usize
30 }
31}
32
33pub type Children = SmallVec<[ExprId; 4]>;
34
35#[derive(Clone, Debug)]
40pub enum ExprNode {
41 Const(f64),
42 Var(VarId),
43 Param(ParamId),
44 Add(Children),
45 Mul(Children),
46 Neg(ExprId),
47 Pow(ExprId, ExprId),
48 Sin(ExprId),
49 Cos(ExprId),
50 Exp(ExprId),
51 Log(ExprId),
52 Linear { coeffs: Vec<(VarId, f64)>, constant: f64 },
53}
54
55#[derive(Clone, Debug, Default)]
56pub struct ExprArena {
57 nodes: Vec<ExprNode>,
58}
59
60impl ExprArena {
61 pub fn new() -> Self {
62 Self::default()
63 }
64
65 pub fn with_capacity(cap: usize) -> Self {
66 Self { nodes: Vec::with_capacity(cap) }
67 }
68
69 #[inline]
70 pub fn len(&self) -> usize {
71 self.nodes.len()
72 }
73
74 #[inline]
75 pub fn is_empty(&self) -> bool {
76 self.nodes.is_empty()
77 }
78
79 pub fn push(&mut self, node: ExprNode) -> ExprId {
83 let id = ExprId(u32::try_from(self.nodes.len()).expect("expression arena overflow"));
84 self.nodes.push(node);
85 id
86 }
87
88 #[inline]
89 pub fn get(&self, id: ExprId) -> &ExprNode {
90 &self.nodes[id.index()]
91 }
92
93 #[inline]
94 pub fn get_mut(&mut self, id: ExprId) -> &mut ExprNode {
95 &mut self.nodes[id.index()]
96 }
97
98 pub fn nodes(&self) -> &[ExprNode] {
99 &self.nodes
100 }
101
102 pub fn constant(&mut self, v: f64) -> ExprId {
103 self.push(ExprNode::Const(v))
104 }
105
106 pub fn var(&mut self, v: VarId) -> ExprId {
107 self.push(ExprNode::Var(v))
108 }
109
110 pub fn param(&mut self, p: ParamId) -> ExprId {
111 self.push(ExprNode::Param(p))
112 }
113
114 pub fn linear(&mut self, coeffs: Vec<(VarId, f64)>, constant: f64) -> ExprId {
115 self.push(ExprNode::Linear { coeffs, constant })
116 }
117}