1use crate::context::{BadNode, Context, IntoNode, Node};
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11
12#[derive(
21 Copy,
22 Clone,
23 Debug,
24 Hash,
25 Eq,
26 PartialEq,
27 Ord,
28 PartialOrd,
29 Serialize,
30 Deserialize,
31)]
32pub enum Var {
33 X,
35 Y,
37 Z,
39 V(VarIndex),
41}
42
43#[derive(
45 Copy,
46 Clone,
47 Debug,
48 Hash,
49 Eq,
50 PartialEq,
51 Ord,
52 PartialOrd,
53 Serialize,
54 Deserialize,
55)]
56#[serde(transparent)]
57pub struct VarIndex(u64);
58
59impl Var {
60 #[allow(clippy::new_without_default)]
66 pub fn new() -> Self {
67 let v: u64 = rand::random();
68 Var::V(VarIndex(v))
69 }
70
71 pub fn index(&self) -> Option<VarIndex> {
73 if let Var::V(i) = *self { Some(i) } else { None }
74 }
75}
76
77impl std::fmt::Display for Var {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 match self {
80 Var::X => write!(f, "X"),
81 Var::Y => write!(f, "Y"),
82 Var::Z => write!(f, "Z"),
83 Var::V(VarIndex(v)) if *v < 256 => write!(f, "v_{v}"),
84 Var::V(VarIndex(v)) => write!(f, "V({v:x})"),
85 }
86 }
87}
88
89impl IntoNode for Var {
90 fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
91 Ok(ctx.var(self))
92 }
93}
94
95#[derive(Debug, Default, Serialize, Deserialize)]
105pub struct VarMap {
106 x: Option<usize>,
107 y: Option<usize>,
108 z: Option<usize>,
109 v: HashMap<VarIndex, usize>,
110}
111
112#[allow(missing_docs)]
113impl VarMap {
114 pub fn new() -> Self {
115 Self::default()
116 }
117 pub fn len(&self) -> usize {
118 self.x.is_some() as usize
119 + self.y.is_some() as usize
120 + self.z.is_some() as usize
121 + self.v.len()
122 }
123 pub fn is_empty(&self) -> bool {
124 self.x.is_none()
125 && self.y.is_none()
126 && self.z.is_none()
127 && self.v.is_empty()
128 }
129 pub fn get(&self, v: &Var) -> Option<usize> {
130 match v {
131 Var::X => self.x,
132 Var::Y => self.y,
133 Var::Z => self.z,
134 Var::V(v) => self.v.get(v).cloned(),
135 }
136 }
137
138 pub fn has_free_vars(&self) -> bool {
140 !self.v.is_empty()
141 }
142
143 pub fn insert(&mut self, v: Var) {
147 let next = self.len();
148 match v {
149 Var::X => self.x.get_or_insert(next),
150 Var::Y => self.y.get_or_insert(next),
151 Var::Z => self.z.get_or_insert(next),
152 Var::V(v) => self.v.entry(v).or_insert(next),
153 };
154 }
155
156 pub fn check_tracing_arguments<T>(
158 &self,
159 vars: &[T],
160 ) -> Result<(), TracingArgError> {
161 if vars.len() < self.len() {
162 Err(TracingArgError::BadVarSlice(BadVarSlice {
165 actual: vars.len(),
166 expected: self.len(),
167 }))
168 } else {
169 Ok(())
170 }
171 }
172
173 pub fn check_bulk_arguments<T, V: std::ops::Deref<Target = [T]>>(
175 &self,
176 vars: &[V],
177 ) -> Result<(), BulkArgError> {
178 if vars.len() < self.len() {
181 Err(BulkArgError::BadVarSlice(BadVarSlice {
182 actual: vars.len(),
183 expected: self.len(),
184 }))
185 } else {
186 let Some(n) = vars.first().map(|v| v.len()) else {
187 return Ok(());
188 };
189 if let Some((i, v)) =
190 vars.iter().enumerate().find(|(_i, v)| v.len() != n)
191 {
192 Err(MismatchedSlices {
193 index_a: 0,
194 length_a: n,
195 index_b: i,
196 length_b: v.len(),
197 }
198 .into())
199 } else {
200 Ok(())
201 }
202 }
203 }
204
205 pub fn iter(&self) -> impl Iterator<Item = (Var, usize)> {
207 [
208 self.x.map(|x| (Var::X, x)),
209 self.y.map(|y| (Var::Y, y)),
210 self.z.map(|z| (Var::Z, z)),
211 ]
212 .into_iter()
213 .flatten()
214 .chain(self.v.iter().map(|(v, k)| (Var::V(*v), *k)))
215 }
216}
217
218#[derive(thiserror::Error, Debug)]
220#[error(
221 "variable slice length ({actual}) does not match \
222 expected count ({expected})"
223)]
224pub struct BadVarSlice {
225 pub actual: usize,
227 pub expected: usize,
229}
230
231#[derive(thiserror::Error, Debug)]
233#[error(
234 "slice lengths are mismatched: \
235 slice {index_a} has {length_a} items; \
236 slice {index_b} has {length_b} items;"
237)]
238pub struct MismatchedSlices {
239 index_a: usize,
240 length_a: usize,
241 index_b: usize,
242 length_b: usize,
243}
244
245#[derive(thiserror::Error, Debug)]
247pub enum BulkArgError {
248 #[error(transparent)]
250 BadVarSlice(#[from] BadVarSlice),
251
252 #[error(transparent)]
254 MismatchedSlices(#[from] MismatchedSlices),
255}
256
257#[derive(thiserror::Error, Debug)]
259pub enum TracingArgError {
260 #[error(transparent)]
262 BadVarSlice(BadVarSlice),
263}
264
265impl std::ops::Index<&Var> for VarMap {
266 type Output = usize;
267 fn index(&self, v: &Var) -> &Self::Output {
268 match v {
269 Var::X => self.x.as_ref().unwrap(),
270 Var::Y => self.y.as_ref().unwrap(),
271 Var::Z => self.z.as_ref().unwrap(),
272 Var::V(v) => &self.v[v],
273 }
274 }
275}
276
277#[cfg(test)]
278mod test {
279 use super::*;
280
281 #[test]
282 fn var_identity() {
283 let v1 = Var::new();
284 let v2 = Var::new();
285 assert_ne!(v1, v2);
286 }
287
288 #[test]
289 fn var_map() {
290 let v = Var::new();
291 let mut m = VarMap::new();
292 assert!(m.get(&v).is_none());
293 m.insert(v);
294 assert_eq!(m.get(&v), Some(0));
295 m.insert(v);
296 assert_eq!(m.get(&v), Some(0));
297
298 let u = Var::new();
299 assert!(m.get(&u).is_none());
300 m.insert(u);
301 assert_eq!(m.get(&u), Some(1));
302 }
303}