1#[cfg(test)]
69use std::sync::Arc;
70
71use crate::ast::Value;
72use crate::iteration::comprehension::ast::Comprehension;
73use crate::iteration::comprehension::cardinality::{Interval, ProductMeasure};
74use crate::iteration::comprehension::eval_source::{EvalContext, SourceEval};
75use crate::iteration::comprehension::measure::AxisMeasure;
76use crate::iteration::comprehension::metadata::{IndexFn, cycle_length};
77use crate::iteration::comprehension::predicate::CompiledPredicate;
78use crate::iteration::comprehension::source::Source;
79use crate::iteration::comprehension::strategies::{Selection, ranked_filter, shape_input};
80use crate::iteration::comprehension::strategy::StrategyName;
81#[cfg(test)]
82use crate::kernel::PolydatKernel;
83use crate::kernel::interp::Lookup;
84
85pub type RuntimeTuple = Vec<(String, Value)>;
94
95struct EvaluatedNode {
105 tuples: Vec<RuntimeTuple>,
106 index_fn: Option<IndexFn>,
107}
108
109#[derive(Debug, Clone)]
111pub enum RuntimeError {
112 SourceEval {
115 var: String,
117 source: String,
119 message: String,
121 },
122 FilterEval {
124 predicate: String,
126 message: String,
128 },
129 OrderEval {
131 strategy: StrategyName,
133 message: String,
135 },
136 StrategyRejectsInput {
139 strategy: StrategyName,
141 index_fn: Option<IndexFn>,
143 },
144 UnsupportedShape(String),
149 ZipLengthMismatch {
153 lengths: Vec<u64>,
155 },
156}
157
158impl std::fmt::Display for RuntimeError {
159 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160 match self {
161 RuntimeError::SourceEval {
162 var,
163 source,
164 message,
165 } => {
166 write!(f, "for_each clause '{var} in {source}': {message}")
167 }
168 RuntimeError::FilterEval { predicate, message } => {
169 write!(f, "comprehension filter '{predicate}': {message}")
170 }
171 RuntimeError::OrderEval { strategy, message } => {
172 write!(f, "order strategy {strategy:?}: {message}")
173 }
174 RuntimeError::StrategyRejectsInput { strategy, index_fn } => write!(
175 f,
176 "order strategy {strategy:?} rejects input shape {index_fn:?} \
177 (V4: per-strategy IndexFn contract; see comprehension_forms.md §3.6's \
178 strategy table)"
179 ),
180 RuntimeError::UnsupportedShape(msg) => write!(f, "{msg}"),
181 RuntimeError::ZipLengthMismatch { lengths } => {
182 write!(f, "zip strict: child lengths differ ({lengths:?})")
183 }
184 }
185 }
186}
187
188impl std::error::Error for RuntimeError {}
189
190pub fn evaluate_for_iteration(
199 comp: &Comprehension,
200 scope: &dyn Lookup,
201) -> Result<Vec<RuntimeTuple>, RuntimeError> {
202 evaluate_indexed(comp, scope).map(|t| t.to_vec())
203}
204
205pub fn evaluate_for_iteration_reported(
220 comp: &Comprehension,
221 scope: &dyn Lookup,
222) -> Result<EvaluatedIteration, RuntimeError> {
223 let mut state = EvalState::new(comp, scope);
224 let (node, _) = state.index_node(comp, &[])?;
225 Ok(EvaluatedIteration {
226 tuples: IndexedTuples { node }.to_vec(),
227 clauses: state.yields,
228 })
229}
230
231pub fn evaluate_indexed(
251 comp: &Comprehension,
252 scope: &dyn Lookup,
253) -> Result<IndexedTuples, RuntimeError> {
254 let mut state = EvalState::new(comp, scope);
255 let (node, _) = state.index_node(comp, &[])?;
256 Ok(IndexedTuples { node })
257}
258
259pub fn evaluate_for_iteration_materialized(
267 comp: &Comprehension,
268 scope: &dyn Lookup,
269) -> Result<EvaluatedIteration, RuntimeError> {
270 let mut state = EvalState::new(comp, scope);
271 let tuples = state.evaluate_node(comp, &[])?.tuples;
272 Ok(EvaluatedIteration {
273 tuples,
274 clauses: state.yields,
275 })
276}
277
278#[derive(Debug, Clone)]
281pub struct IndexedTuples {
282 node: Indexed,
283}
284
285impl IndexedTuples {
286 pub fn len(&self) -> u64 {
288 self.node.len()
289 }
290
291 pub fn is_empty(&self) -> bool {
293 self.len() == 0
294 }
295
296 pub fn get(&self, i: u64) -> Option<RuntimeTuple> {
298 if i >= self.len() {
299 return None;
300 }
301 let mut out = RuntimeTuple::new();
302 self.node.append_at(i, &mut out);
303 Some(out)
304 }
305
306 pub fn iter(&self) -> impl Iterator<Item = RuntimeTuple> + '_ {
308 (0..self.len()).filter_map(|i| self.get(i))
309 }
310
311 pub fn to_vec(&self) -> Vec<RuntimeTuple> {
313 self.iter().collect()
314 }
315}
316
317#[derive(Debug, Clone)]
320enum Indexed {
321 Tuples(Vec<RuntimeTuple>),
324 Clause { name: String, values: ClauseValues },
326 Product {
330 children: Vec<Indexed>,
331 lens: Vec<u64>,
332 len: u64,
333 },
334 Lockstep { children: Vec<Indexed>, len: u64 },
336 Cycle { children: Vec<Indexed>, len: u64 },
340 Concat { children: Vec<Indexed>, len: u64 },
342 Select {
344 child: Box<Indexed>,
345 selection: Selection,
346 },
347}
348
349#[derive(Debug, Clone)]
351enum ClauseValues {
352 Range { lo: i64, step: i64, len: u64 },
354 List(Vec<Value>),
356}
357
358impl ClauseValues {
359 fn len(&self) -> u64 {
360 match self {
361 ClauseValues::Range { len, .. } => *len,
362 ClauseValues::List(values) => values.len() as u64,
363 }
364 }
365
366 fn at(&self, i: u64) -> Value {
367 match self {
368 ClauseValues::Range { lo, step, .. } => {
371 Value::U64((i128::from(*lo) + i128::from(i) * i128::from(*step)) as i64 as u64)
372 }
373 ClauseValues::List(values) => values[i as usize].clone(),
374 }
375 }
376}
377
378impl Indexed {
379 fn len(&self) -> u64 {
380 match self {
381 Indexed::Tuples(tuples) => tuples.len() as u64,
382 Indexed::Clause { values, .. } => values.len(),
383 Indexed::Product { len, .. }
384 | Indexed::Lockstep { len, .. }
385 | Indexed::Cycle { len, .. }
386 | Indexed::Concat { len, .. } => *len,
387 Indexed::Select { selection, .. } => selection.len(),
388 }
389 }
390
391 fn append_at(&self, i: u64, out: &mut RuntimeTuple) {
394 match self {
395 Indexed::Tuples(tuples) => out.extend(tuples[i as usize].iter().cloned()),
396 Indexed::Clause { name, values } => out.push((name.clone(), values.at(i))),
397 Indexed::Product { children, lens, .. } => {
398 let mut digits = vec![0u64; lens.len()];
399 let mut rest = i;
400 for (d, len) in digits.iter_mut().zip(lens).rev() {
401 *d = rest % len;
402 rest /= len;
403 }
404 for (child, d) in children.iter().zip(digits) {
405 child.append_at(d, out);
406 }
407 }
408 Indexed::Lockstep { children, .. } => {
409 for child in children {
410 child.append_at(i, out);
411 }
412 }
413 Indexed::Cycle { children, .. } => {
414 for child in children {
415 child.append_at(i % child.len(), out);
416 }
417 }
418 Indexed::Concat { children, .. } => {
419 let mut offset = i;
420 for child in children {
421 let len = child.len();
422 if offset < len {
423 child.append_at(offset, out);
424 return;
425 }
426 offset -= len;
427 }
428 }
429 Indexed::Select { child, selection } => {
430 if let Some(p) = selection.get(i) {
431 child.append_at(p, out);
432 }
433 }
434 }
435 }
436}
437
438#[derive(Debug, Clone, PartialEq, Eq)]
454pub struct ClauseYield {
455 pub var: String,
457 pub source: Option<String>,
459 pub evaluations: usize,
461 pub values: usize,
463}
464
465#[derive(Debug, Clone)]
468pub struct EvaluatedIteration {
469 pub tuples: Vec<RuntimeTuple>,
471 pub clauses: Vec<ClauseYield>,
473}
474
475struct EvalState<'a> {
476 scope: &'a dyn Lookup,
479 yields: Vec<ClauseYield>,
483 by_leaf: std::collections::HashMap<usize, usize>,
488 mult: usize,
493}
494
495impl<'a> EvalState<'a> {
496 fn new(comp: &Comprehension, scope: &'a dyn Lookup) -> Self {
498 let mut state = EvalState {
499 scope,
500 yields: Vec::new(),
501 by_leaf: std::collections::HashMap::new(),
502 mult: 1,
503 };
504 state.enumerate_leaves(comp);
505 state
506 }
507
508 fn enumerate_leaves(&mut self, node: &Comprehension) {
509 match node {
510 Comprehension::Clause { name, source } => {
511 self.by_leaf
512 .insert(std::ptr::from_ref(source) as usize, self.yields.len());
513 self.yields.push(ClauseYield {
514 var: name.clone(),
515 source: source.to_text(),
516 evaluations: 0,
517 values: 0,
518 });
519 }
520 Comprehension::Cartesian { children }
521 | Comprehension::Zip { children, .. }
522 | Comprehension::Union { children } => {
523 for child in children {
524 self.enumerate_leaves(child);
525 }
526 }
527 Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
528 self.enumerate_leaves(child);
529 }
530 }
531 }
532
533 fn record_yield(&mut self, source: &Source, values: usize) {
536 if let Some(&i) = self.by_leaf.get(&(std::ptr::from_ref(source) as usize)) {
537 self.yields[i].evaluations = self.yields[i].evaluations.saturating_add(self.mult);
538 self.yields[i].values = self.yields[i]
539 .values
540 .saturating_add(values.saturating_mul(self.mult));
541 }
542 }
543}
544
545impl EvalState<'_> {
546 fn evaluate_node(
547 &mut self,
548 node: &Comprehension,
549 prefix: &[(String, Value)],
550 ) -> Result<EvaluatedNode, RuntimeError> {
551 match node {
552 Comprehension::Clause { name, source } => self.evaluate_clause(name, source, prefix),
553 Comprehension::Cartesian { children } => self.evaluate_cartesian(children, prefix),
554 Comprehension::Zip { children, mode } => self.evaluate_zip(children, *mode, prefix),
555 Comprehension::Union { children } => self.evaluate_union(children, prefix),
556 Comprehension::Filter { child, predicate } => {
557 let inner = self.evaluate_node(child, prefix)?;
558 self.apply_filter(inner, predicate)
559 }
560 Comprehension::Order {
561 child,
562 strategy,
563 truncation,
564 seed,
565 } => {
566 let child = shape_input(child, *strategy);
569 if has_continuous_axis(child) {
574 let sampled =
575 self.sample_space(child, prefix, *strategy, *truncation, *seed)?;
576 return Ok(EvaluatedNode {
577 index_fn: Some(selected(sampled.tuples.len() as u64)),
578 tuples: sampled.tuples,
579 });
580 }
581 if let Some((child, predicate)) = ranked_filter(child, *strategy) {
584 let inner = self.evaluate_node(child, prefix)?;
585 let predicate = CompiledPredicate::new(predicate);
586 let mut survivors = Vec::new();
587 for (p, tuple) in inner.tuples.iter().enumerate() {
588 if predicate.keeps(tuple, self.scope)? {
589 survivors.push(p as u64);
590 }
591 }
592 let selection = surviving_selection(
593 *strategy,
594 inner.index_fn.as_ref(),
595 inner.tuples.len() as u64,
596 *truncation,
597 *seed,
598 &survivors,
599 )?;
600 return Ok(EvaluatedNode {
601 tuples: selection
602 .iter()
603 .map(|p| inner.tuples[p as usize].clone())
604 .collect(),
605 index_fn: Some(selected(selection.len())),
606 });
607 }
608 let inner = self.evaluate_node(child, prefix)?;
609 self.apply_order(inner, *strategy, *truncation, *seed)
610 }
611 }
612 }
613
614 fn evaluate_source(
617 &mut self,
618 name: &str,
619 source: &Source,
620 prefix: &[(String, Value)],
621 ) -> Result<crate::iteration::comprehension::eval_source::EvaluatedSource, RuntimeError> {
622 let ctx = EvalContext {
623 var_name: name,
624 scope: self.scope,
625 prefix,
626 };
627 let evaluated = source.evaluate(Some(&ctx)).map_err(|e| match e {
628 crate::iteration::comprehension::eval_source::EvalError::EvalFailed {
629 var,
630 source,
631 message,
632 } => RuntimeError::SourceEval {
633 var,
634 source,
635 message,
636 },
637 crate::iteration::comprehension::eval_source::EvalError::NeedsContext => {
638 RuntimeError::UnsupportedShape(format!(
639 "clause '{name}': source requires kernel context but evaluator \
640 provided none — internal bug in runtime walker"
641 ))
642 }
643 })?;
644 self.record_yield(source, evaluated.values.len());
645 Ok(evaluated)
646 }
647
648 fn evaluate_clause(
649 &mut self,
650 name: &str,
651 source: &Source,
652 prefix: &[(String, Value)],
653 ) -> Result<EvaluatedNode, RuntimeError> {
654 let evaluated = self.evaluate_source(name, source, prefix)?;
655
656 if evaluated.values.is_empty() {
657 return Ok(EvaluatedNode {
663 tuples: Vec::new(),
664 index_fn: Some(evaluated.index_fn),
665 });
666 }
667 let tuples: Vec<RuntimeTuple> = evaluated
668 .values
669 .into_iter()
670 .map(|v| vec![(name.to_string(), v)])
671 .collect();
672 Ok(EvaluatedNode {
673 tuples,
674 index_fn: Some(evaluated.index_fn),
675 })
676 }
677
678 fn evaluate_cartesian(
679 &mut self,
680 children: &[Comprehension],
681 prefix: &[(String, Value)],
682 ) -> Result<EvaluatedNode, RuntimeError> {
683 if children.is_empty() {
684 return Ok(EvaluatedNode {
685 tuples: vec![Vec::new()],
686 index_fn: Some(IndexFn::Lattice {
687 axis_sizes: vec![1],
688 }),
689 });
690 }
691 let mut child_index_fns: Vec<(Option<IndexFn>, u64)> = Vec::with_capacity(children.len());
692 let mut dependent_observed = false;
693 let result_tuples = self.evaluate_cartesian_rec(
694 children.len(),
695 children,
696 prefix,
697 &mut child_index_fns,
698 &mut dependent_observed,
699 )?;
700
701 let combined = if dependent_observed {
707 None
708 } else {
709 let index_fns: Vec<Option<IndexFn>> =
710 child_index_fns.into_iter().map(|(idx, _)| idx).collect();
711 combine_cartesian_index_fn(&index_fns)
712 };
713 Ok(EvaluatedNode {
714 tuples: result_tuples,
715 index_fn: combined,
716 })
717 }
718
719 fn evaluate_cartesian_rec(
720 &mut self,
721 child_count: usize,
722 children: &[Comprehension],
723 prefix: &[(String, Value)],
724 child_index_fns: &mut Vec<(Option<IndexFn>, u64)>,
725 dependent_observed: &mut bool,
726 ) -> Result<Vec<RuntimeTuple>, RuntimeError> {
727 if children.is_empty() {
728 return Ok(vec![Vec::new()]);
729 }
730 let (head, tail) = children.split_first().unwrap();
731 let head_eval = self.evaluate_node(head, prefix)?;
732 let head_axis_len = head_eval.tuples.len() as u64;
733 let depth = child_count - children.len();
738 match child_index_fns.get(depth) {
739 None => child_index_fns.push((head_eval.index_fn.clone(), head_axis_len)),
740 Some((_, first)) if *first != head_axis_len => *dependent_observed = true,
741 Some(_) => {}
742 }
743
744 if tail.is_empty() {
745 return Ok(head_eval.tuples);
746 }
747 let mut out = Vec::new();
748 for head_tuple in head_eval.tuples {
749 let mut extended_prefix: Vec<(String, Value)> = prefix.to_vec();
750 extended_prefix.extend(head_tuple.iter().cloned());
751 let tail_tuples = self.evaluate_cartesian_rec(
752 child_count,
753 tail,
754 &extended_prefix,
755 child_index_fns,
756 dependent_observed,
757 )?;
758 for tail_tuple in tail_tuples {
759 let mut merged = head_tuple.clone();
760 merged.extend(tail_tuple);
761 out.push(merged);
762 }
763 }
764 Ok(out)
765 }
766
767 fn evaluate_zip(
768 &mut self,
769 children: &[Comprehension],
770 mode: crate::iteration::comprehension::strategy::ZipMode,
771 prefix: &[(String, Value)],
772 ) -> Result<EvaluatedNode, RuntimeError> {
773 use crate::iteration::comprehension::strategy::ZipMode;
774 if children.is_empty() {
775 return Ok(EvaluatedNode {
776 tuples: vec![Vec::new()],
777 index_fn: Some(IndexFn::Lockstep { length: 1 }),
778 });
779 }
780 let per_child: Vec<EvaluatedNode> = children
781 .iter()
782 .map(|c| self.evaluate_node(c, prefix))
783 .collect::<Result<_, _>>()?;
784 let lengths: Vec<usize> = per_child.iter().map(|n| n.tuples.len()).collect();
785 let iter_count = match mode {
786 ZipMode::Strict => {
787 let first = lengths.first().copied().unwrap_or(0);
788 if lengths.iter().any(|&n| n != first) {
789 return Err(RuntimeError::ZipLengthMismatch {
790 lengths: lengths.iter().map(|&n| n as u64).collect(),
791 });
792 }
793 first
794 }
795 ZipMode::Truncate => lengths.iter().copied().min().unwrap_or(0),
796 ZipMode::Cycle => {
797 let counts: Vec<u64> = lengths.iter().map(|&n| n as u64).collect();
798 cycle_length(&counts) as usize
799 }
800 };
801 let mut tuples = Vec::with_capacity(iter_count);
802 for i in 0..iter_count {
803 let mut bindings: RuntimeTuple = Vec::new();
804 for (child, &len) in per_child.iter().zip(lengths.iter()) {
805 let idx = match mode {
806 ZipMode::Cycle => i % len,
807 _ => i,
808 };
809 bindings.extend(child.tuples[idx].iter().cloned());
810 }
811 tuples.push(bindings);
812 }
813 let index_fn = match mode {
814 ZipMode::Strict | ZipMode::Truncate => Some(IndexFn::Lockstep {
815 length: iter_count as u64,
816 }),
817 ZipMode::Cycle => Some(IndexFn::Modular {
818 axis_sizes: lengths.iter().map(|n| *n as u64).collect(),
819 }),
820 };
821 Ok(EvaluatedNode { tuples, index_fn })
822 }
823
824 fn evaluate_union(
825 &mut self,
826 children: &[Comprehension],
827 prefix: &[(String, Value)],
828 ) -> Result<EvaluatedNode, RuntimeError> {
829 let mut tuples = Vec::new();
830 let mut segment_sizes = Vec::with_capacity(children.len());
831 let mut all_segments_addressable = true;
832 for child in children {
833 let sub = self.evaluate_node(child, prefix)?;
834 segment_sizes.push(sub.tuples.len() as u64);
835 if sub.index_fn.is_none() {
836 all_segments_addressable = false;
837 }
838 tuples.extend(sub.tuples);
839 }
840 let index_fn = if all_segments_addressable {
841 Some(IndexFn::Concatenation { segment_sizes })
842 } else {
843 None
844 };
845 Ok(EvaluatedNode { tuples, index_fn })
846 }
847
848 fn apply_filter(
849 &mut self,
850 input: EvaluatedNode,
851 predicate: &str,
852 ) -> Result<EvaluatedNode, RuntimeError> {
853 let predicate = CompiledPredicate::new(predicate);
854 let mut out = Vec::with_capacity(input.tuples.len());
855 for tuple in input.tuples {
856 if predicate.keeps(&tuple, self.scope)? {
857 out.push(tuple);
858 }
859 }
860 Ok(EvaluatedNode {
863 tuples: out,
864 index_fn: None,
865 })
866 }
867
868 fn sample_space(
881 &mut self,
882 child: &Comprehension,
883 prefix: &[(String, Value)],
884 strategy: StrategyName,
885 truncation: Option<u64>,
886 seed: Option<u64>,
887 ) -> Result<EvaluatedNode, RuntimeError> {
888 let mut space = SampleSpace::default();
889 self.collect_sample_space(child, prefix, &mut space, &mut Vec::new())?;
890 let discrete_axes: Vec<u64> = space
891 .axes
892 .iter()
893 .filter_map(|a| match a {
894 SampleAxis::Discrete(tuples) => Some(tuples.len() as u64),
895 SampleAxis::Continuous { .. } => None,
896 })
897 .collect();
898 if discrete_axes.contains(&0) {
899 return Ok(EvaluatedNode {
900 tuples: Vec::new(),
901 index_fn: None,
902 });
903 }
904 let (intervals, measures): (Vec<Interval>, Vec<ProductMeasure>) = space
905 .axes
906 .iter()
907 .filter_map(|a| match a {
908 SampleAxis::Continuous {
909 interval, measure, ..
910 } => Some((
911 interval.clone(),
912 match measure {
913 AxisMeasure::Uniform => ProductMeasure::Uniform,
914 AxisMeasure::Named { name, .. } => ProductMeasure::Named(*name),
915 },
916 )),
917 SampleAxis::Discrete(_) => None,
918 })
919 .unzip();
920 let sequence = !matches!(strategy, StrategyName::Extrema);
921 let index_fn = if !sequence {
926 IndexFn::Lattice {
927 axis_sizes: space
928 .axes
929 .iter()
930 .map(|a| match a {
931 SampleAxis::Discrete(tuples) => tuples.len() as u64,
932 SampleAxis::Continuous { .. } => 2,
933 })
934 .collect(),
935 }
936 } else if discrete_axes.is_empty() {
937 IndexFn::Continuous {
938 intervals,
939 measure: ProductMeasure::Product(measures),
940 }
941 } else {
942 IndexFn::Hybrid {
943 discrete_axes,
944 continuous_axes: intervals,
945 measure: ProductMeasure::Product(measures),
946 }
947 };
948
949 let mut want = truncation;
950 let mut rounds = 0;
951 loop {
952 let multi_indices = draw_sample(&index_fn, strategy, want, seed)?;
953 let drawn = multi_indices.len() as u64;
954 let tuples = multi_indices
955 .iter()
956 .map(|mi| space.realize(mi, strategy))
957 .collect();
958 let mut node = EvaluatedNode {
959 tuples,
960 index_fn: None,
961 };
962 for predicate in &space.predicates {
963 node = self.apply_filter(node, predicate)?;
964 }
965 let (Some(n), Some(asked)) = (truncation, want) else {
966 return Ok(node);
967 };
968 let enough = node.tuples.len() as u64 >= n;
969 let exhausted = drawn < asked;
970 if !sequence {
971 return Ok(node);
972 }
973 if enough || exhausted || rounds >= SAMPLE_ROUNDS {
974 node.tuples.truncate(n as usize);
975 return Ok(node);
976 }
977
978 want = Some(asked.saturating_mul(2));
979 rounds += 1;
980 }
981 }
982
983 fn collect_sample_space(
991 &mut self,
992 c: &Comprehension,
993 prefix: &[(String, Value)],
994 space: &mut SampleSpace,
995 bound: &mut Vec<String>,
996 ) -> Result<(), RuntimeError> {
997 let measure_error = |name: &str, message: String| RuntimeError::SourceEval {
998 var: name.to_string(),
999 source: "<continuous>".to_string(),
1000 message,
1001 };
1002 match c {
1003 Comprehension::Clause {
1004 name,
1005 source: Source::ContinuousInterval { interval, measure },
1006 } => {
1007 let measure =
1008 AxisMeasure::from_product(measure, 0).map_err(|m| measure_error(name, m))?;
1009 space.axes.push(SampleAxis::Continuous {
1010 name: name.clone(),
1011 interval: interval.clone(),
1012 measure,
1013 });
1014 bound.push(name.clone());
1015 }
1016 Comprehension::Clause {
1017 name,
1018 source:
1019 Source::Distribution {
1020 distribution,
1021 support,
1022 params,
1023 },
1024 } => {
1025 let measure = AxisMeasure::named(*distribution, params)
1026 .map_err(|m| measure_error(name, m))?;
1027 space.axes.push(SampleAxis::Continuous {
1028 name: name.clone(),
1029 interval: support.clone(),
1030 measure,
1031 });
1032 bound.push(name.clone());
1033 }
1034 Comprehension::Clause { name, source } => {
1035 let references = c.referenced_source_names();
1036 if let Some(dep) = bound.iter().find(|b| references.contains(*b)) {
1037 return Err(RuntimeError::UnsupportedShape(format!(
1038 "clause '{name}' references '{dep}' beside a continuous axis; \
1039 a sampled cartesian is independent (comprehension_forms.md §5, V4)"
1040 )));
1041 }
1042 let node = self.evaluate_clause(name, source, prefix)?;
1043 space.axes.push(SampleAxis::Discrete(node.tuples));
1044 bound.push(name.clone());
1045 }
1046 Comprehension::Cartesian { children } => {
1047 for child in children {
1048 self.collect_sample_space(child, prefix, space, bound)?;
1049 }
1050 }
1051 Comprehension::Filter { child, predicate } => {
1052 self.collect_sample_space(child, prefix, space, bound)?;
1053 space.predicates.push(predicate.clone());
1054 }
1055 Comprehension::Zip { .. }
1056 | Comprehension::Union { .. }
1057 | Comprehension::Order { .. } => {
1058 let node = self.evaluate_node(c, prefix)?;
1059 bound.extend(c.coordinate_names());
1060 space.axes.push(SampleAxis::Discrete(node.tuples));
1061 }
1062 }
1063 Ok(())
1064 }
1065
1066 fn apply_order(
1067 &mut self,
1068 input: EvaluatedNode,
1069 strategy: StrategyName,
1070 truncation: Option<u64>,
1071 seed: Option<u64>,
1072 ) -> Result<EvaluatedNode, RuntimeError> {
1073 let selection = order_selection(
1077 strategy,
1078 input.index_fn.as_ref(),
1079 input.tuples.len() as u64,
1080 truncation,
1081 seed,
1082 )?;
1083 let out = selection
1084 .iter()
1085 .map(|p| input.tuples[p as usize].clone())
1086 .collect();
1087 Ok(EvaluatedNode {
1088 tuples: out,
1089 index_fn: order_output(strategy, truncation, input.index_fn, selection.len()),
1090 })
1091 }
1092}
1093
1094fn selected(len: u64) -> IndexFn {
1098 IndexFn::Lattice {
1099 axis_sizes: vec![len],
1100 }
1101}
1102
1103fn order_output(
1108 strategy: StrategyName,
1109 truncation: Option<u64>,
1110 input: Option<IndexFn>,
1111 len: u64,
1112) -> Option<IndexFn> {
1113 match (strategy, truncation) {
1114 (StrategyName::Lex, None) => input,
1115 (StrategyName::Lex, Some(_)) => input.map(|_| selected(len)),
1116 _ => Some(selected(len)),
1117 }
1118}
1119
1120impl EvalState<'_> {
1125 fn index_node(
1126 &mut self,
1127 node: &Comprehension,
1128 prefix: &[(String, Value)],
1129 ) -> Result<(Indexed, Option<IndexFn>), RuntimeError> {
1130 match node {
1131 Comprehension::Clause { name, source } => self.index_clause(name, source, prefix),
1132 Comprehension::Cartesian { children } => self.index_cartesian(children, prefix),
1133 Comprehension::Zip { children, mode } => self.index_zip(children, *mode, prefix),
1134 Comprehension::Union { children } => self.index_union(children, prefix),
1135 Comprehension::Filter { child, predicate } => {
1136 let (inner, _) = self.index_node(child, prefix)?;
1139 let predicate = CompiledPredicate::new(predicate);
1140 let mut kept = Vec::new();
1141 let mut tuple = RuntimeTuple::new();
1142 for i in 0..inner.len() {
1143 tuple.clear();
1144 inner.append_at(i, &mut tuple);
1145 if predicate.keeps(&tuple, self.scope)? {
1146 kept.push(tuple.clone());
1147 }
1148 }
1149 Ok((Indexed::Tuples(kept), None))
1150 }
1151 Comprehension::Order {
1152 child,
1153 strategy,
1154 truncation,
1155 seed,
1156 } => {
1157 let child = shape_input(child, *strategy);
1158 if has_continuous_axis(child) {
1159 let sampled =
1160 self.sample_space(child, prefix, *strategy, *truncation, *seed)?;
1161 let len = sampled.tuples.len() as u64;
1162 return Ok((Indexed::Tuples(sampled.tuples), Some(selected(len))));
1163 }
1164 if let Some((child, predicate)) = ranked_filter(child, *strategy) {
1169 let (inner, index_fn) = self.index_node(child, prefix)?;
1170 let predicate = CompiledPredicate::new(predicate);
1171 let mut survivors = Vec::new();
1172 let mut tuple = RuntimeTuple::new();
1173 for p in 0..inner.len() {
1174 tuple.clear();
1175 inner.append_at(p, &mut tuple);
1176 if predicate.keeps(&tuple, self.scope)? {
1177 survivors.push(p);
1178 }
1179 }
1180 let selection = surviving_selection(
1181 *strategy,
1182 index_fn.as_ref(),
1183 inner.len(),
1184 *truncation,
1185 *seed,
1186 &survivors,
1187 )?;
1188 let len = selection.len();
1189 return Ok((
1190 Indexed::Select {
1191 child: Box::new(inner),
1192 selection,
1193 },
1194 Some(selected(len)),
1195 ));
1196 }
1197 let (inner, index_fn) = self.index_node(child, prefix)?;
1198 let selection = order_selection(
1199 *strategy,
1200 index_fn.as_ref(),
1201 inner.len(),
1202 *truncation,
1203 *seed,
1204 )?;
1205 let len = selection.len();
1206 Ok((
1207 Indexed::Select {
1208 child: Box::new(inner),
1209 selection,
1210 },
1211 order_output(*strategy, *truncation, index_fn, len),
1212 ))
1213 }
1214 }
1215 }
1216
1217 fn index_clause(
1218 &mut self,
1219 name: &str,
1220 source: &Source,
1221 prefix: &[(String, Value)],
1222 ) -> Result<(Indexed, Option<IndexFn>), RuntimeError> {
1223 let values = match source {
1224 Source::IntRange { lo, hi, step } => {
1227 let step = (*step).max(1);
1228 let len = if hi <= lo {
1229 0
1230 } else {
1231 ((i128::from(*hi) - i128::from(*lo)) as u128).div_ceil(step as u128) as u64
1232 };
1233 self.record_yield(source, len as usize);
1234 ClauseValues::Range { lo: *lo, step, len }
1235 }
1236 _ => {
1237 let evaluated = self.evaluate_source(name, source, prefix)?;
1238 let index_fn = evaluated.index_fn;
1239 return Ok((
1240 Indexed::Clause {
1241 name: name.to_string(),
1242 values: ClauseValues::List(evaluated.values),
1243 },
1244 Some(index_fn),
1245 ));
1246 }
1247 };
1248 let len = values.len();
1249 Ok((
1250 Indexed::Clause {
1251 name: name.to_string(),
1252 values,
1253 },
1254 Some(IndexFn::Lattice {
1255 axis_sizes: vec![len],
1256 }),
1257 ))
1258 }
1259
1260 fn index_cartesian(
1266 &mut self,
1267 children: &[Comprehension],
1268 prefix: &[(String, Value)],
1269 ) -> Result<(Indexed, Option<IndexFn>), RuntimeError> {
1270 if children.is_empty() || references_an_earlier_axis(children) {
1271 let node = self.evaluate_cartesian(children, prefix)?;
1272 return Ok((Indexed::Tuples(node.tuples), node.index_fn));
1273 }
1274 let base = self.mult;
1275 let mut parts = Vec::with_capacity(children.len());
1276 let mut index_fns = Vec::with_capacity(children.len());
1277 let mut lens = Vec::with_capacity(children.len());
1278 let mut len: u64 = 1;
1279 for child in children {
1280 let evaluated = self.index_node(child, prefix);
1281 let (part, index_fn) = match evaluated {
1282 Ok(done) => done,
1283 Err(e) => {
1284 self.mult = base;
1285 return Err(e);
1286 }
1287 };
1288 let part_len = part.len();
1289 parts.push(part);
1290 index_fns.push(index_fn);
1291 lens.push(part_len);
1292 len = match len.checked_mul(part_len) {
1293 Some(n) => n,
1294 None => {
1295 self.mult = base;
1296 return Err(RuntimeError::UnsupportedShape(format!(
1297 "cartesian of {lens:?} tuples exceeds 2^64"
1298 )));
1299 }
1300 };
1301 if part_len == 0 {
1302 break;
1303 }
1304 self.mult = self
1305 .mult
1306 .saturating_mul(usize::try_from(part_len).unwrap_or(usize::MAX));
1307 }
1308 self.mult = base;
1309 let index_fn = combine_cartesian_index_fn(&index_fns);
1310 if len == 0 {
1311 return Ok((Indexed::Tuples(Vec::new()), index_fn));
1312 }
1313 Ok((
1314 Indexed::Product {
1315 children: parts,
1316 lens,
1317 len,
1318 },
1319 index_fn,
1320 ))
1321 }
1322
1323 fn index_zip(
1324 &mut self,
1325 children: &[Comprehension],
1326 mode: crate::iteration::comprehension::strategy::ZipMode,
1327 prefix: &[(String, Value)],
1328 ) -> Result<(Indexed, Option<IndexFn>), RuntimeError> {
1329 use crate::iteration::comprehension::strategy::ZipMode;
1330 if children.is_empty() {
1331 let node = self.evaluate_zip(children, mode, prefix)?;
1332 return Ok((Indexed::Tuples(node.tuples), node.index_fn));
1333 }
1334 let mut parts = Vec::with_capacity(children.len());
1335 for child in children {
1336 parts.push(self.index_node(child, prefix)?.0);
1337 }
1338 let lengths: Vec<u64> = parts.iter().map(Indexed::len).collect();
1339 let len = match mode {
1340 ZipMode::Strict => {
1341 let first = lengths[0];
1342 if lengths.iter().any(|&n| n != first) {
1343 return Err(RuntimeError::ZipLengthMismatch { lengths });
1344 }
1345 first
1346 }
1347 ZipMode::Truncate => lengths.iter().copied().min().unwrap_or(0),
1348 ZipMode::Cycle => cycle_length(&lengths),
1349 };
1350 Ok(match mode {
1351 ZipMode::Strict | ZipMode::Truncate => (
1352 Indexed::Lockstep {
1353 children: parts,
1354 len,
1355 },
1356 Some(IndexFn::Lockstep { length: len }),
1357 ),
1358 ZipMode::Cycle => (
1359 Indexed::Cycle {
1360 children: parts,
1361 len,
1362 },
1363 Some(IndexFn::Modular {
1364 axis_sizes: lengths,
1365 }),
1366 ),
1367 })
1368 }
1369
1370 fn index_union(
1371 &mut self,
1372 children: &[Comprehension],
1373 prefix: &[(String, Value)],
1374 ) -> Result<(Indexed, Option<IndexFn>), RuntimeError> {
1375 let mut parts = Vec::with_capacity(children.len());
1376 let mut segment_sizes = Vec::with_capacity(children.len());
1377 let mut all_segments_addressable = true;
1378 for child in children {
1379 let (part, index_fn) = self.index_node(child, prefix)?;
1380 all_segments_addressable &= index_fn.is_some();
1381 segment_sizes.push(part.len());
1382 parts.push(part);
1383 }
1384 let len = segment_sizes
1385 .iter()
1386 .try_fold(0u64, |acc, n| acc.checked_add(*n))
1387 .ok_or_else(|| {
1388 RuntimeError::UnsupportedShape(format!(
1389 "union of {segment_sizes:?} tuples exceeds 2^64"
1390 ))
1391 })?;
1392 let index_fn = all_segments_addressable.then_some(IndexFn::Concatenation { segment_sizes });
1393 Ok((
1394 Indexed::Concat {
1395 children: parts,
1396 len,
1397 },
1398 index_fn,
1399 ))
1400 }
1401}
1402
1403fn references_an_earlier_axis(children: &[Comprehension]) -> bool {
1407 let mut bound: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
1408 for child in children {
1409 if child
1410 .referenced_source_names()
1411 .iter()
1412 .any(|n| bound.contains(n))
1413 {
1414 return true;
1415 }
1416 collect_clause_names(child, &mut bound);
1417 }
1418 false
1419}
1420
1421fn collect_clause_names(c: &Comprehension, out: &mut std::collections::BTreeSet<String>) {
1423 match c {
1424 Comprehension::Clause { name, .. } => {
1425 out.insert(name.clone());
1426 }
1427 Comprehension::Cartesian { children }
1428 | Comprehension::Zip { children, .. }
1429 | Comprehension::Union { children } => {
1430 for child in children {
1431 collect_clause_names(child, out);
1432 }
1433 }
1434 Comprehension::Filter { child, .. } | Comprehension::Order { child, .. } => {
1435 collect_clause_names(child, out);
1436 }
1437 }
1438}
1439
1440fn order_selection(
1448 strategy: StrategyName,
1449 index_fn: Option<&IndexFn>,
1450 cardinality: u64,
1451 truncation: Option<u64>,
1452 seed: Option<u64>,
1453) -> Result<Selection, RuntimeError> {
1454 let dispatch = crate::iteration::comprehension::strategies::for_name(strategy);
1455 if !dispatch.accepts_input(index_fn) {
1456 return Err(RuntimeError::StrategyRejectsInput {
1457 strategy,
1458 index_fn: index_fn.cloned(),
1459 });
1460 }
1461 let fallback;
1462 let index_fn = match index_fn {
1463 Some(idx) => idx,
1464 None => {
1465 fallback = IndexFn::Lattice {
1466 axis_sizes: vec![cardinality],
1467 };
1468 &fallback
1469 }
1470 };
1471 Ok(dispatch.select(index_fn, cardinality, truncation, seed))
1472}
1473
1474fn surviving_selection(
1479 strategy: StrategyName,
1480 index_fn: Option<&IndexFn>,
1481 cardinality: u64,
1482 truncation: Option<u64>,
1483 seed: Option<u64>,
1484 survivors: &[u64],
1485) -> Result<Selection, RuntimeError> {
1486 let dispatch = crate::iteration::comprehension::strategies::for_name(strategy);
1487 let Some(index_fn) = index_fn.filter(|idx| dispatch.accepts_input(Some(idx))) else {
1488 return Err(RuntimeError::StrategyRejectsInput {
1489 strategy,
1490 index_fn: index_fn.cloned(),
1491 });
1492 };
1493 Ok(dispatch.select_surviving(index_fn, cardinality, truncation, seed, survivors))
1494}
1495
1496const SAMPLE_ROUNDS: u32 = 6;
1499
1500const UNIT_SCALE: f64 = (1u64 << 53) as f64;
1503
1504enum SampleAxis {
1509 Discrete(Vec<RuntimeTuple>),
1510 Continuous {
1511 name: String,
1512 interval: Interval,
1513 measure: AxisMeasure,
1514 },
1515}
1516
1517#[derive(Default)]
1521struct SampleSpace {
1522 axes: Vec<SampleAxis>,
1523 predicates: Vec<String>,
1524}
1525
1526impl SampleSpace {
1527 fn realize(&self, mi: &[u64], strategy: StrategyName) -> RuntimeTuple {
1535 let extrema = matches!(strategy, StrategyName::Extrema);
1536 let discrete_count = self
1537 .axes
1538 .iter()
1539 .filter(|a| matches!(a, SampleAxis::Discrete(_)))
1540 .count();
1541 let (mut d, mut c) = (0, if extrema { 0 } else { discrete_count });
1542 let mut out = RuntimeTuple::new();
1543 for axis in &self.axes {
1544 match axis {
1545 SampleAxis::Discrete(tuples) => {
1546 let pos = mi.get(d).copied().unwrap_or(0) as usize;
1547 d += 1;
1548 if extrema {
1549 c += 1;
1550 }
1551 if let Some(t) = tuples.get(pos) {
1552 out.extend(t.iter().cloned());
1553 }
1554 }
1555 SampleAxis::Continuous {
1556 name,
1557 interval,
1558 measure,
1559 } => {
1560 let code = mi.get(c).copied().unwrap_or(0);
1561 c += 1;
1562 if extrema {
1563 d += 1;
1564 }
1565 let x = if extrema {
1566 measure.endpoint(interval, code == 1)
1567 } else {
1568 measure.map_unit(code as f64 / UNIT_SCALE, interval)
1569 };
1570 out.push((name.clone(), Value::F64(x)));
1571 }
1572 }
1573 }
1574 out
1575 }
1576}
1577
1578fn draw_sample(
1583 index_fn: &IndexFn,
1584 strategy: StrategyName,
1585 count: Option<u64>,
1586 seed: Option<u64>,
1587) -> Result<Vec<Vec<u64>>, RuntimeError> {
1588 use crate::iteration::comprehension::strategies::{
1589 extrema::extrema_multi_indices, halton::try_halton_multi_indices,
1590 lhs::try_lhs_multi_indices, shuffle::try_shuffle_multi_indices,
1591 sobol::try_sobol_multi_indices,
1592 };
1593 if matches!(strategy, StrategyName::Extrema) {
1594 return Ok(extrema_multi_indices(index_fn, count));
1595 }
1596 let Some(n) = count else {
1597 return Err(RuntimeError::OrderEval {
1598 strategy,
1599 message: "a continuous source has no finite tuple set; give the order a count, \
1600 as in `order halton/16`"
1601 .into(),
1602 });
1603 };
1604 let drawn = match strategy {
1608 StrategyName::Halton => try_halton_multi_indices(index_fn, Some(n)),
1609 StrategyName::Sobol => try_sobol_multi_indices(index_fn, Some(n)),
1610 StrategyName::Lhs => try_lhs_multi_indices(index_fn, Some(n), seed),
1611 StrategyName::Shuffle => try_shuffle_multi_indices(index_fn, Some(n), seed),
1612 other => {
1613 return Err(RuntimeError::OrderEval {
1614 strategy: other,
1615 message: "a continuous source needs a sampling strategy: halton, sobol, lhs, \
1616 shuffle, or extrema"
1617 .into(),
1618 });
1619 }
1620 };
1621 drawn.map_err(|message| RuntimeError::OrderEval { strategy, message })
1622}
1623
1624pub(crate) fn has_continuous_axis(c: &Comprehension) -> bool {
1629 match c {
1630 Comprehension::Clause { source, .. } => matches!(
1631 source,
1632 Source::ContinuousInterval { .. } | Source::Distribution { .. }
1633 ),
1634 Comprehension::Cartesian { children } => children.iter().any(has_continuous_axis),
1635 Comprehension::Filter { child, .. } => has_continuous_axis(child),
1636 Comprehension::Zip { .. } | Comprehension::Union { .. } | Comprehension::Order { .. } => {
1637 false
1638 }
1639 }
1640}
1641
1642fn combine_cartesian_index_fn(children: &[Option<IndexFn>]) -> Option<IndexFn> {
1643 let mut axis_sizes = Vec::new();
1644 for opt in children {
1645 match opt {
1646 Some(IndexFn::Lattice { axis_sizes: a }) => axis_sizes.extend(a.iter().copied()),
1647 Some(IndexFn::Lockstep { length }) => axis_sizes.push(*length),
1648 _ => return None,
1651 }
1652 }
1653 Some(IndexFn::Lattice { axis_sizes })
1654}
1655
1656#[cfg(test)]
1657mod tests {
1658 use super::*;
1659 use crate::iteration::comprehension::source::LiteralValue;
1660
1661 fn empty_kernel() -> Arc<PolydatKernel> {
1662 Arc::new(crate::dsl::compile_polydat_interpreter("\n").unwrap())
1663 }
1664
1665 fn canonical_with_k() -> Arc<PolydatKernel> {
1670 Arc::new(crate::dsl::compile_polydat_interpreter("extern k: u64\n").unwrap())
1671 }
1672
1673 fn clause(name: &str, source: Source) -> Comprehension {
1674 Comprehension::Clause {
1675 name: name.into(),
1676 source,
1677 }
1678 }
1679
1680 fn empty_literal() -> Source {
1681 Source::Literal { values: Vec::new() }
1682 }
1683
1684 #[test]
1686 fn every_leaf_reports_what_it_yielded() {
1687 let comp = Comprehension::Cartesian {
1688 children: vec![
1689 clause(
1690 "a",
1691 Source::IntRange {
1692 lo: 0,
1693 hi: 3,
1694 step: 1,
1695 },
1696 ),
1697 clause(
1698 "b",
1699 Source::Literal {
1700 values: vec![LiteralValue::Int(7), LiteralValue::Int(8)],
1701 },
1702 ),
1703 ],
1704 };
1705 let scope = empty_kernel();
1706
1707 let out = evaluate_for_iteration_reported(&comp, &*scope).unwrap();
1708 assert_eq!(out.tuples.len(), 6, "3 x 2");
1709 assert_eq!(out.clauses.len(), 2, "one entry per leaf, in tree order");
1710 assert_eq!(out.clauses[0].var, "a");
1711 assert_eq!(out.clauses[0].values, 3);
1712 assert_eq!(out.clauses[1].var, "b");
1713 assert_eq!(out.clauses[1].evaluations, 3);
1716 assert_eq!(out.clauses[1].values, 6);
1717 }
1718
1719 #[test]
1722 fn an_empty_clause_is_reached_and_yields_nothing() {
1723 let comp = Comprehension::Cartesian {
1724 children: vec![
1725 clause(
1726 "a",
1727 Source::IntRange {
1728 lo: 0,
1729 hi: 2,
1730 step: 1,
1731 },
1732 ),
1733 clause("b", empty_literal()),
1734 ],
1735 };
1736 let scope = empty_kernel();
1737
1738 let out = evaluate_for_iteration_reported(&comp, &*scope).unwrap();
1739 assert!(out.tuples.is_empty(), "an empty clause empties the product");
1740 let culprits: Vec<&str> = out
1741 .clauses
1742 .iter()
1743 .filter(|c| c.evaluations > 0 && c.values == 0)
1744 .map(|c| c.var.as_str())
1745 .collect();
1746 assert_eq!(culprits, ["b"], "only the empty clause is named");
1747 }
1748
1749 #[test]
1752 fn a_clause_behind_an_empty_one_is_never_reached() {
1753 let comp = Comprehension::Cartesian {
1754 children: vec![
1755 clause("outer", empty_literal()),
1756 clause(
1757 "inner",
1758 Source::IntRange {
1759 lo: 0,
1760 hi: 9,
1761 step: 1,
1762 },
1763 ),
1764 ],
1765 };
1766 let scope = empty_kernel();
1767
1768 let out = evaluate_for_iteration_reported(&comp, &*scope).unwrap();
1769 assert!(out.tuples.is_empty());
1770 let by = |v: &str| {
1771 out.clauses
1772 .iter()
1773 .find(|c| c.var == v)
1774 .expect("every leaf is present whether reached or not")
1775 };
1776 assert_eq!(by("outer").evaluations, 1);
1777 assert_eq!(by("outer").values, 0);
1778 assert_eq!(
1779 by("inner").evaluations,
1780 0,
1781 "never reached: the cause is `outer`, not this"
1782 );
1783 assert_eq!(by("inner").values, 0);
1784 }
1785
1786 #[test]
1789 fn clauses_sharing_a_name_across_a_union_are_counted_apart() {
1790 let comp = Comprehension::Union {
1791 children: vec![
1792 clause(
1793 "k",
1794 Source::Literal {
1795 values: vec![LiteralValue::Int(1)],
1796 },
1797 ),
1798 clause("k", empty_literal()),
1799 ],
1800 };
1801 let scope = empty_kernel();
1802
1803 let out = evaluate_for_iteration_reported(&comp, &*scope).unwrap();
1804 assert_eq!(out.clauses.len(), 2, "two leaves, one name");
1805 assert_eq!(out.clauses[0].values, 1);
1806 assert_eq!(out.clauses[1].values, 0);
1807 assert_eq!(out.clauses[1].evaluations, 1, "reached, and empty");
1808 }
1809
1810 #[test]
1813 fn the_plain_entry_point_agrees_with_the_reported_one() {
1814 let comp = Comprehension::Cartesian {
1815 children: vec![
1816 clause(
1817 "a",
1818 Source::IntRange {
1819 lo: 1,
1820 hi: 4,
1821 step: 1,
1822 },
1823 ),
1824 clause(
1825 "b",
1826 Source::Literal {
1827 values: vec![LiteralValue::Int(5)],
1828 },
1829 ),
1830 ],
1831 };
1832 let scope = empty_kernel();
1833
1834 let plain = evaluate_for_iteration(&comp, &*scope).unwrap();
1835 let reported = evaluate_for_iteration_reported(&comp, &*scope).unwrap();
1836 assert_eq!(plain, reported.tuples);
1837 }
1838
1839 #[test]
1840 fn int_range_yields_values() {
1841 let comp = Comprehension::Clause {
1842 name: "k".into(),
1843 source: Source::IntRange {
1844 lo: 1,
1845 hi: 5,
1846 step: 1,
1847 },
1848 };
1849 let canonical = empty_kernel();
1850
1851 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1852 assert_eq!(tuples.len(), 4);
1853 assert_eq!(tuples[0][0].1, Value::U64(1));
1854 assert_eq!(tuples[3][0].1, Value::U64(4));
1855 }
1856
1857 #[test]
1858 fn literal_list_yields_values() {
1859 let comp = Comprehension::Clause {
1860 name: "x".into(),
1861 source: Source::Literal {
1862 values: vec![LiteralValue::Int(10), LiteralValue::Int(20)],
1863 },
1864 };
1865 let canonical = empty_kernel();
1866
1867 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1868 assert_eq!(tuples.len(), 2);
1869 }
1870
1871 #[test]
1872 fn cartesian_produces_product() {
1873 let comp = Comprehension::cartesian(vec![
1874 Comprehension::Clause {
1875 name: "x".into(),
1876 source: Source::IntRange {
1877 lo: 1,
1878 hi: 3,
1879 step: 1,
1880 },
1881 },
1882 Comprehension::Clause {
1883 name: "y".into(),
1884 source: Source::IntRange {
1885 lo: 10,
1886 hi: 30,
1887 step: 10,
1888 },
1889 },
1890 ]);
1891 let canonical = empty_kernel();
1892
1893 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1894 assert_eq!(tuples.len(), 4);
1896 }
1897
1898 #[test]
1899 fn union_produces_concatenation() {
1900 let comp = Comprehension::union(vec![
1901 Comprehension::Clause {
1902 name: "k".into(),
1903 source: Source::Literal {
1904 values: vec![LiteralValue::Int(1)],
1905 },
1906 },
1907 Comprehension::Clause {
1908 name: "k".into(),
1909 source: Source::Literal {
1910 values: vec![LiteralValue::Int(10), LiteralValue::Int(20)],
1911 },
1912 },
1913 ]);
1914 let canonical = empty_kernel();
1915
1916 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1917 assert_eq!(tuples.len(), 3);
1918 }
1919
1920 #[test]
1921 fn filter_drops_non_matching() {
1922 let comp = Comprehension::filter(
1923 Comprehension::Clause {
1924 name: "k".into(),
1925 source: Source::IntRange {
1926 lo: 1,
1927 hi: 6,
1928 step: 1,
1929 },
1930 },
1931 "{k} > 3",
1932 );
1933 let canonical = canonical_with_k();
1934
1935 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1936 assert_eq!(tuples.len(), 2);
1938 }
1939
1940 #[test]
1941 fn order_lex_truncate() {
1942 let comp = Comprehension::order(
1943 Comprehension::Clause {
1944 name: "k".into(),
1945 source: Source::IntRange {
1946 lo: 1,
1947 hi: 100,
1948 step: 1,
1949 },
1950 },
1951 StrategyName::Lex,
1952 Some(5),
1953 );
1954 let canonical = empty_kernel();
1955
1956 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
1957 assert_eq!(tuples.len(), 5);
1958 }
1959
1960 #[test]
1964 fn a_multi_name_head_is_one_lattice_axis() {
1965 use crate::iteration::comprehension::strategy::ZipMode;
1966 let lit = |name: &str, vs: &[i64]| Comprehension::Clause {
1967 name: name.into(),
1968 source: Source::Literal {
1969 values: vs.iter().map(|v| LiteralValue::Int(*v)).collect(),
1970 },
1971 };
1972 let comp = Comprehension::order(
1973 Comprehension::cartesian(vec![
1974 Comprehension::zip(vec![lit("a", &[1, 2]), lit("b", &[3, 4])], ZipMode::Strict),
1975 lit("c", &[5, 6, 7]),
1976 ]),
1977 StrategyName::Halton,
1978 Some(6),
1979 );
1980 let tuples = evaluate_for_iteration(&comp, &*empty_kernel()).unwrap();
1981 assert_eq!(
1982 tuples.len(),
1983 6,
1984 "every tuple of the 2 x 3 product: {tuples:?}"
1985 );
1986 }
1987
1988 #[test]
1994 fn extrema_over_cartesian_uses_indexed_form() {
1995 let comp = Comprehension::order(
1996 Comprehension::cartesian(vec![
1997 Comprehension::Clause {
1998 name: "k".into(),
1999 source: Source::Literal {
2000 values: vec![
2001 LiteralValue::Int(1),
2002 LiteralValue::Int(2),
2003 LiteralValue::Int(3),
2004 ],
2005 },
2006 },
2007 Comprehension::Clause {
2008 name: "limit".into(),
2009 source: Source::Literal {
2010 values: vec![
2011 LiteralValue::Int(10),
2012 LiteralValue::Int(20),
2013 LiteralValue::Int(30),
2014 ],
2015 },
2016 },
2017 ]),
2018 StrategyName::Extrema,
2019 Some(1),
2023 );
2024 let canonical = empty_kernel();
2025
2026 let tuples = evaluate_for_iteration(&comp, &*canonical).unwrap();
2027 assert_eq!(tuples.len(), 4);
2029 for t in &tuples {
2031 assert_eq!(t.len(), 2);
2032 let k = match &t[0].1 {
2033 Value::U64(n) => *n,
2034 other => panic!("expected u64 k, got {other:?}"),
2035 };
2036 let lim = match &t[1].1 {
2037 Value::U64(n) => *n,
2038 other => panic!("expected u64 limit, got {other:?}"),
2039 };
2040 assert!(k == 1 || k == 3, "expected extreme k, got {k}");
2041 assert!(lim == 10 || lim == 30, "expected extreme limit, got {lim}");
2042 }
2043 }
2044}