1use super::ir::{ArrayNode, Assignment, ConstraintEvalError, Expr, FuncId, Ty, eval, mask};
41use num_bigint::BigUint;
42use num_traits::Zero;
43use std::fmt;
44
45pub const MAX_STEPS: u64 = 1 << 22;
47pub const MAX_DEPTH: usize = 64;
49
50#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
52pub enum Stmt {
53 Assign(u32, Expr),
55 If(Expr, Vec<Stmt>, Vec<Stmt>),
57 Loop {
63 cond: Option<Expr>,
65 body: Vec<Stmt>,
67 step: Vec<Stmt>,
69 test_first: bool,
71 },
72 Break,
74 Continue,
76 Return,
78}
79
80#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
82pub struct Func {
83 pub name: String,
85 pub slots: Vec<Ty>,
88 pub n_params: usize,
90 pub ret: u32,
92 pub body: Vec<Stmt>,
94}
95
96#[derive(Clone, Debug, PartialEq, Eq)]
98pub enum FuncError {
99 Eval(ConstraintEvalError),
101 Malformed(&'static str),
104 TooDeep {
106 func: String,
108 },
109 StepLimit {
111 func: String,
113 },
114}
115
116impl fmt::Display for FuncError {
117 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118 match self {
119 FuncError::Eval(e) => write!(f, "a function called from a constraint failed: {e}"),
120 FuncError::Malformed(w) => write!(f, "malformed function call in a constraint: {w}"),
121 FuncError::TooDeep { func } => write!(
122 f,
123 "function `{func}` called from a constraint nests calls deeper than {MAX_DEPTH}"
124 ),
125 FuncError::StepLimit { func } => write!(
126 f,
127 "function `{func}` called from a constraint ran more than {MAX_STEPS} statements \
128 (a runaway loop?)"
129 ),
130 }
131 }
132}
133
134impl std::error::Error for FuncError {}
135
136impl From<ConstraintEvalError> for FuncError {
137 fn from(e: ConstraintEvalError) -> Self {
138 FuncError::Eval(e)
139 }
140}
141
142#[derive(Clone, Copy, PartialEq, Eq)]
144enum Flow {
145 Normal,
146 Break,
147 Continue,
148 Return,
149}
150
151struct Machine<'a> {
152 funcs: &'a [Func],
153 asg: Assignment,
155 steps: u64,
156}
157
158impl Machine<'_> {
159 fn func(&self, id: FuncId) -> Result<&Func, FuncError> {
160 self.funcs.get(id as usize).ok_or(FuncError::Malformed(
161 "a call of a function outside the table",
162 ))
163 }
164
165 fn resolve(&mut self, e: &Expr, depth: usize) -> Result<Expr, FuncError> {
167 if !has_call(e) {
168 return Ok(e.clone());
169 }
170 match e {
171 Expr::Array(ArrayNode::Call { func, args, ty }) => {
172 let mut vals = Vec::with_capacity(args.len());
173 for a in args {
174 vals.push(self.eval_in(a, depth)?);
175 }
176 let r = self.run(*func, vals, depth + 1)?;
177 Ok(Expr::Const(*ty, r))
178 }
179 other => other.try_map_children(&mut |c| self.resolve(c, depth)),
180 }
181 }
182
183 fn eval_in(&mut self, e: &Expr, depth: usize) -> Result<BigUint, FuncError> {
184 let r = self.resolve(e, depth)?;
185 Ok(eval(&r, &self.asg)?)
186 }
187
188 fn run(&mut self, id: FuncId, args: Vec<BigUint>, depth: usize) -> Result<BigUint, FuncError> {
189 let (n_params, ret, slots, name) = {
190 let f = self.func(id)?;
191 (f.n_params, f.ret, f.slots.clone(), f.name.clone())
192 };
193 if depth >= MAX_DEPTH {
194 return Err(FuncError::TooDeep { func: name });
195 }
196 if args.len() != n_params {
197 return Err(FuncError::Malformed(
198 "a call with the wrong number of arguments",
199 ));
200 }
201 let mut frame: Vec<BigUint> = vec![BigUint::zero(); slots.len()];
202 for (k, a) in args.into_iter().enumerate() {
203 frame[k] = a & mask(slots[k].width);
204 }
205 let saved = std::mem::replace(&mut self.asg.vars, frame);
206 let body = self.func(id)?.body.clone();
207 let flow = self.exec(&body, id, depth);
208 let frame = std::mem::replace(&mut self.asg.vars, saved);
209 flow?;
210 Ok(frame[ret as usize].clone() & mask(slots[ret as usize].width))
211 }
212
213 fn tick(&mut self, id: FuncId) -> Result<(), FuncError> {
214 self.steps += 1;
215 if self.steps > MAX_STEPS {
216 return Err(FuncError::StepLimit {
217 func: self.func(id)?.name.clone(),
218 });
219 }
220 Ok(())
221 }
222
223 fn exec(&mut self, body: &[Stmt], id: FuncId, depth: usize) -> Result<Flow, FuncError> {
224 for s in body {
225 self.tick(id)?;
226 match s {
227 Stmt::Assign(slot, e) => {
228 let v = self.eval_in(e, depth)?;
229 let w = self
230 .func(id)?
231 .slots
232 .get(*slot as usize)
233 .ok_or(FuncError::Malformed(
234 "an assignment to a slot outside the frame",
235 ))?
236 .width;
237 self.asg.vars[*slot as usize] = v & mask(w);
238 }
239 Stmt::If(c, t, e) => {
240 let taken = !self.eval_in(c, depth)?.is_zero();
241 let flow = self.exec(if taken { t } else { e }, id, depth)?;
242 if flow != Flow::Normal {
243 return Ok(flow);
244 }
245 }
246 Stmt::Loop {
247 cond,
248 body,
249 step,
250 test_first,
251 } => {
252 let mut first = true;
253 loop {
254 self.tick(id)?;
255 if (*test_first || !first)
256 && let Some(c) = cond
257 && self.eval_in(c, depth)?.is_zero()
258 {
259 break;
260 }
261 first = false;
262 match self.exec(body, id, depth)? {
263 Flow::Break => break,
264 Flow::Return => return Ok(Flow::Return),
265 Flow::Normal | Flow::Continue => {}
266 }
267 if self.exec(step, id, depth)? == Flow::Return {
268 return Ok(Flow::Return);
269 }
270 }
271 }
272 Stmt::Break => return Ok(Flow::Break),
273 Stmt::Continue => return Ok(Flow::Continue),
274 Stmt::Return => return Ok(Flow::Return),
275 }
276 }
277 Ok(Flow::Normal)
278 }
279}
280
281pub fn has_call(e: &Expr) -> bool {
283 e.any_node(&mut |n| matches!(n, Expr::Array(ArrayNode::Call { .. })))
284}
285
286pub fn call(
290 funcs: &[Func],
291 id: FuncId,
292 args: &[BigUint],
293 states: &[Option<BigUint>],
294) -> Result<BigUint, FuncError> {
295 let mut m = Machine {
296 funcs,
297 asg: Assignment {
298 vars: Vec::new(),
299 states: states.to_vec(),
300 },
301 steps: 0,
302 };
303 m.run(id, args.to_vec(), 0)
304}
305
306fn is_ground(e: &Expr) -> bool {
309 !e.any_node(&mut |n| matches!(n, Expr::Var(..) | Expr::Array(_)))
310}
311
312pub fn fold_ground_calls(
317 e: &Expr,
318 funcs: &[Func],
319 states: &[Option<BigUint>],
320) -> Result<Expr, FuncError> {
321 if !has_call(e) {
322 return Ok(e.clone());
323 }
324 match e {
325 Expr::Array(ArrayNode::Call { func, args, ty }) => {
326 let args = args
327 .iter()
328 .map(|a| fold_ground_calls(a, funcs, states))
329 .collect::<Result<Vec<_>, _>>()?;
330 if args.iter().all(is_ground) {
331 let asg = Assignment {
332 vars: Vec::new(),
333 states: states.to_vec(),
334 };
335 let vals = args
336 .iter()
337 .map(|a| eval(a, &asg))
338 .collect::<Result<Vec<_>, _>>()?;
339 let r = call(funcs, *func, &vals, states)?;
340 Ok(Expr::Const(*ty, r))
341 } else {
342 Ok(Expr::Array(ArrayNode::Call {
343 func: *func,
344 args,
345 ty: *ty,
346 }))
347 }
348 }
349 Expr::Array(ArrayNode::Pick {
353 index,
354 low,
355 elems,
356 ty,
357 }) => Ok(Expr::Array(ArrayNode::Pick {
358 index: Box::new(fold_ground_calls(index, funcs, states)?),
359 low: *low,
360 elems: elems.clone(),
361 ty: *ty,
362 })),
363 other => other.try_map_children(&mut |c| fold_ground_calls(c, funcs, states)),
364 }
365}
366
367#[cfg(test)]
368pub(crate) mod tests {
369 use super::*;
370 use crate::constraints::ir::{BinOp, UnOp};
371
372 fn u(w: u32) -> Ty {
373 Ty::unsigned(w)
374 }
375 fn s32() -> Ty {
376 Ty::signed(32)
377 }
378 fn var(i: u32, t: Ty) -> Expr {
379 Expr::Var(i, t)
380 }
381 fn k(t: Ty, v: u64) -> Expr {
382 Expr::constant(t, v)
383 }
384 fn bin(op: BinOp, a: Expr, b: Expr, t: Ty) -> Expr {
385 Expr::Binary(op, Box::new(a), Box::new(b), t)
386 }
387 fn n(v: u64) -> BigUint {
388 BigUint::from(v)
389 }
390
391 pub(crate) fn count_ones() -> Func {
394 let w = var(0, u(10));
396 let c = var(1, s32());
397 Func {
398 name: "count_ones".into(),
399 slots: vec![u(10), s32()],
400 n_params: 1,
401 ret: 1,
402 body: vec![
403 Stmt::Assign(1, k(s32(), 0)),
404 Stmt::Loop {
405 cond: Some(bin(BinOp::Ne, w.clone(), k(u(10), 0), Ty::BIT)),
406 body: vec![Stmt::Assign(
407 1,
408 bin(
409 BinOp::Add,
410 c.clone(),
411 Expr::Resize(
412 Box::new(bin(BinOp::And, w.clone(), k(u(10), 1), u(10))),
413 s32(),
414 ),
415 s32(),
416 ),
417 )],
418 step: vec![Stmt::Assign(
419 0,
420 bin(BinOp::Shr, w.clone(), k(u(32), 1), u(10)),
421 )],
422 test_first: true,
423 },
424 ],
425 }
426 }
427
428 #[test]
429 fn count_ones_counts_the_set_bits() {
430 let f = [count_ones()];
431 for v in [0u64, 1, 0b1011, 0x3ff, 0b10_0000_0001] {
432 let r = call(&f, 0, &[n(v)], &[]).unwrap();
433 assert_eq!(r, n(u64::from(v.count_ones())), "count_ones({v:#b})");
434 }
435 }
436
437 #[test]
438 fn a_runaway_loop_is_a_loud_error_not_a_hang() {
439 let f = [Func {
440 name: "spin".into(),
441 slots: vec![s32()],
442 n_params: 0,
443 ret: 0,
444 body: vec![Stmt::Loop {
445 cond: None,
446 body: Vec::new(),
447 step: Vec::new(),
448 test_first: true,
449 }],
450 }];
451 assert!(matches!(
452 call(&f, 0, &[], &[]),
453 Err(FuncError::StepLimit { .. })
454 ));
455 }
456
457 #[test]
458 fn recursion_works_and_unbounded_recursion_is_an_error() {
459 let rec = Expr::Array(ArrayNode::Call {
461 func: 0,
462 args: vec![bin(BinOp::Sub, var(0, s32()), k(s32(), 1), s32())],
463 ty: s32(),
464 });
465 let fact = Func {
466 name: "fact".into(),
467 slots: vec![s32(), s32()],
468 n_params: 1,
469 ret: 1,
470 body: vec![Stmt::If(
471 bin(BinOp::Le, var(0, s32()), k(s32(), 1), Ty::BIT),
472 vec![Stmt::Assign(1, k(s32(), 1)), Stmt::Return],
473 vec![Stmt::Assign(1, bin(BinOp::Mul, var(0, s32()), rec, s32()))],
474 )],
475 };
476 assert_eq!(
477 call(std::slice::from_ref(&fact), 0, &[n(5)], &[]).unwrap(),
478 n(120)
479 );
480 let endless = Func {
481 body: vec![Stmt::Assign(
482 1,
483 Expr::Array(ArrayNode::Call {
484 func: 0,
485 args: vec![var(0, s32())],
486 ty: s32(),
487 }),
488 )],
489 ..fact
490 };
491 assert!(matches!(
492 call(&[endless], 0, &[n(1)], &[]),
493 Err(FuncError::TooDeep { .. })
494 ));
495 }
496
497 #[test]
498 fn break_continue_and_do_while_follow_the_loop_rules() {
499 let i = var(0, s32());
502 let r = var(1, s32());
503 let inc = Stmt::Assign(0, bin(BinOp::Add, i.clone(), k(s32(), 1), s32()));
504 let f = Func {
505 name: "dw".into(),
506 slots: vec![s32(), s32()],
507 n_params: 0,
508 ret: 1,
509 body: vec![Stmt::Loop {
510 cond: Some(bin(BinOp::Lt, i.clone(), k(s32(), 100), Ty::BIT)),
511 body: vec![
512 inc,
513 Stmt::If(
514 bin(BinOp::Eq, i.clone(), k(s32(), 2), Ty::BIT),
515 vec![Stmt::Continue],
516 Vec::new(),
517 ),
518 Stmt::If(
519 bin(BinOp::Eq, i.clone(), k(s32(), 5), Ty::BIT),
520 vec![Stmt::Break],
521 Vec::new(),
522 ),
523 Stmt::Assign(1, bin(BinOp::Add, r.clone(), i.clone(), s32())),
524 ],
525 step: Vec::new(),
526 test_first: false,
527 }],
528 };
529 assert_eq!(call(&[f], 0, &[], &[]).unwrap(), n(8));
531 }
532
533 #[test]
534 fn a_function_reads_state_and_an_unreadable_state_is_an_error() {
535 let f = Func {
537 name: "f".into(),
538 slots: vec![s32(), s32()],
539 n_params: 1,
540 ret: 1,
541 body: vec![Stmt::Assign(
542 1,
543 bin(BinOp::Add, var(0, s32()), Expr::State(0, s32()), s32()),
544 )],
545 };
546 assert_eq!(
547 call(std::slice::from_ref(&f), 0, &[n(2)], &[Some(n(40))]).unwrap(),
548 n(42)
549 );
550 assert_eq!(
551 call(&[f], 0, &[n(2)], &[None]),
552 Err(FuncError::Eval(ConstraintEvalError::StateUnknown(0)))
553 );
554 }
555
556 #[test]
557 fn division_by_zero_inside_a_function_is_an_error() {
558 let f = Func {
559 name: "d".into(),
560 slots: vec![s32(), s32()],
561 n_params: 1,
562 ret: 1,
563 body: vec![Stmt::Assign(
564 1,
565 bin(BinOp::Div, k(s32(), 1), var(0, s32()), s32()),
566 )],
567 };
568 assert_eq!(
569 call(&[f], 0, &[n(0)], &[]),
570 Err(FuncError::Eval(ConstraintEvalError::DivideByZero))
571 );
572 }
573
574 #[test]
575 fn only_calls_with_fixed_arguments_are_folded() {
576 let f = [count_ones()];
577 let ground = Expr::Array(ArrayNode::Call {
578 func: 0,
579 args: vec![k(u(10), 0b111)],
580 ty: s32(),
581 });
582 let open = Expr::Array(ArrayNode::Call {
583 func: 0,
584 args: vec![var(0, u(10))],
585 ty: s32(),
586 });
587 let e = bin(BinOp::Add, ground, open.clone(), s32());
588 let folded = fold_ground_calls(&e, &f, &[]).unwrap();
589 assert_eq!(folded, bin(BinOp::Add, k(s32(), 3), open, s32()));
590 let neg = Expr::Unary(
592 UnOp::Neg,
593 Box::new(Expr::Array(ArrayNode::Call {
594 func: 0,
595 args: vec![k(u(10), 1)],
596 ty: s32(),
597 })),
598 s32(),
599 );
600 assert_eq!(
601 fold_ground_calls(&neg, &f, &[]).unwrap(),
602 Expr::Unary(UnOp::Neg, Box::new(k(s32(), 1)), s32())
603 );
604 }
605}