1use super::dist::{DefaultWeight, Dist, Item, SELECTOR_BASE, Span};
39use super::func::{self, Func, FuncError};
40use super::guard::{self, GuardValue};
41use super::ir::{
42 ArrayId, ArrayNode, Assignment, BinOp, ConstraintEvalError, Expr, LoopId, ReduceOp, Ty, UnOp,
43 VarId, eval, mask,
44};
45use super::solve::{SolveError, Solver};
46use num_bigint::BigUint;
47use num_traits::{One, Zero};
48use std::collections::HashMap;
49
50pub const SIZE_TIERS: [u64; 2] = [1024, 65536];
52
53pub const MAX_SIZE_RETRIES: usize = 256;
56
57#[derive(Clone, Debug)]
59pub struct ArrayState {
60 pub elem_ty: Ty,
62 pub active: bool,
65 pub cur: Vec<Option<BigUint>>,
67 pub domain: Option<Vec<BigUint>>,
69 pub rows: Option<RowsState>,
73 pub template_of: Option<ArrayId>,
76}
77
78#[derive(Clone, Debug)]
81pub struct RowsState {
82 pub template: ArrayId,
84 pub cur_rows: Vec<Vec<Option<BigUint>>>,
86 pub row_base: ArrayId,
89}
90
91impl ArrayState {
92 pub fn flat(elem_ty: Ty, active: bool, cur: Vec<Option<BigUint>>) -> ArrayState {
94 ArrayState {
95 elem_ty,
96 active,
97 cur,
98 domain: None,
99 rows: None,
100 template_of: None,
101 }
102 }
103
104 pub fn with_domain(mut self, domain: Option<Vec<BigUint>>) -> ArrayState {
106 self.domain = domain;
107 self
108 }
109}
110
111#[derive(Clone, Debug, PartialEq, Eq)]
113pub enum ExpandError {
114 IndexOutOfRange {
116 array: ArrayId,
118 index: i64,
120 size: usize,
122 },
123 LeafIndexOutOfRange {
126 index: i64,
128 size: usize,
130 },
131 NonConstantIndex,
134 GuardError,
136 StateUnknown(u32),
138 Func(FuncError),
140 DistRandomBound,
144 Malformed(&'static str),
146}
147
148#[derive(Clone, Debug, Default)]
150pub struct Expanded {
151 pub conjuncts: Vec<Expr>,
154 pub elem_vars: Vec<(ArrayId, usize)>,
156 pub elem_tys: Vec<Ty>,
158 pub extra_states: Vec<Option<BigUint>>,
161 pub extra_dists: Vec<Dist>,
166}
167
168struct Expander<'a> {
169 arrays: &'a [ArrayState],
170 sizes: &'a [usize],
171 elem_base: Vec<u32>,
173 asg: Assignment,
174 synth: HashMap<(ArrayId, usize), u32>,
175 env: Vec<(LoopId, Expr)>,
176 funcs: &'a [Func],
177 dists: &'a [Dist],
179 extra_dists: Vec<Dist>,
180 instanced: Vec<InstancedDist>,
185 rebind: Vec<(ArrayId, ArrayId)>,
188}
189
190type InstancedDist = (usize, Vec<(LoopId, Expr)>, usize);
192
193fn konst(t: Ty, v: u64) -> Expr {
194 Expr::constant(t, v)
195}
196
197fn has_vars(e: &Expr) -> bool {
198 let mut any = false;
199 e.for_each_var(&mut |_, _| any = true);
200 any
201}
202
203impl Expander<'_> {
204 fn array(&self, a: ArrayId) -> Result<&ArrayState, ExpandError> {
205 self.arrays.get(a as usize).ok_or(ExpandError::Malformed(
206 "an array id outside the program's table",
207 ))
208 }
209
210 fn res(&self, a: ArrayId) -> ArrayId {
212 self.rebind
213 .iter()
214 .rev()
215 .find(|(t, _)| *t == a)
216 .map_or(a, |(_, r)| *r)
217 }
218
219 fn size_of(&self, a: ArrayId) -> Result<usize, ExpandError> {
220 let a = self.res(a);
221 self.sizes
222 .get(a as usize)
223 .copied()
224 .ok_or(ExpandError::Malformed("no size for an array"))
225 }
226
227 fn const_value(&self, e: &Expr) -> Result<BigUint, ExpandError> {
229 if has_vars(e) {
230 return Err(ExpandError::NonConstantIndex);
231 }
232 eval(e, &self.asg).map_err(|err| match err {
233 ConstraintEvalError::StateUnknown(s) => ExpandError::StateUnknown(s),
234 _ => ExpandError::NonConstantIndex,
235 })
236 }
237
238 fn index_value(&mut self, index: &Expr) -> Result<i64, ExpandError> {
241 let idx = self.expand(index)?;
242 let v = self.const_value(&idx)?;
243 let t = idx.ty();
244 Ok(if t.signed && v.bit(u64::from(t.width) - 1) {
245 -i64::try_from((BigUint::one() << t.width) - &v).unwrap_or(i64::MAX)
246 } else {
247 i64::try_from(&v).unwrap_or(i64::MAX)
248 })
249 }
250
251 fn element(&mut self, a: ArrayId, idx: usize, ty: Ty) -> Result<Expr, ExpandError> {
254 let a = self.res(a);
255 let arr = self.array(a)?;
256 if arr.elem_ty != ty {
257 return Err(ExpandError::Malformed("an element read at the wrong type"));
258 }
259 let (active, cur) = (arr.active, arr.cur.get(idx).cloned());
260 if active {
261 let id = self.elem_base[a as usize] + u32::try_from(idx).expect("an element index");
262 return Ok(Expr::Var(id, ty));
263 }
264 let Some(cur) = cur else {
265 return Err(ExpandError::Malformed("an element of a stale array view"));
266 };
267 if let Some(&id) = self.synth.get(&(a, idx)) {
268 return Ok(Expr::State(id, ty));
269 }
270 let id = u32::try_from(self.asg.states.len()).expect("a state id");
271 self.asg.states.push(cur);
272 self.synth.insert((a, idx), id);
273 Ok(Expr::State(id, ty))
274 }
275
276 fn lookup(&self, id: LoopId) -> Result<Expr, ExpandError> {
277 self.env
278 .iter()
279 .rev()
280 .find(|(l, _)| *l == id)
281 .map(|(_, e)| e.clone())
282 .ok_or(ExpandError::Malformed("a loop variable outside its loop"))
283 }
284
285 fn expand(&mut self, e: &Expr) -> Result<Expr, ExpandError> {
286 match e {
287 Expr::Guard(p, then, otherwise, t) => {
288 let p = self.expand(p)?;
289 match guard::value(&p, &self.asg) {
290 GuardValue::True => self.expand(then),
291 GuardValue::False => self.expand(otherwise),
292 GuardValue::Error => Err(ExpandError::GuardError),
293 GuardValue::Random => Ok(Expr::Guard(
294 Box::new(p),
295 Box::new(self.expand(then)?),
296 Box::new(self.expand(otherwise)?),
297 *t,
298 )),
299 }
300 }
301 Expr::Ite(c, x, y, t) => {
302 let c = self.expand(c)?;
303 if !has_vars(&c)
304 && let Ok(v) = eval(&c, &self.asg)
305 {
306 return if v.is_zero() {
307 self.expand(y)
308 } else {
309 self.expand(x)
310 };
311 }
312 Ok(Expr::Ite(
313 Box::new(c),
314 Box::new(self.expand(x)?),
315 Box::new(self.expand(y)?),
316 *t,
317 ))
318 }
319 Expr::Binary(op @ (BinOp::LogAnd | BinOp::LogOr | BinOp::LogImpl), l, r, t) => {
320 let l = self.expand(l)?;
321 if !has_vars(&l)
322 && let Ok(v) = eval(&l, &self.asg)
323 {
324 let l_true = !v.is_zero();
325 let decided = match op {
326 BinOp::LogAnd | BinOp::LogImpl => !l_true,
327 _ => l_true,
328 };
329 if decided {
330 return Ok(konst(*t, u64::from(*op != BinOp::LogAnd)));
331 }
332 }
333 Ok(Expr::Binary(
334 *op,
335 Box::new(l),
336 Box::new(self.expand(r)?),
337 *t,
338 ))
339 }
340 Expr::Array(n) => self.array_node(n),
341 Expr::Var(v, t) if *v >= SELECTOR_BASE => {
342 let j = (*v - SELECTOR_BASE) as usize;
343 match self.dists.get(j) {
344 Some(d) if d.subject.is_some() => {
345 let k = match self
346 .instanced
347 .iter()
348 .find(|(tj, env, _)| *tj == j && *env == self.env)
349 {
350 Some(&(_, _, k)) => k,
351 None => {
352 let inst = self.instantiate(d)?;
353 self.extra_dists.push(inst);
354 let k = self.extra_dists.len() - 1;
355 self.instanced.push((j, self.env.clone(), k));
356 k
357 }
358 };
359 Ok(Expr::Var(SELECTOR_BASE + (self.dists.len() + k) as u32, *t))
360 }
361 _ => Ok(e.clone()),
362 }
363 }
364 other => other.try_map_children(&mut |c| self.expand(c)),
365 }
366 }
367
368 fn instantiate(&mut self, d: &Dist) -> Result<Dist, ExpandError> {
373 let Some(subject) = &d.subject else {
374 return Err(ExpandError::Malformed("an instance of a plain dist"));
375 };
376 let subject = self.expand(subject)?;
377 let mut vars: Vec<VarId> = Vec::new();
378 subject.for_each_var(&mut |v, _| vars.push(v));
379 vars.sort_unstable();
380 vars.dedup();
381 let plain_var = match &subject {
382 Expr::Var(v, _) => Some(*v),
383 _ => None,
384 };
385 let fixed = |this: &mut Self, e: &Expr| -> Result<Expr, ExpandError> {
386 let x = this.expand(e)?;
387 if has_vars(&x) {
388 return Err(ExpandError::DistRandomBound);
389 }
390 Ok(x)
391 };
392 let mut items = Vec::with_capacity(d.items.len());
393 for it in &d.items {
394 let span = match &it.span {
395 Span::Point(v) => Span::Point(fixed(self, v)?),
396 Span::Range { lo, hi } => Span::Range {
397 lo: lo.as_ref().map(|b| fixed(self, b)).transpose()?,
398 hi: hi.as_ref().map(|b| fixed(self, b)).transpose()?,
399 },
400 };
401 items.push(Item {
402 member: konst(Ty::BIT, 1),
403 ty: it.ty,
404 span,
405 weight: fixed(self, &it.weight)?,
406 per_range: it.per_range,
407 });
408 }
409 let default = match &d.default {
410 Some(w) => Some(DefaultWeight {
411 weight: fixed(self, &w.weight)?,
412 per_range: w.per_range,
413 }),
414 None => None,
415 };
416 Ok(Dist {
417 subject_ty: d.subject_ty,
418 plain_var,
419 vars,
420 items,
421 default,
422 subject: None,
423 })
424 }
425
426 fn row_array(&mut self, template: ArrayId, row: &Expr) -> Result<ArrayId, ExpandError> {
429 let outer = self
430 .array(template)?
431 .template_of
432 .ok_or(ExpandError::Malformed("a row scope over a non-template"))?;
433 let k = self.index_value(row)?;
434 let size = self.size_of(outer)?;
435 if k < 0 || k as usize >= size {
436 return Err(ExpandError::IndexOutOfRange {
437 array: outer,
438 index: k,
439 size,
440 });
441 }
442 let base = self
443 .array(outer)?
444 .rows
445 .as_ref()
446 .ok_or(ExpandError::Malformed("rows of an array without rows"))?
447 .row_base;
448 Ok(base + k as ArrayId)
449 }
450
451 fn array_node(&mut self, n: &ArrayNode) -> Result<Expr, ExpandError> {
452 match n {
453 ArrayNode::Call { func, args, ty } => {
454 let args = args
455 .iter()
456 .map(|a| self.expand(a))
457 .collect::<Result<Vec<_>, _>>()?;
458 let call = Expr::Array(ArrayNode::Call {
459 func: *func,
460 args,
461 ty: *ty,
462 });
463 func::fold_ground_calls(&call, self.funcs, &self.asg.states).map_err(|e| match e {
466 FuncError::Eval(ConstraintEvalError::StateUnknown(s)) => {
467 ExpandError::StateUnknown(s)
468 }
469 other => ExpandError::Func(other),
470 })
471 }
472 ArrayNode::Size(a) => Ok(konst(Ty::signed(32), self.size_of(*a)? as u64)),
473 ArrayNode::Loop(id, t) => {
474 let v = self.lookup(*id)?;
475 if v.ty() != *t {
476 return Err(ExpandError::Malformed(
477 "a loop variable bound at another type",
478 ));
479 }
480 Ok(v)
481 }
482 ArrayNode::Pick {
483 index,
484 low,
485 elems,
486 ty,
487 } => {
488 let i = self.index_value(index)?;
489 let pos = i
490 .checked_sub(*low)
491 .and_then(|p| usize::try_from(p).ok())
492 .filter(|&p| p < elems.len());
493 let Some(pos) = pos else {
494 return Err(ExpandError::LeafIndexOutOfRange {
495 index: i,
496 size: elems.len(),
497 });
498 };
499 let x = self.expand(&elems[pos])?;
500 if x.ty() != *ty {
501 return Err(ExpandError::Malformed("a pick element of another type"));
502 }
503 Ok(x)
504 }
505 ArrayNode::Elem { array, index, ty } => {
506 let i = self.index_value(index)?;
507 let size = self.size_of(*array)?;
508 if i < 0 || i as usize >= size {
509 return Err(ExpandError::IndexOutOfRange {
510 array: self.res(*array),
511 index: i,
512 size,
513 });
514 }
515 self.element(*array, i as usize, *ty)
516 }
517 ArrayNode::RowScope {
518 template,
519 row,
520 body,
521 } => {
522 let target = self.row_array(*template, row)?;
523 self.rebind.push((*template, target));
524 let r = self.expand(body);
525 self.rebind.pop();
526 r
527 }
528 ArrayNode::Distinct { scalars, arrays } => {
529 let mut members = Vec::new();
530 for s in scalars {
531 members.push(self.expand(s)?);
532 }
533 for a in arrays {
534 let size = self.size_of(*a)?;
535 let elem_ty = self.array(self.res(*a))?.elem_ty;
536 for i in 0..size {
537 members.push(self.element(*a, i, elem_ty)?);
538 }
539 }
540 if members.iter().any(|m| m.ty() != members[0].ty()) {
541 return Err(ExpandError::Malformed(
542 "the members of a unique group are of different types",
543 ));
544 }
545 Ok(if members.len() < 2 {
546 konst(Ty::BIT, 1)
547 } else {
548 Expr::AllDistinct(members)
549 })
550 }
551 ArrayNode::Foreach { array, var, body } => {
552 let size = self.size_of(*array)?;
553 let mut acc: Option<Expr> = None;
554 for i in 0..size {
555 self.env.push((*var, konst(Ty::signed(32), i as u64)));
556 let b = self.expand(body);
557 self.env.pop();
558 let b = b?;
559 acc = Some(match acc {
560 None => b,
561 Some(a) => Expr::Binary(BinOp::LogAnd, Box::new(a), Box::new(b), Ty::BIT),
562 });
563 }
564 Ok(acc.unwrap_or_else(|| konst(Ty::BIT, 1)))
565 }
566 ArrayNode::Reduce {
567 op,
568 array,
569 var,
570 index_var,
571 body,
572 } => {
573 let size = self.size_of(*array)?;
574 let elem_ty = self.array(self.res(*array))?.elem_ty;
575 let ty = body.ty();
576 let mut acc: Option<Expr> = None;
577 for i in 0..size {
578 let item = self.element(*array, i, elem_ty)?;
579 self.env.push((*var, item));
580 self.env.push((*index_var, konst(Ty::signed(32), i as u64)));
581 let b = self.expand(body);
582 self.env.pop();
583 self.env.pop();
584 let b = b?;
585 acc = Some(match acc {
586 None => b,
587 Some(a) => {
588 let bop = match op {
589 ReduceOp::Sum => BinOp::Add,
590 ReduceOp::Product => BinOp::Mul,
591 ReduceOp::And => BinOp::And,
592 ReduceOp::Or => BinOp::Or,
593 ReduceOp::Xor => BinOp::Xor,
594 };
595 Expr::Binary(bop, Box::new(a), Box::new(b), ty)
596 }
597 });
598 }
599 Ok(acc.unwrap_or_else(|| {
600 Expr::Const(
601 ty,
602 match op {
603 ReduceOp::Product => BigUint::one(),
604 ReduceOp::And => mask(ty.width),
605 ReduceOp::Sum | ReduceOp::Or | ReduceOp::Xor => BigUint::zero(),
606 },
607 )
608 }))
609 }
610 }
611 }
612}
613
614fn domain_constraint(v: Expr, domain: &[BigUint]) -> Expr {
616 let t = v.ty();
617 let mut acc: Option<Expr> = None;
618 for d in domain {
619 let eq = Expr::Binary(
620 BinOp::Eq,
621 Box::new(v.clone()),
622 Box::new(Expr::Const(t, d.clone() & mask(t.width))),
623 Ty::BIT,
624 );
625 acc = Some(match acc {
626 None => eq,
627 Some(a) => Expr::Binary(BinOp::LogOr, Box::new(a), Box::new(eq), Ty::BIT),
628 });
629 }
630 acc.unwrap_or_else(|| konst(Ty::BIT, 0))
631}
632
633pub fn expand(
639 conjuncts: &[Expr],
640 arrays: &[ArrayState],
641 sizes: &[usize],
642 n_vars0: u32,
643 states: &[Option<BigUint>],
644) -> Result<Expanded, ExpandError> {
645 expand_with(conjuncts, arrays, sizes, n_vars0, states, &[], &[])
646}
647
648pub fn expand_with(
653 conjuncts: &[Expr],
654 arrays: &[ArrayState],
655 sizes: &[usize],
656 n_vars0: u32,
657 states: &[Option<BigUint>],
658 funcs: &[Func],
659 dists: &[Dist],
660) -> Result<Expanded, ExpandError> {
661 if arrays.len() != sizes.len() {
662 return Err(ExpandError::Malformed("one size per array"));
663 }
664 let mut out = Expanded::default();
665 let mut elem_base = vec![0u32; arrays.len()];
666 let mut next = n_vars0;
667 for (a, arr) in arrays.iter().enumerate() {
668 if !arr.active {
669 continue;
670 }
671 elem_base[a] = next;
672 for i in 0..sizes[a] {
673 out.elem_vars.push((a as ArrayId, i));
674 out.elem_tys.push(arr.elem_ty);
675 next += 1;
676 if let Some(dom) = &arr.domain {
677 let id = elem_base[a] + i as u32;
678 out.conjuncts
679 .push(domain_constraint(Expr::Var(id, arr.elem_ty), dom));
680 }
681 }
682 }
683 let mut ex = Expander {
684 arrays,
685 sizes,
686 elem_base,
687 asg: Assignment {
688 vars: Vec::new(),
689 states: states.to_vec(),
690 },
691 synth: HashMap::new(),
692 env: Vec::new(),
693 funcs,
694 dists,
695 extra_dists: Vec::new(),
696 instanced: Vec::new(),
697 rebind: Vec::new(),
698 };
699 for c in conjuncts {
700 let e = ex.expand(c)?;
701 out.conjuncts.push(e);
702 }
703 out.extra_dists = std::mem::take(&mut ex.extra_dists);
704 out.extra_states = ex.asg.states.split_off(states.len());
705 Ok(out)
706}
707
708#[derive(Clone, Debug)]
710pub struct Lowered {
711 pub conjuncts: Vec<Expr>,
714 pub arrays: Vec<ArrayState>,
717 pub bases: Vec<(ArrayId, ArrayId)>,
719}
720
721struct Lowerer<'a> {
722 arrays: &'a [ArrayState],
723 sizes: &'a [usize],
724 states: &'a [Option<BigUint>],
725 bases: &'a [(ArrayId, ArrayId)],
726 env: Vec<(LoopId, i64)>,
727 rebind: Vec<(ArrayId, ArrayId)>,
728}
729
730impl Lowerer<'_> {
731 fn res(&self, a: ArrayId) -> ArrayId {
732 self.rebind
733 .iter()
734 .rev()
735 .find(|(t, _)| *t == a)
736 .map_or(a, |(_, r)| *r)
737 }
738
739 fn base_of(&self, outer: ArrayId) -> Option<ArrayId> {
740 self.bases
741 .iter()
742 .find(|(o, _)| *o == outer)
743 .map(|(_, b)| *b)
744 }
745
746 fn fixed(&self, row: &Expr) -> Option<i64> {
749 let row = replace_sizes(row, &mut |a| {
750 let a = self.res(a);
751 konst(
752 Ty::signed(32),
753 self.sizes.get(a as usize).copied().unwrap_or(0) as u64,
754 )
755 });
756 let asg = Assignment {
757 vars: Vec::new(),
758 states: self.states.to_vec(),
759 };
760 let v = eval(&row, &asg).ok()?;
761 let t = row.ty();
762 Some(if t.signed && v.bit(u64::from(t.width) - 1) {
763 -i64::try_from((BigUint::one() << t.width) - &v).ok()?
764 } else {
765 i64::try_from(&v).ok()?
766 })
767 }
768
769 fn low(&mut self, e: &Expr) -> Result<Expr, ExpandError> {
770 let Expr::Array(n) = e else {
771 return e.try_map_children(&mut |c| self.low(c));
772 };
773 match n {
774 ArrayNode::Loop(id, t) => Ok(match self.env.iter().rev().find(|(l, _)| l == id) {
775 Some(&(_, i)) => Expr::constant(*t, i as u64),
776 None => e.clone(),
777 }),
778 ArrayNode::Size(a) => Ok(Expr::Array(ArrayNode::Size(self.res(*a)))),
779 ArrayNode::Elem { array, index, ty } => Ok(Expr::Array(ArrayNode::Elem {
780 array: self.res(*array),
781 index: Box::new(self.low(index)?),
782 ty: *ty,
783 })),
784 ArrayNode::Distinct { scalars, arrays } => Ok(Expr::Array(ArrayNode::Distinct {
785 scalars: scalars
786 .iter()
787 .map(|x| self.low(x))
788 .collect::<Result<_, _>>()?,
789 arrays: arrays.iter().map(|a| self.res(*a)).collect(),
790 })),
791 ArrayNode::Reduce {
792 op,
793 array,
794 var,
795 index_var,
796 body,
797 } => Ok(Expr::Array(ArrayNode::Reduce {
798 op: *op,
799 array: self.res(*array),
800 var: *var,
801 index_var: *index_var,
802 body: Box::new(self.low(body)?),
803 })),
804 ArrayNode::Foreach { array, var, body } => {
805 let a = self.res(*array);
806 if self.base_of(a).is_none() {
807 return Ok(Expr::Array(ArrayNode::Foreach {
808 array: a,
809 var: *var,
810 body: Box::new(self.low(body)?),
811 }));
812 }
813 let mut acc: Option<Expr> = None;
814 for i in 0..self.sizes[a as usize] {
815 self.env.push((*var, i as i64));
816 let b = self.low(body);
817 self.env.pop();
818 let b = b?;
819 acc = Some(match acc {
820 None => b,
821 Some(x) => Expr::Binary(BinOp::LogAnd, Box::new(x), Box::new(b), Ty::BIT),
822 });
823 }
824 Ok(acc.unwrap_or_else(|| konst(Ty::BIT, 1)))
825 }
826 ArrayNode::RowScope {
827 template,
828 row,
829 body,
830 } => {
831 let row = self.low(row)?;
832 let outer = self
833 .arrays
834 .get(*template as usize)
835 .and_then(|t| t.template_of)
836 .ok_or(ExpandError::Malformed("a row scope over a non-template"))?;
837 let (Some(k), Some(base)) = (self.fixed(&row), self.base_of(outer)) else {
838 return Ok(Expr::Array(ArrayNode::RowScope {
840 template: *template,
841 row: Box::new(row),
842 body: Box::new(self.low(body)?),
843 }));
844 };
845 let size = self.sizes[outer as usize];
846 if k < 0 || k as usize >= size {
847 return Err(ExpandError::IndexOutOfRange {
848 array: outer,
849 index: k,
850 size,
851 });
852 }
853 self.rebind.push((*template, base + k as ArrayId));
854 let r = self.low(body);
855 self.rebind.pop();
856 r
857 }
858 ArrayNode::Call { .. } | ArrayNode::Pick { .. } => {
859 e.try_map_children(&mut |c| self.low(c))
860 }
861 }
862 }
863}
864
865pub fn lower_rows(
871 conjuncts: &[Expr],
872 arrays: &[ArrayState],
873 sizes: &[usize],
874 states: &[Option<BigUint>],
875) -> Result<Lowered, ExpandError> {
876 let mut out: Vec<ArrayState> = arrays.to_vec();
877 let mut bases = Vec::new();
878 let mut sizes2: Vec<usize> = sizes.to_vec();
879 for (a, arr) in arrays.iter().enumerate() {
880 let Some(rows) = &arr.rows else { continue };
881 let n = sizes[a];
882 let base = out.len() as ArrayId;
883 bases.push((a as ArrayId, base));
884 for k in 0..n {
885 out.push(ArrayState {
886 elem_ty: arr.elem_ty,
887 active: arr.active,
888 cur: rows.cur_rows.get(k).cloned().unwrap_or_default(),
889 domain: arr.domain.clone(),
890 rows: None,
891 template_of: None,
892 });
893 sizes2.push(rows.cur_rows.get(k).map_or(0, Vec::len));
894 }
895 out[a].active = false;
896 out[a].cur = vec![Some(BigUint::zero()); n];
897 if let Some(r) = out[a].rows.as_mut() {
898 r.row_base = base;
899 }
900 }
901 let mut lw = Lowerer {
902 arrays,
903 sizes: &sizes2,
904 states,
905 bases: &bases,
906 env: Vec::new(),
907 rebind: Vec::new(),
908 };
909 let lowered = conjuncts
910 .iter()
911 .map(|c| lw.low(c))
912 .collect::<Result<Vec<_>, _>>()?;
913 Ok(Lowered {
914 conjuncts: lowered,
915 arrays: out,
916 bases,
917 })
918}
919
920#[derive(Clone, Debug, PartialEq, Eq)]
922pub enum SizeError {
923 Unsatisfiable,
925 TooLarge,
927 GuardError,
929 Solve(SolveError),
931}
932
933#[derive(Clone, Debug, PartialEq, Eq)]
935pub struct SizeSolution {
936 pub sizes: Vec<usize>,
938 pub sized: Vec<ArrayId>,
940 pub tuple: Vec<usize>,
942}
943
944fn size_eligible(c: &Expr) -> bool {
947 !c.any_node(&mut |n| match n {
948 Expr::Array(
949 ArrayNode::Elem { .. }
950 | ArrayNode::Loop(..)
951 | ArrayNode::Foreach { .. }
952 | ArrayNode::Reduce { .. }
953 | ArrayNode::Distinct { .. }
954 | ArrayNode::Pick { .. }
956 | ArrayNode::RowScope { .. }
958 | ArrayNode::Call { .. },
961 ) => true,
962 Expr::Var(v, _) => *v >= super::dist::SELECTOR_BASE,
963 _ => false,
964 })
965}
966
967fn replace_sizes(e: &Expr, f: &mut impl FnMut(ArrayId) -> Expr) -> Expr {
968 match e {
969 Expr::Array(ArrayNode::Size(a)) => f(*a),
970 other => other
971 .try_map_children::<()>(&mut |c| Ok(replace_sizes(c, f)))
972 .expect("infallible"),
973 }
974}
975
976fn find(parent: &mut [usize], x: usize) -> usize {
977 let mut r = x;
978 while parent[r] != r {
979 r = parent[r];
980 }
981 let mut c = x;
982 while parent[c] != r {
983 let n = parent[c];
984 parent[c] = r;
985 c = n;
986 }
987 r
988}
989
990pub fn solve_sizes(
996 conjuncts: &[Expr],
997 arrays: &[ArrayState],
998 var_tys: &[Ty],
999 states: &[Option<BigUint>],
1000 blocked: &[Vec<usize>],
1001 seed: &mut u64,
1002) -> Result<SizeSolution, SizeError> {
1003 let current: Vec<usize> = arrays.iter().map(|a| a.cur.len()).collect();
1004 let n_fields = var_tys.len() as u32;
1005 let mut sized: Vec<ArrayId> = Vec::new();
1007 let mut eligible: Vec<&Expr> = Vec::new();
1008 for c in conjuncts {
1009 if !size_eligible(c) {
1010 continue;
1011 }
1012 c.any_node(&mut |n| {
1013 if let Expr::Array(ArrayNode::Size(a)) = n
1014 && arrays.get(*a as usize).is_some_and(|x| x.active)
1015 && !sized.contains(a)
1016 {
1017 sized.push(*a);
1018 }
1019 false
1020 });
1021 eligible.push(c);
1022 }
1023 if sized.is_empty() {
1024 return Ok(SizeSolution {
1025 sizes: current,
1026 sized,
1027 tuple: Vec::new(),
1028 });
1029 }
1030 sized.sort_unstable();
1031 let size_var = |a: ArrayId| -> u32 {
1032 n_fields + u32::try_from(sized.iter().position(|x| *x == a).expect("sized")).unwrap()
1033 };
1034 let total = n_fields as usize + sized.len();
1036 let subbed: Vec<Expr> = eligible
1037 .iter()
1038 .map(|c| {
1039 replace_sizes(c, &mut |a| {
1040 if arrays[a as usize].active {
1041 Expr::Var(size_var(a), Ty::signed(32))
1042 } else {
1043 konst(Ty::signed(32), current[a as usize] as u64)
1044 }
1045 })
1046 })
1047 .collect();
1048 let mut parent: Vec<usize> = (0..total).collect();
1049 let vars_of: Vec<Vec<VarId>> = subbed
1050 .iter()
1051 .map(|c| {
1052 let mut v = Vec::new();
1053 c.for_each_var(&mut |x, _| v.push(x));
1054 v
1055 })
1056 .collect();
1057 for vs in &vars_of {
1058 for w in vs.windows(2) {
1059 let (a, b) = (
1060 find(&mut parent, w[0] as usize),
1061 find(&mut parent, w[1] as usize),
1062 );
1063 parent[a] = b;
1064 }
1065 }
1066 let size_roots: Vec<usize> = (0..sized.len())
1067 .map(|k| find(&mut parent, n_fields as usize + k))
1068 .collect();
1069 let mut keep: Vec<Expr> = Vec::new();
1070 let mut used = vec![false; total];
1071 for k in 0..sized.len() {
1072 used[n_fields as usize + k] = true;
1073 }
1074 for (c, vs) in subbed.iter().zip(&vars_of) {
1075 let Some(&v0) = vs.first() else { continue };
1076 let root = find(&mut parent, v0 as usize);
1077 if size_roots.contains(&root) {
1078 for &v in vs {
1079 used[v as usize] = true;
1080 }
1081 keep.push(c.clone());
1082 }
1083 }
1084 let mut new_id = vec![0u32; total];
1086 let mut tys: Vec<Ty> = Vec::new();
1087 for (v, u) in used.iter().enumerate() {
1088 if *u {
1089 new_id[v] = tys.len() as u32;
1090 tys.push(if v < n_fields as usize {
1091 var_tys[v]
1092 } else {
1093 Ty::signed(32)
1094 });
1095 }
1096 }
1097 let mut keep: Vec<Expr> = keep
1098 .iter()
1099 .map(|c| c.substitute_vars(&mut |v, t| Expr::Var(new_id[v as usize], t)))
1100 .collect();
1101 let svar = |k: usize| Expr::Var(new_id[n_fields as usize + k], Ty::signed(32));
1102 for tup in blocked {
1104 let mut same: Option<Expr> = None;
1105 for (k, &sz) in tup.iter().enumerate() {
1106 let eq = Expr::Binary(
1107 BinOp::Eq,
1108 Box::new(svar(k)),
1109 Box::new(konst(Ty::signed(32), sz as u64)),
1110 Ty::BIT,
1111 );
1112 same = Some(match same {
1113 None => eq,
1114 Some(a) => Expr::Binary(BinOp::LogAnd, Box::new(a), Box::new(eq), Ty::BIT),
1115 });
1116 }
1117 if let Some(s) = same {
1118 keep.push(Expr::Unary(UnOp::LogNot, Box::new(s), Ty::BIT));
1119 }
1120 }
1121 let assignment = Assignment {
1122 vars: Vec::new(),
1123 states: states.to_vec(),
1124 };
1125 let order: Vec<VarId> = (0..tys.len() as u32).collect();
1126 let tiers: [u64; 3] = [SIZE_TIERS[0], SIZE_TIERS[1], i32::MAX as u64];
1127 for (ti, bound) in tiers.iter().enumerate() {
1128 let mut cs = keep.clone();
1129 for k in 0..sized.len() {
1130 cs.push(Expr::Binary(
1131 BinOp::Ge,
1132 Box::new(svar(k)),
1133 Box::new(konst(Ty::signed(32), 0)),
1134 Ty::BIT,
1135 ));
1136 cs.push(Expr::Binary(
1137 BinOp::Le,
1138 Box::new(svar(k)),
1139 Box::new(konst(Ty::signed(32), *bound)),
1140 Ty::BIT,
1141 ));
1142 }
1143 let mut resolved = Vec::with_capacity(cs.len());
1144 for c in &cs {
1145 resolved.push(guard::resolve(c, &assignment).map_err(|_| SizeError::GuardError)?);
1146 }
1147 let solver = match Solver::new(&tys, states, &resolved) {
1148 Ok(s) => s,
1149 Err(SolveError::Unsatisfiable) => continue,
1150 Err(e) => return Err(SizeError::Solve(e)),
1151 };
1152 let values = match solver.sample(&order, seed) {
1153 Ok(v) => v,
1154 Err(SolveError::Unsatisfiable) => continue,
1155 Err(e) => return Err(SizeError::Solve(e)),
1156 };
1157 if ti == 2 {
1158 return Err(SizeError::TooLarge);
1159 }
1160 let mut sizes = current;
1161 let mut tuple = Vec::with_capacity(sized.len());
1162 for (k, &a) in sized.iter().enumerate() {
1163 let v = &values[new_id[n_fields as usize + k] as usize];
1164 let n = usize::try_from(v).expect("a bounded size");
1165 sizes[a as usize] = n;
1166 tuple.push(n);
1167 }
1168 return Ok(SizeSolution {
1169 sizes,
1170 sized,
1171 tuple,
1172 });
1173 }
1174 Err(SizeError::Unsatisfiable)
1175}
1176
1177#[cfg(test)]
1178mod tests {
1179 use super::*;
1180
1181 fn s32() -> Ty {
1182 Ty::signed(32)
1183 }
1184 fn u(w: u32) -> Ty {
1185 Ty::unsigned(w)
1186 }
1187 fn size(a: ArrayId) -> Expr {
1188 Expr::Array(ArrayNode::Size(a))
1189 }
1190 fn elem(a: ArrayId, i: Expr, t: Ty) -> Expr {
1191 Expr::Array(ArrayNode::Elem {
1192 array: a,
1193 index: Box::new(i),
1194 ty: t,
1195 })
1196 }
1197 fn lp(id: LoopId) -> Expr {
1198 Expr::Array(ArrayNode::Loop(id, s32()))
1199 }
1200 fn bin(op: BinOp, a: Expr, b: Expr, t: Ty) -> Expr {
1201 Expr::Binary(op, Box::new(a), Box::new(b), t)
1202 }
1203 fn arr(active: bool, cur: &[u64]) -> ArrayState {
1204 ArrayState {
1205 elem_ty: s32(),
1206 active,
1207 cur: cur.iter().map(|v| Some(BigUint::from(*v))).collect(),
1208 domain: None,
1209 rows: None,
1210 template_of: None,
1211 }
1212 }
1213 fn ev(e: &Expr, vars: &[u64], states: &[Option<BigUint>]) -> bool {
1214 let a = Assignment {
1215 vars: vars.iter().map(|v| BigUint::from(*v)).collect(),
1216 states: states.to_vec(),
1217 };
1218 !eval(e, &a).expect("evaluates").is_zero()
1219 }
1220
1221 #[test]
1224 fn foreach_unrolls_over_the_size() {
1225 let body = bin(
1226 BinOp::Gt,
1227 elem(0, lp(0), s32()),
1228 bin(BinOp::Mul, konst(s32(), 2), lp(0), s32()),
1229 Ty::BIT,
1230 );
1231 let c = Expr::Array(ArrayNode::Foreach {
1232 array: 0,
1233 var: 0,
1234 body: Box::new(body),
1235 });
1236 let x = expand(&[c], &[arr(true, &[])], &[3], 0, &[]).expect("expands");
1237 assert_eq!(x.elem_vars, vec![(0, 0), (0, 1), (0, 2)]);
1238 assert_eq!(x.conjuncts.len(), 1);
1239 assert!(ev(&x.conjuncts[0], &[1, 3, 5], &[]));
1240 assert!(!ev(&x.conjuncts[0], &[1, 2, 5], &[]));
1241 assert!(!ev(&x.conjuncts[0], &[0, 3, 5], &[]));
1242 }
1243
1244 #[test]
1247 fn a_constant_guard_keeps_the_last_neighbour_unread() {
1248 let pred = bin(
1249 BinOp::Lt,
1250 lp(0),
1251 bin(BinOp::Sub, size(0), konst(s32(), 1), s32()),
1252 Ty::BIT,
1253 );
1254 let then = bin(
1255 BinOp::Lt,
1256 elem(0, lp(0), s32()),
1257 elem(0, bin(BinOp::Add, lp(0), konst(s32(), 1), s32()), s32()),
1258 Ty::BIT,
1259 );
1260 let body = Expr::Guard(
1261 Box::new(pred),
1262 Box::new(then),
1263 Box::new(konst(Ty::BIT, 1)),
1264 Ty::BIT,
1265 );
1266 let c = Expr::Array(ArrayNode::Foreach {
1267 array: 0,
1268 var: 0,
1269 body: Box::new(body),
1270 });
1271 let x = expand(&[c], &[arr(true, &[])], &[4], 0, &[]).expect("expands");
1272 assert_eq!(x.elem_vars.len(), 4);
1273 assert!(ev(&x.conjuncts[0], &[1, 2, 3, 4], &[]));
1274 assert!(!ev(&x.conjuncts[0], &[1, 3, 2, 4], &[]));
1275 }
1276
1277 #[test]
1278 fn an_unguarded_out_of_range_index_is_a_loud_error() {
1279 let c = bin(
1280 BinOp::Eq,
1281 elem(0, konst(s32(), 5), s32()),
1282 konst(s32(), 0),
1283 Ty::BIT,
1284 );
1285 assert_eq!(
1286 expand(&[c], &[arr(true, &[])], &[2], 0, &[]).unwrap_err(),
1287 ExpandError::IndexOutOfRange {
1288 array: 0,
1289 index: 5,
1290 size: 2
1291 }
1292 );
1293 }
1294
1295 #[test]
1298 fn a_pick_selects_the_element_its_index_names() {
1299 let pick = |idx: Expr| {
1300 Expr::Array(ArrayNode::Pick {
1301 index: Box::new(idx),
1302 low: 2,
1303 elems: vec![Expr::Var(0, u(8)), Expr::Var(1, u(8)), Expr::Var(2, u(8))],
1304 ty: u(8),
1305 })
1306 };
1307 let states = [Some(BigUint::from(3u32))];
1308 let c = bin(
1309 BinOp::Eq,
1310 pick(Expr::State(0, s32())),
1311 konst(u(8), 9),
1312 Ty::BIT,
1313 );
1314 let x = expand(&[c], &[], &[], 3, &states).expect("expands");
1315 assert_eq!(x.conjuncts.len(), 1);
1316 assert!(ev(&x.conjuncts[0], &[0, 9, 0], &states));
1317 assert!(!ev(&x.conjuncts[0], &[9, 0, 9], &states));
1318 }
1319
1320 #[test]
1324 fn an_out_of_range_pick_is_an_error_only_when_live() {
1325 let pick = Expr::Array(ArrayNode::Pick {
1326 index: Box::new(Expr::State(0, s32())),
1327 low: 0,
1328 elems: vec![Expr::Var(0, u(8)), Expr::Var(1, u(8))],
1329 ty: u(8),
1330 });
1331 let body = bin(BinOp::Eq, pick, konst(u(8), 1), Ty::BIT);
1332 let guarded = Expr::Guard(
1333 Box::new(bin(
1334 BinOp::Lt,
1335 Expr::State(0, s32()),
1336 konst(s32(), 2),
1337 Ty::BIT,
1338 )),
1339 Box::new(body.clone()),
1340 Box::new(konst(Ty::BIT, 1)),
1341 Ty::BIT,
1342 );
1343 let out = [Some(BigUint::from(7u32))];
1344 let x = expand(&[guarded], &[], &[], 2, &out).expect("the guard drops the pick");
1345 assert!(ev(&x.conjuncts[0], &[], &out));
1346 assert_eq!(
1347 expand(std::slice::from_ref(&body), &[], &[], 2, &out).unwrap_err(),
1348 ExpandError::LeafIndexOutOfRange { index: 7, size: 2 }
1349 );
1350 let neg = [Some((BigUint::from(1u32) << 32) - 1u32)];
1352 assert_eq!(
1353 expand(&[body], &[], &[], 2, &neg).unwrap_err(),
1354 ExpandError::LeafIndexOutOfRange { index: -1, size: 2 }
1355 );
1356 }
1357
1358 #[test]
1361 fn an_inactive_array_is_state() {
1362 let c = bin(
1363 BinOp::Eq,
1364 elem(0, konst(s32(), 1), s32()),
1365 konst(s32(), 7),
1366 Ty::BIT,
1367 );
1368 let x = expand(&[c], &[arr(false, &[3, 7])], &[2], 0, &[]).expect("expands");
1369 assert!(x.elem_vars.is_empty());
1370 assert_eq!(x.extra_states, vec![Some(BigUint::from(7u32))]);
1371 assert!(ev(&x.conjuncts[0], &[], &x.extra_states));
1372 }
1373
1374 #[test]
1375 fn reductions_fold_and_empty_ones_are_the_identity() {
1376 let sum = |op| {
1377 Expr::Array(ArrayNode::Reduce {
1378 op,
1379 array: 0,
1380 var: 0,
1381 index_var: 1,
1382 body: Box::new(Expr::Array(ArrayNode::Loop(0, s32()))),
1383 })
1384 };
1385 let c = bin(BinOp::Eq, sum(ReduceOp::Sum), konst(s32(), 6), Ty::BIT);
1386 let x = expand(&[c], &[arr(true, &[])], &[3], 0, &[]).expect("expands");
1387 assert!(ev(&x.conjuncts[0], &[1, 2, 3], &[]));
1388 assert!(!ev(&x.conjuncts[0], &[1, 2, 4], &[]));
1389 for (op, want) in [
1390 (ReduceOp::Sum, 0u64),
1391 (ReduceOp::Product, 1),
1392 (ReduceOp::Or, 0),
1393 (ReduceOp::Xor, 0),
1394 (ReduceOp::And, 0xFFFF_FFFF),
1395 ] {
1396 let c = bin(BinOp::Eq, sum(op), konst(s32(), want), Ty::BIT);
1397 let x = expand(&[c], &[arr(true, &[])], &[0], 0, &[]).expect("expands");
1398 assert!(ev(&x.conjuncts[0], &[], &[]), "{op:?}");
1399 }
1400 }
1401
1402 #[test]
1404 fn a_size_constraint_resizes_the_array() {
1405 let c = bin(
1406 BinOp::Eq,
1407 size(0),
1408 Expr::Resize(Box::new(Expr::Var(0, u(8))), s32()),
1409 Ty::BIT,
1410 );
1411 let bound = bin(BinOp::Lt, Expr::Var(0, u(8)), konst(u(8), 5), Ty::BIT);
1412 let mut seen = std::collections::BTreeSet::new();
1413 let mut seed = 1u64;
1414 for _ in 0..200 {
1415 let s = solve_sizes(
1416 &[c.clone(), bound.clone()],
1417 &[arr(true, &[9, 9])],
1418 &[u(8)],
1419 &[],
1420 &[],
1421 &mut seed,
1422 )
1423 .expect("solves");
1424 assert_eq!(s.sized, vec![0]);
1425 assert!(s.sizes[0] < 5);
1426 seen.insert(s.sizes[0]);
1427 }
1428 assert_eq!(seen.len(), 5, "every size 0..5 is reachable: {seen:?}");
1429 }
1430
1431 #[test]
1432 fn an_unconstrained_size_keeps_the_array() {
1433 let mut seed = 1u64;
1434 let s = solve_sizes(&[], &[arr(true, &[1, 2, 3])], &[], &[], &[], &mut seed).unwrap();
1435 assert_eq!(s.sizes, vec![3]);
1436 assert!(s.sized.is_empty());
1437 }
1438
1439 #[test]
1440 fn an_impossible_size_is_unsatisfiable_and_a_huge_one_is_too_large() {
1441 let mut seed = 1u64;
1442 let lt = |n: u64| bin(BinOp::Lt, size(0), konst(s32(), n), Ty::BIT);
1443 let gt = |n: u64| bin(BinOp::Gt, size(0), konst(s32(), n), Ty::BIT);
1444 let a = [arr(true, &[])];
1445 assert_eq!(
1446 solve_sizes(&[lt(3), gt(5)], &a, &[], &[], &[], &mut seed).unwrap_err(),
1447 SizeError::Unsatisfiable
1448 );
1449 assert_eq!(
1450 solve_sizes(&[gt(10_000_000)], &a, &[], &[], &[], &mut seed).unwrap_err(),
1451 SizeError::TooLarge
1452 );
1453 let s = solve_sizes(&[gt(2000)], &a, &[], &[], &[], &mut seed).unwrap();
1455 assert!(s.sizes[0] > 2000);
1456 }
1457
1458 #[test]
1459 fn blocked_tuples_are_excluded() {
1460 let c = bin(BinOp::Le, size(0), konst(s32(), 1), Ty::BIT);
1461 let a = [arr(true, &[])];
1462 let mut seed = 3u64;
1463 let s = solve_sizes(
1464 std::slice::from_ref(&c),
1465 &a,
1466 &[],
1467 &[],
1468 &[vec![0]],
1469 &mut seed,
1470 )
1471 .unwrap();
1472 assert_eq!(s.sizes, vec![1]);
1473 assert_eq!(
1474 solve_sizes(&[c], &a, &[], &[], &[vec![0], vec![1]], &mut seed).unwrap_err(),
1475 SizeError::Unsatisfiable
1476 );
1477 }
1478
1479 #[test]
1482 fn the_size_problem_holds_only_connected_size_constraints() {
1483 let elem_c = bin(
1484 BinOp::Eq,
1485 elem(0, konst(s32(), 0), s32()),
1486 konst(s32(), 1),
1487 Ty::BIT,
1488 );
1489 let unrelated = bin(BinOp::Eq, Expr::Var(0, u(8)), konst(u(8), 200), Ty::BIT);
1490 let c = bin(BinOp::Eq, size(0), konst(s32(), 2), Ty::BIT);
1491 let mut seed = 1u64;
1492 let s = solve_sizes(
1493 &[elem_c, unrelated, c],
1494 &[arr(true, &[])],
1495 &[u(8)],
1496 &[],
1497 &[],
1498 &mut seed,
1499 )
1500 .unwrap();
1501 assert_eq!(s.sizes, vec![2]);
1502 }
1503}