use std::collections::HashMap;
use super::ast::{Clause, Comprehension, ShellOrigin, TraversalOrder, ZipMode};
pub fn parse_clause(s: &str) -> Result<Clause, String> {
let bytes = s.as_bytes();
let mut depth: i32 = 0;
let mut i: usize = 0;
while i + 4 <= bytes.len() {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth += 1; i += 1; }
b')' | b']' | b'}' => { depth -= 1; i += 1; }
b' ' if depth == 0
&& bytes.get(i + 1) == Some(&b'i')
&& bytes.get(i + 2) == Some(&b'n')
&& bytes.get(i + 3) == Some(&b' ') =>
{
let lhs = s[..i].trim();
let rhs = s[i + 4..].trim();
return parse_clause_from_sides(lhs, rhs, s);
}
_ => { i += 1; }
}
}
Err(format!("invalid for_each clause: '{s}' (expected 'var in expr')"))
}
fn parse_clause_from_sides(lhs: &str, rhs: &str, whole: &str) -> Result<Clause, String> {
let lhs_paren = is_paren_wrapped(lhs);
let (rhs_inner, mode) = match strip_zip_mode_prefix(rhs) {
Some((inner, m)) => (inner, m),
None => (rhs.to_string(), ZipMode::Strict),
};
let rhs_paren = is_paren_wrapped(&rhs_inner);
let rhs_explicit_paren = mode != ZipMode::Strict || rhs_paren;
if lhs_paren || rhs_explicit_paren {
if !(lhs_paren && rhs_explicit_paren) {
return Err(format!(
"parallel-iter clause '{whole}' requires parentheses on both sides \
(e.g. '(a, b) in (e1, e2)' or '(a, b) in zip_truncate(e1, e2)')"
));
}
let vars = split_paren_group(lhs);
let exprs = split_paren_group(&rhs_inner);
for v in &vars {
if !is_simple_ident(v) {
return Err(format!(
"parallel-iter clause '{whole}': '{v}' is not a valid variable name"
));
}
}
if vars.len() < 2 {
return Err(format!(
"parallel-iter clause '{whole}' requires ≥ 2 variables \
(use 'var in expr' for the single-var form)"
));
}
if vars.len() != exprs.len() {
return Err(format!(
"parallel-iter clause '{whole}': {} variables but {} expressions",
vars.len(), exprs.len()
));
}
Ok(Clause::parallel_with_mode(mode, vars, exprs))
} else {
Ok(Clause::new(lhs, rhs))
}
}
fn strip_zip_mode_prefix(rhs: &str) -> Option<(String, ZipMode)> {
for (prefix, mode) in [
("zip_truncate", ZipMode::Truncate),
("zip_cycle", ZipMode::Cycle),
] {
if let Some(rest) = rhs.strip_prefix(prefix) {
let trimmed = rest.trim_start();
if trimmed.starts_with('(') && is_paren_wrapped(trimmed) {
return Some((trimmed.to_string(), mode));
}
}
}
None
}
fn is_paren_wrapped(s: &str) -> bool {
let bytes = s.as_bytes();
if bytes.first() != Some(&b'(') || bytes.last() != Some(&b')') {
return false;
}
let mut depth: i32 = 0;
for (i, &b) in bytes.iter().enumerate() {
match b {
b'(' => depth += 1,
b')' => {
depth -= 1;
if depth == 0 {
return i == bytes.len() - 1;
}
}
_ => {}
}
}
false
}
fn split_paren_group(s: &str) -> Vec<String> {
let inner = &s[1..s.len() - 1];
let bytes = inner.as_bytes();
let mut parts: Vec<String> = Vec::new();
let mut start: usize = 0;
let mut depth: i32 = 0;
let mut i: usize = 0;
while i < bytes.len() {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth += 1; i += 1; }
b')' | b']' | b'}' => { depth -= 1; i += 1; }
b',' if depth == 0 => {
parts.push(inner[start..i].trim().to_string());
start = i + 1;
i += 1;
}
_ => { i += 1; }
}
}
let tail = inner[start..].trim();
if !tail.is_empty() {
parts.push(tail.to_string());
}
parts
}
fn is_simple_ident(s: &str) -> bool {
let mut chars = s.chars();
match chars.next() {
Some(c) if c.is_ascii_alphabetic() || c == '_' => {}
_ => return false,
}
chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
}
pub fn parse_comprehension_text(text: &str) -> Result<Comprehension, String> {
let (head, order_text) = split_at_order(text);
let (clause_text, filter) = split_at_where(&head);
let clauses = parse_clause_list(&clause_text)?;
let subspaces: Vec<Vec<Clause>> = clauses.into_iter()
.map(|c| vec![c])
.collect();
let mut comp = comprehension_from_subspaces(subspaces);
if let Some(predicate) = filter {
comp = comp.with_filter(predicate);
}
if let Some(spec) = order_text {
comp = comp.with_order(parse_order_spec(&spec)?);
}
comp.validate().map_err(|errs| errs.join("; "))?;
Ok(comp)
}
#[deprecated(note = "use Comprehension::validate() — single source of truth for AST invariants")]
pub fn validate_order_for_mode(
mode: &super::ast::ComprehensionMode,
order: &Option<TraversalOrder>,
) -> Result<(), String> {
super::ast::check_order_for_mode(mode, order)
}
pub fn parse_order_spec(text: &str) -> Result<TraversalOrder, String> {
let trimmed = text.trim();
if trimmed.is_empty() {
return Err("order spec is empty".to_string());
}
if let Some(open) = trimmed.find('(') {
if !trimmed.ends_with(')') {
return Err(format!(
"order spec '{trimmed}': unbalanced parens — expected `name(...)`"
));
}
let name = trimmed[..open].trim();
let body = &trimmed[open + 1..trimmed.len() - 1];
return build_order_from_keyword(name, body);
}
let (name, n_opt) = match trimmed.find('/') {
Some(slash) => {
let n_text = trimmed[slash + 1..].trim();
let n = n_text.parse::<usize>().map_err(|_| format!(
"order spec '{trimmed}': '/N' suffix must be a non-negative integer, got '{n_text}'"
))?;
(trimmed[..slash].trim(), Some(n))
}
None => (trimmed, None),
};
build_order_from_terse(name, n_opt)
}
fn build_order_from_terse(name: &str, n: Option<usize>) -> Result<TraversalOrder, String> {
match name {
"lex" => Ok(TraversalOrder::Lex { count: n }),
"reverse_lex" => Ok(TraversalOrder::ReverseLex { count: n }),
"diagonal" => Ok(TraversalOrder::Diagonal { count: n }),
"antidiagonal" => Ok(TraversalOrder::Antidiagonal { count: n }),
"extrema" => Ok(TraversalOrder::Extrema { strata: n }),
"shells" => Ok(TraversalOrder::Shells {
origin: ShellOrigin::Outer,
depth: n,
}),
"halton" => Ok(TraversalOrder::Halton { count: n }),
"sobol" => Ok(TraversalOrder::Sobol { count: n }),
"lhs" => Ok(TraversalOrder::Lhs { count: n, seed: None }),
"custom" => Err(format!(
"order spec 'custom': use 'custom(<function>)' to name the GK function"
)),
other => Err(format!(
"order spec: unknown strategy '{other}' — \
expected one of lex/reverse_lex/diagonal/antidiagonal/extrema/shells/halton/sobol/lhs/custom"
)),
}
}
fn build_order_from_keyword(name: &str, body: &str) -> Result<TraversalOrder, String> {
let args = parse_keyword_args(body)?;
let count = args.iter().find_map(|(k, v)| (k == "count").then(|| v.parse::<usize>().ok()).flatten());
let depth = args.iter().find_map(|(k, v)| (k == "depth").then(|| v.parse::<usize>().ok()).flatten());
let strata = args.iter().find_map(|(k, v)| (k == "strata").then(|| v.parse::<usize>().ok()).flatten());
let seed = args.iter().find_map(|(k, v)| (k == "seed").then(|| v.parse::<u64>().ok()).flatten());
match name {
"lex" => Ok(TraversalOrder::Lex { count }),
"reverse_lex" => Ok(TraversalOrder::ReverseLex { count }),
"diagonal" => Ok(TraversalOrder::Diagonal { count }),
"antidiagonal" => Ok(TraversalOrder::Antidiagonal { count }),
"extrema" => Ok(TraversalOrder::Extrema { strata }),
"shells" => {
let origin = match args.iter().find_map(|(k, v)| (k == "origin").then_some(v.as_str())) {
Some("outer") | None => ShellOrigin::Outer,
Some("center") => ShellOrigin::Center,
Some("corner") => ShellOrigin::Corner,
Some(other) => return Err(format!(
"order shells: unknown origin '{other}' — expected outer/center/corner"
)),
};
Ok(TraversalOrder::Shells { origin, depth })
}
"halton" => Ok(TraversalOrder::Halton { count }),
"sobol" => Ok(TraversalOrder::Sobol { count }),
"lhs" => Ok(TraversalOrder::Lhs { count, seed }),
"space_filling" => {
let strategy = args.iter()
.find(|(k, _)| k.is_empty())
.map(|(_, v)| v.as_str())
.ok_or_else(|| "space_filling: missing strategy name (halton/sobol/lhs)".to_string())?;
match strategy {
"halton" => Ok(TraversalOrder::Halton { count }),
"sobol" => Ok(TraversalOrder::Sobol { count }),
"lhs" => Ok(TraversalOrder::Lhs { count, seed }),
other => Err(format!(
"space_filling: unknown strategy '{other}' — expected halton/sobol/lhs"
)),
}
}
"custom" => {
let function = args.iter()
.find(|(k, _)| k.is_empty())
.map(|(_, v)| v.clone())
.ok_or_else(|| "custom: missing function name".to_string())?;
Ok(TraversalOrder::Custom { function })
}
other => Err(format!(
"order spec: unknown strategy '{other}'"
)),
}
}
fn parse_keyword_args(body: &str) -> Result<Vec<(String, String)>, String> {
let mut out = Vec::new();
let bytes = body.as_bytes();
let n = bytes.len();
let mut start = 0;
let mut i = 0;
let mut depth: u32 = 0;
let push = |s: &str, out: &mut Vec<(String, String)>| {
let trimmed = s.trim();
if trimmed.is_empty() { return; }
let (k, v) = if let Some(eq) = trimmed.find('=') {
(trimmed[..eq].trim().to_string(), trimmed[eq + 1..].trim().to_string())
} else {
(String::new(), trimmed.to_string())
};
let v = strip_quotes(&v);
out.push((k, v));
};
while i < n {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth = depth.saturating_add(1); i += 1; }
b')' | b']' | b'}' => { depth = depth.saturating_sub(1); i += 1; }
b',' if depth == 0 => {
push(&body[start..i], &mut out);
start = i + 1;
i += 1;
}
_ => { i += 1; }
}
}
push(&body[start..], &mut out);
Ok(out)
}
fn strip_quotes(s: &str) -> String {
let s = s.trim();
if (s.starts_with('"') && s.ends_with('"') && s.len() >= 2)
|| (s.starts_with('\'') && s.ends_with('\'') && s.len() >= 2)
{
s[1..s.len() - 1].to_string()
} else {
s.to_string()
}
}
pub fn split_at_order(text: &str) -> (String, Option<String>) {
const KEYWORD: &str = " order ";
let bytes = text.as_bytes();
let n = bytes.len();
let mut depth: u32 = 0;
let mut i: usize = 0;
while i < n {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth = depth.saturating_add(1); i += 1; }
b')' | b']' | b'}' => { depth = depth.saturating_sub(1); i += 1; }
b' ' if depth == 0 && text.is_char_boundary(i)
&& text[i..].starts_with(KEYWORD) =>
{
let head = text[..i].to_string();
let spec = text[i + KEYWORD.len()..].trim().to_string();
if spec.is_empty() {
return (text.to_string(), None);
}
return (head, Some(spec));
}
_ => { i += 1; }
}
}
(text.to_string(), None)
}
pub fn split_at_where(text: &str) -> (String, Option<String>) {
const KEYWORD: &str = " where ";
let bytes = text.as_bytes();
let n = bytes.len();
let mut depth: u32 = 0;
let mut i: usize = 0;
while i < n {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth = depth.saturating_add(1); i += 1; }
b')' | b']' | b'}' => { depth = depth.saturating_sub(1); i += 1; }
b' ' if depth == 0 && text.is_char_boundary(i)
&& text[i..].starts_with(KEYWORD) =>
{
let prefix = text[..i].to_string();
let suffix = text[i + KEYWORD.len()..].trim().to_string();
if suffix.is_empty() {
return (text.to_string(), None);
}
return (prefix, Some(suffix));
}
_ => { i += 1; }
}
}
(text.to_string(), None)
}
pub fn parse_clause_list(text: &str) -> Result<Vec<Clause>, String> {
let mut out = Vec::new();
for part in split_respecting_parens(text) {
let trimmed = part.trim();
if trimmed.is_empty() { continue; }
out.push(parse_clause(trimmed)?);
}
Ok(out)
}
pub fn comprehension_from_subspaces(subspaces: Vec<Vec<Clause>>) -> Comprehension {
let mut counts: HashMap<&str, usize> = HashMap::new();
for set in &subspaces {
for clause in set {
for v in &clause.vars {
*counts.entry(v.as_str()).or_insert(0) += 1;
}
}
}
let any_repeat = counts.values().any(|c| *c > 1);
if any_repeat {
Comprehension::union(subspaces)
} else {
let flat: Vec<Clause> = subspaces.into_iter().flatten().collect();
Comprehension::cartesian(flat)
}
}
pub fn split_respecting_parens(s: &str) -> Vec<String> {
let bytes = s.as_bytes();
let mut parts: Vec<String> = Vec::new();
let mut start: usize = 0;
let mut i: usize = 0;
let mut depth: u32 = 0;
while i < bytes.len() {
let ch = bytes[i];
match ch {
b'(' | b'[' | b'{' => { depth = depth.saturating_add(1); i += 1; }
b')' | b']' | b'}' => { depth = depth.saturating_sub(1); i += 1; }
b',' if depth == 0 => {
if is_clause_boundary(&s[i + 1..]) {
parts.push(s[start..i].to_string());
start = i + 1;
i += 1;
} else {
i += 1;
}
}
_ => { i += 1; }
}
}
let tail = &s[start..];
if !tail.trim().is_empty() {
parts.push(tail.to_string());
}
parts
}
fn is_clause_boundary(tail: &str) -> bool {
let trimmed = tail.trim_start();
if trimmed.starts_with('(') {
let bytes = trimmed.as_bytes();
let mut depth: i32 = 0;
for (i, &b) in bytes.iter().enumerate() {
match b {
b'(' => depth += 1,
b')' => {
depth -= 1;
if depth == 0 {
let after = &trimmed[i + 1..];
return after.starts_with(" in ");
}
}
_ => {}
}
}
return false;
}
let mut ident_end = 0;
for (i, c) in trimmed.char_indices() {
if i == 0 {
if !(c.is_ascii_alphabetic() || c == '_') { return false; }
} else if !(c.is_ascii_alphanumeric() || c == '_') {
ident_end = i;
break;
}
ident_end = i + c.len_utf8();
}
if ident_end == 0 { return false; }
let after = &trimmed[ident_end..];
after.starts_with(" in ")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_single_clause() {
let c = parse_clause("k in {k_values}").unwrap();
assert_eq!(c.var(), "k");
assert_eq!(c.expr(), "{k_values}");
}
#[test]
fn parse_clause_rejects_malformed() {
let err = parse_clause("not a clause").unwrap_err();
assert!(err.contains("invalid for_each clause"));
}
#[test]
fn split_respects_parens_in_function_call() {
let s = "profile in matching_profiles('{dataset}', '{prefix}')";
let parts = split_respecting_parens(s);
assert_eq!(parts.len(), 1);
assert_eq!(parts[0], s);
}
#[test]
fn split_respects_inner_value_list_commas() {
let s = "k in 10, limit in 10,20,30";
let parts = split_respecting_parens(s);
assert_eq!(parts.len(), 2);
assert_eq!(parts[0], "k in 10");
assert_eq!(parts[1].trim(), "limit in 10,20,30");
}
#[test]
fn split_handles_multiple_real_boundaries() {
let s = "a in 1, b in 2, c in 3";
let parts = split_respecting_parens(s);
assert_eq!(parts.len(), 3);
}
#[test]
fn parse_clause_list_distinct_names() {
let clauses = parse_clause_list("k in {k_values}, limit in {k_{k}_limits}").unwrap();
assert_eq!(clauses.len(), 2);
assert_eq!(clauses[0].var(), "k");
assert_eq!(clauses[1].var(), "limit");
}
#[test]
fn parse_clause_list_paren_safe() {
let clauses = parse_clause_list(
"profile in matching_profiles('{dataset}', '{prefix}')"
).unwrap();
assert_eq!(clauses.len(), 1);
assert_eq!(clauses[0].var(), "profile");
assert_eq!(clauses[0].expr(),
"matching_profiles('{dataset}', '{prefix}')");
}
#[test]
fn comprehension_from_subspaces_distinct_names_flattens_to_cartesian() {
let subspaces = vec![
vec![Clause::new("k", "10")],
vec![Clause::new("limit", "20")],
];
let c = comprehension_from_subspaces(subspaces);
assert!(c.is_cartesian());
assert_eq!(c.coordinate_names(), vec!["k", "limit"]);
assert_eq!(c.flat_clauses().len(), 2);
}
#[test]
fn comprehension_from_subspaces_repeated_names_yields_union() {
let subspaces = vec![
vec![Clause::new("k", "10"), Clause::new("limit", "10,20,30")],
vec![Clause::new("k", "100"), Clause::new("limit", "100,200,300")],
];
let c = comprehension_from_subspaces(subspaces);
assert!(c.is_union());
assert_eq!(c.coordinate_names(), vec!["k", "limit"]);
assert_eq!(c.flat_clauses().len(), 4);
}
#[test]
fn split_at_where_simple() {
let (clauses, filter) = split_at_where("k in 10,100 where k > 5");
assert_eq!(clauses, "k in 10,100");
assert_eq!(filter, Some("k > 5".to_string()));
}
#[test]
fn split_at_where_no_predicate() {
let (clauses, filter) = split_at_where("k in 10,100, limit in 10,20,30");
assert_eq!(clauses, "k in 10,100, limit in 10,20,30");
assert_eq!(filter, None);
}
#[test]
fn split_at_where_inside_parens_is_ignored() {
let (clauses, filter) = split_at_where(
"p in pick(profiles, where='ann') where p == 'x'"
);
assert_eq!(clauses, "p in pick(profiles, where='ann')");
assert_eq!(filter, Some("p == 'x'".to_string()));
}
#[test]
fn parse_comprehension_text_no_filter() {
let comp = parse_comprehension_text("k in 10,100, limit in 10,20,30").unwrap();
assert!(comp.is_cartesian());
assert_eq!(comp.coordinate_names(), vec!["k", "limit"]);
assert_eq!(comp.filter, None);
}
#[test]
fn parse_comprehension_text_with_filter() {
let comp = parse_comprehension_text(
"k in 10,100, limit in 10,20,30 where k * limit < 1000"
).unwrap();
assert!(comp.is_cartesian());
assert_eq!(comp.coordinate_names(), vec!["k", "limit"]);
assert_eq!(comp.filter, Some("k * limit < 1000".to_string()));
}
#[test]
fn parse_comprehension_text_repeated_var_yields_union_with_filter() {
let comp = parse_comprehension_text(
"k in 1, k in 2 where k > 0"
).unwrap();
assert!(comp.is_union());
assert_eq!(comp.filter, Some("k > 0".to_string()));
}
#[test]
fn split_at_order_simple() {
let (head, order) = split_at_order("k in 1..10 order extrema/1");
assert_eq!(head, "k in 1..10");
assert_eq!(order, Some("extrema/1".to_string()));
}
#[test]
fn split_at_order_no_order() {
let (head, order) = split_at_order("k in 1..10 where {k} > 5");
assert_eq!(head, "k in 1..10 where {k} > 5");
assert_eq!(order, None);
}
#[test]
fn split_at_order_inside_parens_is_ignored() {
let (head, order) = split_at_order(
"p in pick(profiles, order='ann') order lex"
);
assert_eq!(head, "p in pick(profiles, order='ann')");
assert_eq!(order, Some("lex".to_string()));
}
#[test]
fn parse_order_spec_bare() {
match parse_order_spec("lex").unwrap() {
TraversalOrder::Lex { count: None } => {}
other => panic!("expected Lex, got {other:?}"),
}
match parse_order_spec("extrema").unwrap() {
TraversalOrder::Extrema { strata: None } => {}
other => panic!("expected Extrema, got {other:?}"),
}
match parse_order_spec("shells").unwrap() {
TraversalOrder::Shells { origin: ShellOrigin::Outer, depth: None } => {}
other => panic!("expected Shells outer/None, got {other:?}"),
}
}
#[test]
fn parse_order_spec_terse() {
match parse_order_spec("extrema/1").unwrap() {
TraversalOrder::Extrema { strata: Some(1) } => {}
other => panic!("expected Extrema strata=1, got {other:?}"),
}
match parse_order_spec("shells/2").unwrap() {
TraversalOrder::Shells { origin: ShellOrigin::Outer, depth: Some(2) } => {}
other => panic!("expected Shells outer/2, got {other:?}"),
}
match parse_order_spec("halton/64").unwrap() {
TraversalOrder::Halton { count: Some(64) } => {}
other => panic!("expected Halton count=64, got {other:?}"),
}
match parse_order_spec("lex/100").unwrap() {
TraversalOrder::Lex { count: Some(100) } => {}
other => panic!("expected Lex count=100, got {other:?}"),
}
}
#[test]
fn parse_order_spec_keyword() {
match parse_order_spec("shells(origin=center, depth=3)").unwrap() {
TraversalOrder::Shells { origin: ShellOrigin::Center, depth: Some(3) } => {}
other => panic!("expected Shells center/3, got {other:?}"),
}
match parse_order_spec("lhs(count=20, seed=42)").unwrap() {
TraversalOrder::Lhs { count: Some(20), seed: Some(42) } => {}
other => panic!("expected Lhs count=20 seed=42, got {other:?}"),
}
match parse_order_spec("space_filling(sobol, count=64)").unwrap() {
TraversalOrder::Sobol { count: Some(64) } => {}
other => panic!("expected Sobol count=64, got {other:?}"),
}
}
#[test]
fn parse_order_spec_unknown_strategy_errors() {
let err = parse_order_spec("zigzag").unwrap_err();
assert!(err.contains("unknown strategy"), "got: {err}");
}
#[test]
fn parse_comprehension_text_with_order() {
let comp = parse_comprehension_text(
"k in 1..10 where {k} > 3 order extrema/1"
).unwrap();
assert_eq!(comp.filter, Some("{k} > 3".to_string()));
match comp.order {
Some(TraversalOrder::Extrema { strata: Some(1) }) => {}
other => panic!("expected Extrema strata=1, got {other:?}"),
}
}
#[test]
fn parse_comprehension_text_order_only() {
let comp = parse_comprehension_text(
"k in 1..10, l in 1..10 order halton/50"
).unwrap();
assert_eq!(comp.filter, None);
assert!(matches!(comp.order, Some(TraversalOrder::Halton { count: Some(50) })));
}
#[test]
fn comprehension_from_subspaces_string_form_repeated_var_yields_union() {
let subspaces = vec![
vec![Clause::new("k", "1")],
vec![Clause::new("k", "2")],
];
let c = comprehension_from_subspaces(subspaces);
assert!(c.is_union());
assert_eq!(c.coordinate_names(), vec!["k"]);
}
#[test]
fn union_plus_extrema_is_rejected() {
let err = parse_comprehension_text(
"k in 10, k in 100 order extrema/1"
).unwrap_err();
assert!(err.contains("'extrema'") && err.contains("Union"),
"wrong message: {err}");
assert!(err.contains("Cartesian") || err.contains("lex"),
"should hint at remedy: {err}");
}
#[test]
fn union_plus_halton_is_rejected() {
let err = parse_comprehension_text(
"k in 10, l in 100, k in 200, l in 400 order halton/64"
).unwrap_err();
assert!(err.contains("'halton'") && err.contains("Union"), "{err}");
}
#[test]
fn union_plus_shells_is_rejected() {
let err = parse_comprehension_text(
"k in 10, k in 100 order shells/2"
).unwrap_err();
assert!(err.contains("'shells'") && err.contains("Union"), "{err}");
}
#[test]
fn union_plus_lex_is_accepted() {
let comp = parse_comprehension_text(
"k in 10, k in 100 order lex"
).unwrap();
assert!(comp.is_union());
assert!(matches!(comp.order, Some(TraversalOrder::Lex { count: None })));
}
#[test]
fn union_plus_custom_is_accepted() {
let comp = parse_comprehension_text(
"k in 10, k in 100 order custom(my_fn)"
).unwrap();
assert!(comp.is_union());
assert!(matches!(comp.order, Some(TraversalOrder::Custom { .. })));
}
#[test]
fn cartesian_plus_extrema_remains_valid() {
let comp = parse_comprehension_text(
"k in 1..10, l in 1..10 order extrema/1"
).unwrap();
assert!(comp.is_cartesian());
assert!(matches!(comp.order, Some(TraversalOrder::Extrema { strata: Some(1) })));
}
#[test]
fn validate_rejects_each_index_space_strategy_on_union() {
for (label, ord) in [
("reverse_lex", TraversalOrder::ReverseLex { count: None }),
("diagonal", TraversalOrder::Diagonal { count: None }),
("antidiagonal", TraversalOrder::Antidiagonal { count: None }),
("extrema", TraversalOrder::Extrema { strata: None }),
("shells", TraversalOrder::Shells {
origin: ShellOrigin::Outer, depth: None }),
("halton", TraversalOrder::Halton { count: None }),
("sobol", TraversalOrder::Sobol { count: None }),
("lhs", TraversalOrder::Lhs { count: None, seed: None }),
] {
let comp = Comprehension::union(vec![
vec![Clause::new("k", "10")],
vec![Clause::new("k", "20")],
]).with_order(ord);
let errs = comp.validate().unwrap_err();
assert!(errs.iter().any(|e| e.contains(label)),
"{label}: error should name the strategy: {errs:?}");
}
}
#[test]
fn parse_clause_parallel_two_vars() {
let c = parse_clause("(x, y) in (1..10, 100..1000..100)").unwrap();
assert!(c.is_parallel());
assert_eq!(c.vars, vec!["x".to_string(), "y".to_string()]);
match &c.source {
super::super::ast::ClauseSource::Parallel { exprs, .. } => {
assert_eq!(exprs, &vec!["1..10".to_string(), "100..1000..100".to_string()]);
}
_ => panic!("expected Parallel source"),
}
}
#[test]
fn parse_clause_parallel_three_vars() {
let c = parse_clause("(a, b, c) in (1..3, 10..30..10, 100..300..100)").unwrap();
assert!(c.is_parallel());
assert_eq!(c.vars, vec!["a".to_string(), "b".to_string(), "c".to_string()]);
}
#[test]
fn parse_clause_parallel_with_function_call_rhs() {
let c = parse_clause("(x, y) in (fib(8), pow2(8))").unwrap();
assert!(c.is_parallel());
match &c.source {
super::super::ast::ClauseSource::Parallel { exprs, .. } => {
assert_eq!(exprs, &vec!["fib(8)".to_string(), "pow2(8)".to_string()]);
}
_ => panic!("expected Parallel source"),
}
}
#[test]
fn parse_clause_paren_only_one_side_is_rejected() {
let err = parse_clause("(x, y) in 1..10").unwrap_err();
assert!(err.contains("parentheses on both sides"), "got: {err}");
}
#[test]
fn parse_clause_parallel_count_mismatch_is_rejected() {
let err = parse_clause("(x, y, z) in (1..10, 1..20)").unwrap_err();
assert!(err.contains("3 variables but 2 expressions"), "got: {err}");
}
#[test]
fn parse_clause_parallel_single_var_is_rejected() {
let err = parse_clause("(x) in (1..10)").unwrap_err();
assert!(err.contains("≥ 2 variables"), "got: {err}");
}
#[test]
fn split_respects_parallel_clause_boundaries() {
let s = "(x, y) in (1..2, 10..20..10), z in 100..200..100";
let parts = split_respecting_parens(s);
assert_eq!(parts.len(), 2);
}
#[test]
fn parse_clause_list_mixed_parallel_and_single() {
let clauses = parse_clause_list(
"(x, y) in (1..2, 10..20..10), z in 100..200..100"
).unwrap();
assert_eq!(clauses.len(), 2);
assert!(clauses[0].is_parallel());
assert!(!clauses[1].is_parallel());
assert_eq!(clauses[1].var(), "z");
}
#[test]
fn parse_clause_parallel_invalid_var_name_is_rejected() {
let err = parse_clause("(x, 9b) in (1..10, 1..10)").unwrap_err();
assert!(err.contains("not a valid variable name"), "got: {err}");
}
fn roundtrip_clause(c: Clause) {
let text = c.to_string();
let reparsed = parse_clause(&text)
.unwrap_or_else(|e| panic!("re-parse failed for '{text}': {e}"));
assert_eq!(c, reparsed,
"round-trip diverged: original={c:?}\n text='{text}'\n reparsed={reparsed:?}");
}
#[test]
fn round_trip_single_var_clause() {
roundtrip_clause(Clause::new("k", "1..10"));
roundtrip_clause(Clause::new("limit", "fib(8)"));
}
#[test]
fn round_trip_parallel_strict() {
use super::super::ast::ZipMode;
roundtrip_clause(Clause::parallel(["x", "y"], ["fib(8)", "pow2(8)"]));
roundtrip_clause(Clause::parallel_with_mode(
ZipMode::Strict, ["a", "b", "c"], ["1..3", "10..30..10", "100..300..100"]
));
}
#[test]
fn round_trip_parallel_truncate_and_cycle() {
use super::super::ast::ZipMode;
roundtrip_clause(Clause::parallel_with_mode(
ZipMode::Truncate, ["x", "y"], ["fib(8)", "pow2(4)"]
));
roundtrip_clause(Clause::parallel_with_mode(
ZipMode::Cycle, ["x", "y"], ["fib(4)", "pow2(8)"]
));
}
fn roundtrip_comprehension_text(text: &str) {
let parsed = parse_comprehension_text(text)
.unwrap_or_else(|e| panic!("parse failed for '{text}': {e}"));
let rendered = parsed.to_string();
let reparsed = parse_comprehension_text(&rendered)
.unwrap_or_else(|e| panic!("re-parse failed for '{rendered}' (from '{text}'): {e}"));
assert_eq!(parsed, reparsed,
"round-trip diverged for '{text}':\n rendered='{rendered}'");
}
#[test]
fn round_trip_cartesian_comprehension() {
roundtrip_comprehension_text("k in 1..10");
roundtrip_comprehension_text("k in 1..10, limit in fib(8)");
roundtrip_comprehension_text("k in 1..10 where {k} > 3");
roundtrip_comprehension_text("k in 1..10 order extrema/2");
roundtrip_comprehension_text("k in 1..10, l in 1..10 where {k} != {l} order extrema/1");
}
#[test]
fn round_trip_parallel_iter_through_full_comprehension_text() {
roundtrip_comprehension_text("(x, y) in (fib(5), pow2(5))");
roundtrip_comprehension_text("(x, y) in zip_truncate(fib(8), pow2(4))");
roundtrip_comprehension_text(
"(x, y) in (fib(4), pow2(4)), z in 1..3 order extrema/1"
);
}
#[test]
fn parse_clause_parallel_zip_truncate_mode() {
use super::super::ast::{ClauseSource, ZipMode};
let c = parse_clause("(x, y) in zip_truncate(1..10, fib(8))").unwrap();
match &c.source {
ClauseSource::Parallel { mode, exprs } => {
assert_eq!(*mode, ZipMode::Truncate);
assert_eq!(exprs, &vec!["1..10".to_string(), "fib(8)".to_string()]);
}
_ => panic!("expected Parallel source, got {:?}", c.source),
}
}
#[test]
fn parse_clause_parallel_zip_cycle_mode() {
use super::super::ast::{ClauseSource, ZipMode};
let c = parse_clause("(x, y) in zip_cycle(1..10, 100..1000..100)").unwrap();
match &c.source {
ClauseSource::Parallel { mode, .. } => {
assert_eq!(*mode, ZipMode::Cycle);
}
_ => panic!("expected Parallel source"),
}
}
#[test]
fn parse_clause_parallel_default_mode_is_strict() {
use super::super::ast::{ClauseSource, ZipMode};
let c = parse_clause("(x, y) in (1..10, 100..1000..100)").unwrap();
match &c.source {
ClauseSource::Parallel { mode, .. } => {
assert_eq!(*mode, ZipMode::Strict);
}
_ => panic!("expected Parallel source"),
}
}
}