conjure_cp_core/ast/
comprehension.rs

1#![allow(clippy::arc_with_non_send_sync)]
2
3use std::{collections::BTreeSet, fmt::Display};
4
5use crate::bug_assert;
6use conjure_cp_core::ast::ReturnType;
7use itertools::Itertools as _;
8use parking_lot::RwLockReadGuard;
9use serde::{Deserialize, Serialize};
10use serde_with::serde_as;
11use uniplate::{Biplate, Uniplate};
12
13use super::{
14    DeclarationPtr, Domain, DomainPtr, Expression, Model, Name, Range, SymbolTable, SymbolTablePtr,
15    Typeable,
16    ac_operators::ACOperatorKind,
17    categories::{Category, CategoryOf},
18    serde::{AsId, PtrAsInner},
19};
20
21#[serde_as]
22#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug, Uniplate)]
23#[biplate(to=Expression)]
24#[biplate(to=Name)]
25#[biplate(to=DeclarationPtr)]
26pub enum ComprehensionQualifier {
27    ExpressionGenerator {
28        #[serde_as(as = "AsId")]
29        ptr: DeclarationPtr,
30    },
31    Generator {
32        #[serde_as(as = "AsId")]
33        ptr: DeclarationPtr,
34    },
35    Condition(Expression),
36}
37
38/// A comprehension.
39#[serde_as]
40#[derive(Clone, PartialEq, Eq, Hash, Uniplate, Serialize, Deserialize, Debug)]
41#[biplate(to=Expression)]
42#[biplate(to=SymbolTable)]
43#[biplate(to=SymbolTablePtr)]
44#[non_exhaustive]
45pub struct Comprehension {
46    pub return_expression: Expression,
47    pub qualifiers: Vec<ComprehensionQualifier>,
48    /// When this comprehension appears inside an AC operator, records which operator so
49    /// expansion can apply the correct skip semantics for symbolic guards.
50    #[serde(default)]
51    pub skip_operator: Option<ACOperatorKind>,
52    #[doc(hidden)]
53    #[serde_as(as = "PtrAsInner")]
54    pub symbols: SymbolTablePtr,
55}
56
57impl Comprehension {
58    pub fn domain_of(&self) -> Option<DomainPtr> {
59        let return_expr_domain = self.return_expression.domain_of()?;
60
61        // return a list (matrix with index domain int(1..)) of return_expr elements
62        Some(Domain::matrix(
63            return_expr_domain,
64            vec![Domain::int(vec![Range::UnboundedR(1)])],
65        ))
66    }
67
68    pub fn return_expression(self) -> Expression {
69        self.return_expression
70    }
71
72    pub fn replace_return_expression(&mut self, new_expr: Expression) {
73        self.return_expression = new_expr;
74    }
75
76    pub fn symbols(&self) -> RwLockReadGuard<'_, SymbolTable> {
77        self.symbols.read()
78    }
79
80    pub fn quantified_vars(&self) -> Vec<Name> {
81        self.qualifiers
82            .iter()
83            .filter_map(|q| match q {
84                ComprehensionQualifier::ExpressionGenerator { ptr } => Some(ptr.name().clone()),
85                ComprehensionQualifier::Generator { ptr } => Some(ptr.name().clone()),
86                ComprehensionQualifier::Condition(_) => None,
87            })
88            .collect()
89    }
90
91    pub fn generator_conditions(&self) -> Vec<Expression> {
92        self.qualifiers
93            .iter()
94            .filter_map(|q| match q {
95                ComprehensionQualifier::Condition(c) => Some(c.clone()),
96                ComprehensionQualifier::Generator { .. } => None,
97                ComprehensionQualifier::ExpressionGenerator { .. } => None,
98            })
99            .collect()
100    }
101
102    /// Builds a temporary model containing generator qualifiers and guards.
103    pub fn to_generator_model(&self) -> Model {
104        let mut model = self.empty_model_with_symbols();
105        model.add_constraints(self.generator_conditions());
106        model
107    }
108
109    /// Builds a temporary model containing the return expression only.
110    pub fn to_return_expression_model(&self) -> Model {
111        let mut model = self.empty_model_with_symbols();
112        model.add_constraint(self.return_expression.clone());
113        model
114    }
115
116    fn empty_model_with_symbols(&self) -> Model {
117        let parent = self.symbols.read().parent().clone();
118        let mut model = if let Some(parent) = parent {
119            Model::new_in_parent_scope(parent)
120        } else {
121            Model::default()
122        };
123        *model.symbols_ptr_unchecked_mut() = self.symbols.clone();
124        model
125    }
126
127    /// Adds a guard to the comprehension.
128    ///
129    /// Returns false if the guard references non-quantified decision variables.
130    pub fn add_quantified_guard(&mut self, guard: Expression) -> bool {
131        if self.is_quantified_guard(&guard) {
132            self.qualifiers
133                .push(ComprehensionQualifier::Condition(guard));
134            true
135        } else {
136            false
137        }
138    }
139
140    /// True iff expr does not reference non-quantified decision variables.
141    pub fn is_quantified_guard(&self, expr: &Expression) -> bool {
142        let quantified: BTreeSet<Name> = self.quantified_vars().into_iter().collect();
143        is_quantified_guard(&self.symbols.read(), &quantified, expr)
144    }
145}
146
147impl Typeable for Comprehension {
148    /// A comprehension is a collection, so its type is a matrix of the return expression's type.
149    ///
150    /// This agrees with [`Comprehension::domain_of`], which reports a matrix domain. Reporting the
151    /// element type here instead made every "is this a collection?" test quietly miss
152    /// comprehensions.
153    fn return_type(&self) -> ReturnType {
154        ReturnType::Matrix(Box::new(self.return_expression.return_type()))
155    }
156}
157
158impl Display for Comprehension {
159    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160        let generators_and_guards = self
161            .qualifiers
162            .iter()
163            .map(|qualifier| match qualifier {
164                ComprehensionQualifier::Generator { ptr } => {
165                    let domain = ptr.domain().expect("generator declaration has domain");
166                    format!("{} : {domain}", ptr.name())
167                }
168                ComprehensionQualifier::ExpressionGenerator { ptr } => {
169                    let name = ptr.name();
170                    if let Some(expr) = ptr.as_quantified_expr() {
171                        format!("{name} <- {expr}")
172                    } else {
173                        panic!("Oh nein! Dat is nicht gut!")
174                    }
175                }
176                ComprehensionQualifier::Condition(expr) => format!("{expr}"),
177            })
178            .join(", ");
179
180        write!(
181            f,
182            "[ {} | {generators_and_guards} ]",
183            self.return_expression
184        )
185    }
186}
187
188/// A builder for a comprehension.
189#[derive(Clone, Debug, PartialEq, Eq)]
190pub struct ComprehensionBuilder {
191    qualifiers: Vec<ComprehensionQualifier>,
192    // A single scope for generators and return expression.
193    symbols: SymbolTablePtr,
194    quantified_variables: BTreeSet<Name>,
195}
196
197impl ComprehensionBuilder {
198    pub fn new(symbol_table_ptr: SymbolTablePtr) -> Self {
199        ComprehensionBuilder {
200            qualifiers: vec![],
201            symbols: SymbolTablePtr::with_parent(symbol_table_ptr),
202            quantified_variables: BTreeSet::new(),
203        }
204    }
205
206    /// Backwards-compatible parser API: same table for generators and return expression.
207    pub fn generator_symboltable(&mut self) -> SymbolTablePtr {
208        self.symbols.clone()
209    }
210
211    /// Backwards-compatible parser API: same table for generators and return expression.
212    pub fn return_expr_symboltable(&mut self) -> SymbolTablePtr {
213        self.symbols.clone()
214    }
215
216    pub fn guard(mut self, guard: Expression) -> Self {
217        self.qualifiers
218            .push(ComprehensionQualifier::Condition(guard));
219        self
220    }
221
222    pub fn generator(mut self, declaration: DeclarationPtr) -> Self {
223        let name = declaration.name().clone();
224        bug_assert!(!self.quantified_variables.contains(&name));
225
226        self.quantified_variables.insert(name.clone());
227
228        // insert into comprehension scope as a local quantified variable
229        let quantified_decl = DeclarationPtr::new_quantified(name, declaration.domain().unwrap());
230        self.symbols.write().insert(quantified_decl.clone());
231
232        self.qualifiers.push(ComprehensionQualifier::Generator {
233            ptr: quantified_decl,
234        });
235
236        self
237    }
238
239    pub fn expression_generator(mut self, name: Name, expr: Expression) -> Self {
240        bug_assert!(!self.quantified_variables.contains(&name));
241
242        self.quantified_variables.insert(name.clone());
243
244        // insert into comprehension scope as a local quantified variable
245        let quantified_decl = DeclarationPtr::new_quantified_expr(name, expr);
246        self.symbols.write().insert(quantified_decl.clone());
247
248        self.qualifiers
249            .push(ComprehensionQualifier::ExpressionGenerator {
250                ptr: quantified_decl,
251            });
252
253        self
254    }
255
256    /// Creates a comprehension with the given return expression.
257    ///
258    /// Guards are always stored as [`ComprehensionQualifier::Condition`] entries. When a guard
259    /// references non-quantified decision variables, the enclosing AC operator applies the
260    /// appropriate skip semantics during comprehension expansion.
261    pub fn with_return_value(self, expression: Expression) -> Comprehension {
262        Comprehension {
263            return_expression: expression,
264            qualifiers: self.qualifiers,
265            skip_operator: None,
266            symbols: self.symbols,
267        }
268    }
269}
270
271/// True iff the guard does not reference non-quantified decision variables.
272fn is_quantified_guard(
273    symbols: &SymbolTable,
274    quantified_variables: &BTreeSet<Name>,
275    guard: &Expression,
276) -> bool {
277    guard.universe_bi().iter().all(|name| {
278        quantified_variables.contains(name)
279            || symbols
280                .lookup(name)
281                .is_some_and(|decl| decl.category_of() != Category::Decision)
282    })
283}
conjure_cp_core/ast/comprehension.rs

conjure_cp_core/ast/
comprehension.rs
