Lines
60.76 %
Functions
30 %
use crate::RecoverableParseError;
use crate::errors::FatalParseError;
use crate::expression::parse_expression;
use crate::field;
use crate::parser::ParseContext;
use crate::parser::domain::parse_domain;
use crate::util::named_children;
use conjure_cp_core::ast::ac_operators::ACOperatorKind;
use conjure_cp_core::ast::comprehension::ComprehensionBuilder;
use conjure_cp_core::ast::{DeclarationPtr, Expression, Metadata, Moo, Name};
use std::vec;
use tree_sitter::Node;
pub fn parse_comprehension(
ctx: &mut ParseContext,
node: &Node,
) -> Result<Option<Expression>, FatalParseError> {
// If we're in a set context, add error and return early since comprehensions don't produce sets
if ctx.typechecking_context == crate::util::TypecheckingContext::Set {
ctx.record_error(crate::errors::RecoverableParseError::new(
format!(
"Type error: {}\n\tExpected: set\n\tGot: comprehension",
ctx.source_code[node.start_byte()..node.end_byte()].trim()
),
Some(node.range()),
));
}
// Comprehensions require a symbol table passed in
let symbols_ptr = match ctx.symbols.clone() {
Some(s) => s,
None => {
ctx.record_error(RecoverableParseError::new(
"Comprehensions require a symbol table".to_string(),
return Ok(None);
};
let mut builder = ComprehensionBuilder::new(symbols_ptr);
// We need to track the return expression node separately since it appears first in syntax
// but we need to parse generators first (to get variables in scope)
let mut return_expr_node: Option<Node> = None;
// set return expression node and parse generators/conditions
for child in named_children(node) {
match child.kind() {
"arithmetic_expr" | "bool_expr" | "comparison_expr" | "atom" => {
// Store the return expression node to parse later
return_expr_node = Some(child);
"generator" => {
// Parse the generator variable
let Some(var_node) = field!(recover, ctx, child, "variable") else {
let var_name_str = &ctx.source_code[var_node.start_byte()..var_node.end_byte()];
let var_name = Name::user(var_name_str);
// Parse the domain
let Some(domain_node) = field!(recover, ctx, child, "domain") else {
// Parse with a new context using the generator symbol table
let mut domain_ctx = ctx.with_new_symbols(Some(builder.generator_symboltable()));
let Some(var_domain) = parse_domain(&mut domain_ctx, domain_node)? else {
// Add generator using the builder
let decl = DeclarationPtr::new_find(var_name, var_domain);
builder = builder.generator(decl);
"condition" => {
// Parse the condition expression
let Some(expr_node) = field!(recover, ctx, child, "expression") else {
let generator_symboltable = builder.generator_symboltable();
let mut guard_ctx = ctx.with_new_symbols(Some(generator_symboltable));
let Some(guard_expr) = parse_expression(&mut guard_ctx, expr_node)? else {
// Add the condition as a guard
builder = builder.guard(guard_expr);
_ => {
// Skip other nodes (like punctuation)
// parse the return expression
let return_expr_node = match return_expr_node {
Some(node) => node,
"Comprehension missing return expression".to_string(),
// Use the return expression symbol table which already has quantified variables (as Given) and parent as parent
let mut return_ctx = ctx.with_new_symbols(Some(builder.return_expr_symboltable()));
let Some(return_expr) = parse_expression(&mut return_ctx, return_expr_node)? else {
// Build the comprehension with the return expression and default ACOperatorKind::And
let comprehension = builder.with_return_value(return_expr, Some(ACOperatorKind::And));
Ok(Some(Expression::Comprehension(
Metadata::new(),
Moo::new(comprehension),
)))
/// Parse comprehension-style expressions
/// - `forAll vars : domain . expr` → `And(Comprehension(...))`
/// - `sum vars : domain . expr` → `Sum(Comprehension(...))`
pub fn parse_quantifier_or_aggregate_expr(
// Quantifier and aggregate expressions require a symbol table
"Quantifier and aggregate expressions require a symbol table".to_string(),
// Create the comprehension builder
// First pass: collect domain/collection, variables
let mut domain = None;
let mut collection_node = None;
let mut variables = vec![];
"identifier" => {
let var_name_str = &ctx.source_code[child.start_byte()..child.end_byte()];
variables.push(var_name);
"domain" => {
// Parse with the current symbol table (no need for a new context)
let Some(parsed_domain) = parse_domain(ctx, child)? else {
domain = Some(parsed_domain);
"set_literal" | "matrix" | "tuple" | "record" => {
// Store the collection node to parse later
collection_node = Some(child);
_ => continue,
// We need either a domain or a collection
if domain.is_none() && collection_node.is_none() {
"Quantifier and aggregate expressions require a domain or collection".to_string(),
if variables.is_empty() {
"Quantifier and aggregate expressions require variables".to_string(),
// Get the operator type
let Some(operator_node) = field!(recover, ctx, node, "operator") else {
let operator_str = &ctx.source_code[operator_node.start_byte()..operator_node.end_byte()];
let (ac_operator_kind, wrapper) = match operator_str {
"forAll" => (ACOperatorKind::And, "And"),
"exists" => (ACOperatorKind::Or, "Or"),
"sum" => (ACOperatorKind::Sum, "Sum"),
"min" => (ACOperatorKind::Sum, "Min"), // AC operator doesn't matter for non-boolean aggregates
"max" => (ACOperatorKind::Sum, "Max"),
format!("Unknown operator: {}", operator_str),
Some(operator_node.range()),
// Add variables as generators
if let Some(dom) = domain {
for var_name in variables {
let decl = DeclarationPtr::new_find(var_name, dom.clone());
} else if let Some(_coll_node) = collection_node {
// TODO: support collection domains
"Collection domains in quantifier and aggregate expressions".to_string(),
Some(_coll_node.range()),
// Parse the expression (after variables are in the symbol table)
let Some(expression_node) = field!(recover, ctx, node, "expression") else {
// Parse with a new context using the return expression symbol table
let mut expr_ctx = ctx.with_new_symbols(Some(builder.return_expr_symboltable()));
let Some(expression) = parse_expression(&mut expr_ctx, expression_node)? else {
// Build the comprehension
let comprehension = builder.with_return_value(expression, Some(ac_operator_kind));
let wrapped_comprehension = Expression::Comprehension(Metadata::new(), Moo::new(comprehension));
// Wrap in the appropriate expression type
match wrapper {
"And" => Ok(Some(Expression::And(
Moo::new(wrapped_comprehension),
))),
"Or" => Ok(Some(Expression::Or(
"Sum" => Ok(Some(Expression::Sum(
"Min" => Ok(Some(Expression::Min(
"Max" => Ok(Some(Expression::Max(
_ => unreachable!(),