feat: add Expr::Scope (#683)

Add `Scope` expression variant, allowing submodels to be defined inside a model.

Change the rewriter so that it rewrites expressions inside their
submodels, using that submodels symbol table. This allows rules to apply
to variables defined in a local scope.

This commit is part of on-going work to add lexical scope to Conjure
Oxide, following on from 542c89e (feat!: add `Submodel`, 2025-02-16).

Author: niklasdewally

crates/conjure_core/src/ast/expressions.rs

@@ -16,15 +16,15 @@ use crate::ast::Name;
 use crate::ast::ReturnType;
 use crate::metadata::Metadata;
 
-use super::{Domain, Range, SubModel};
+use super::{Domain, Range, SubModel, Typeable};
 
 /// Represents different types of expressions used to define rules and constraints in the model.
 ///
 /// The `Expression` enum includes operations, constants, and variable references
 /// used to build rules and conditions for the model.
 #[document_compatibility]
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Uniplate)]
-#[uniplate(walk_into=[Atom,Submodel])]
+#[uniplate(walk_into=[Atom,SubModel])]
 #[biplate(to=Literal)]
 #[biplate(to=Metadata)]
 #[biplate(to=Atom)]
@@ -41,6 +41,8 @@ pub enum Expression {
 
     Atomic(Metadata, Atom),
 
+    Scope(Metadata, Box<SubModel>),
+
     /// `|x|` - absolute value of `x`
     #[compatible(JsonInput)]
     Abs(Metadata, Box<Expression>),
@@ -334,6 +336,7 @@ impl Expression {
                 Some(Domain::IntDomain(vec![Range::Single(*n)]))
             }
             Expression::Atomic(_, Atom::Literal(Literal::Bool(_))) => Some(Domain::BoolDomain),
+            Expression::Scope(_, _) => Some(Domain::BoolDomain),
             Expression::Sum(_, exprs) => expr_vec_to_domain_i32(exprs, |x, y| Some(x + y), syms),
             Expression::Product(_, exprs) => {
                 expr_vec_to_domain_i32(exprs, |x, y| Some(x * y), syms)
@@ -512,6 +515,7 @@ impl Expression {
             Expression::Atomic(_, Atom::Literal(Literal::Int(_))) => Some(ReturnType::Int),
             Expression::Atomic(_, Atom::Literal(Literal::Bool(_))) => Some(ReturnType::Bool),
             Expression::Atomic(_, Atom::Reference(_)) => None,
+            Expression::Scope(_, scope) => scope.return_type(),
             Expression::Abs(_, _) => Some(ReturnType::Int),
             Expression::Sum(_, _) => Some(ReturnType::Int),
             Expression::Product(_, _) => Some(ReturnType::Int),
@@ -619,6 +623,7 @@ impl Display for Expression {
                 write!(f, "{}", pretty_expressions_as_top_level(exprs))
             }
             Expression::Atomic(_, atom) => atom.fmt(f),
+            Expression::Scope(_, submodel) => write!(f, "{{\n{submodel}\n}}"),
             Expression::Abs(_, a) => write!(f, "|{}|", a),
             Expression::Sum(_, expressions) => {
                 write!(f, "Sum({})", pretty_vec(expressions))

crates/conjure_core/src/ast/model.rs

@@ -1,4 +1,7 @@
+use std::cell::RefCell;
+use std::collections::HashMap;
 use std::fmt::{Debug, Display};
+use std::rc::Rc;
 use std::sync::{Arc, RwLock};
 
 use derivative::Derivative;
@@ -8,9 +11,10 @@ use uniplate::{Biplate, Tree, Uniplate};
 use crate::ast::Expression;
 use crate::context::Context;
 
+use super::serde::{HasId, ObjId};
 use super::types::Typeable;
-use super::ReturnType;
 use super::SubModel;
+use super::{ReturnType, SymbolTable};
 
 /// An Essence model.
 ///
@@ -88,6 +92,25 @@ impl Biplate<Expression> for Model {
     }
 }
 
+impl Biplate<SubModel> for Model {
+    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
+        let submodel = self.as_submodel().clone();
+
+        let self2 = self.clone();
+        let ctx = Box::new(move |x| {
+            let Tree::One(submodel) = x else {
+                panic!();
+            };
+
+            let mut self3 = self2.clone();
+            self3.replace_submodel(submodel);
+            self3
+        });
+
+        (Tree::One(submodel), ctx)
+    }
+}
+
 impl Display for Model {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         std::fmt::Display::fmt(self.as_submodel(), f)
@@ -108,10 +131,40 @@ pub struct SerdeModel {
 impl SerdeModel {
     /// Initialises the model for rewriting.
     pub fn initialise(self, context: Arc<RwLock<Context<'static>>>) -> Option<Model> {
-        // TODO: Once we have submodels and multiple symbol tables, de-duplicate deserialized
-        // Rc<RefCell<>> symbol tables and declarations using their stored ids.
+        // The definitive versions of each symbol table are stored in the submodels. Parent
+        // pointers store dummy values with the correct ids, but nothing else. We need to replace
+        // these dummy values with pointers to the actual parent symbol tables, using the ids to
+        // know which tables should be equal.
         //
-        // See ast::serde::RcRefCellAsId.
+        // See super::serde::RcRefCellToInner, super::serde::RcRefCellToId.
+
+        // Store the definitive versions of all symbol tables by id.
+        let mut tables: HashMap<ObjId, Rc<RefCell<SymbolTable>>> = HashMap::new();
+
+        // Find the definitive versions by traversing the sub-models.
+        for submodel in self.submodel.universe() {
+            let id = submodel.symbols().id();
+
+            // ids should be unique!
+            assert_eq!(
+                tables.insert(id, submodel.symbols_ptr_unchecked().clone()),
+                None
+            );
+        }
+
+        // Restore parent pointers using `tables`.
+        for table in tables.clone().into_values() {
+            let mut table_mut = table.borrow_mut();
+            let parent_mut = table_mut.parent_mut_unchecked();
+
+            #[allow(clippy::unwrap_used)]
+            if let Some(parent) = parent_mut {
+                let parent_id = parent.borrow().id();
+
+                *parent = tables.get(&parent_id).unwrap().clone();
+            }
+        }
+
         Some(Model {
             submodel: self.submodel,
             context,

crates/conjure_core/src/ast/submodel.rs

@@ -271,3 +271,17 @@ impl Biplate<Expression> for SubModel {
         (tree, ctx)
     }
 }
+
+impl Biplate<SubModel> for SubModel {
+    fn biplate(&self) -> (Tree<SubModel>, Box<dyn Fn(Tree<SubModel>) -> Self>) {
+        (
+            Tree::One(self.clone()),
+            Box::new(move |x| {
+                let Tree::One(x) = x else {
+                    panic!();
+                };
+                x
+            }),
+        )
+    }
+}

crates/conjure_core/src/ast/symbol_table.rs

@@ -202,6 +202,13 @@ impl SymbolTable {
         *(self.next_machine_name.borrow_mut()) += 1;
         Name::MachineName(num) // incremented when inserted
     }
+
+    /// Gets the parent of this symbol table as a mutable reference.
+    ///
+    /// This function provides no sanity checks.
+    pub fn parent_mut_unchecked(&mut self) -> &mut Option<Rc<RefCell<SymbolTable>>> {
+        &mut self.parent
+    }
 }
 
 impl IntoIterator for SymbolTable {

crates/conjure_core/src/rule_engine/mod.rs

@@ -69,6 +69,7 @@ pub use rewrite_naive::rewrite_naive;
 pub use rewriter_common::RewriteError;
 pub use rule::{ApplicationError, ApplicationResult, Reduction, Rule};
 pub use rule_set::RuleSet;
+mod submodel_zipper;
 
 use crate::solver::SolverFamily;
 

crates/conjure_core/src/rule_engine/rewrite.rs

@@ -129,7 +129,7 @@ pub fn rewrite_model<'a>(
         ) {
             debug_assert!(is_vec_bool(&step.new_top)); // All new_top expressions should be boolean
             new_model.as_submodel_mut().constraints_mut()[i] = step.new_expression.clone();
-            step.apply(&mut new_model); // Apply side-effects (e.g., symbol table updates)
+            step.apply(new_model.as_submodel_mut()); // Apply side-effects (e.g., symbol table updates)
         }
 
         // If new constraints are added, continue processing them in the next iterations.
@@ -220,7 +220,7 @@ fn rewrite_iteration(
     let rule_results = apply_all_rules(&expression, model, rules, stats);
     if let Some(result) = choose_rewrite(&rule_results, &expression) {
         // If a rule is applied, mark the expression as dirty
-        log_rule_application(&result, &expression, model);
+        log_rule_application(&result, &expression, model.as_submodel());
         return Some(result.reduction);
     }
 

crates/conjure_core/src/rule_engine/rewrite_naive.rs

@@ -1,10 +1,11 @@
 use super::{resolve_rules::RuleData, RewriteError, RuleSet};
 use crate::{
-    ast::Expression as Expr,
+    ast::{Expression as Expr, SubModel},
     bug,
     rule_engine::{
         get_rules_grouped,
         rewriter_common::{log_rule_application, RuleResult},
+        submodel_zipper::submodel_ctx,
     },
     Model,
 };
@@ -26,11 +27,31 @@ pub fn rewrite_naive<'a>(
         .collect_vec();
 
     let mut model = model.clone();
+    let mut done_something = true;
 
     // Rewrite until there are no more rules left to apply.
-    while let Some(()) =
-        try_rewrite_model(&mut model, &rules_grouped, prop_multiple_equally_applicable)
-    {}
+    while done_something {
+        let mut new_model = None;
+        done_something = false;
+
+        // Rewrite each sub-model in the tree, largest first.
+        for (mut submodel, ctx) in <_ as Biplate<SubModel>>::contexts_bi(&model) {
+            if try_rewrite_model(
+                &mut submodel,
+                &rules_grouped,
+                prop_multiple_equally_applicable,
+            )
+            .is_some()
+            {
+                new_model = Some(ctx(submodel));
+                done_something = true;
+                break;
+            }
+        }
+        if let Some(new_model) = new_model {
+            model = new_model;
+        }
+    }
 
     Ok(model)
 }
@@ -39,23 +60,23 @@ pub fn rewrite_naive<'a>(
 //
 // Returns None if no change was made.
 fn try_rewrite_model(
-    model: &mut Model,
+    submodel: &mut SubModel,
     rules_grouped: &Vec<(u16, Vec<RuleData<'_>>)>,
     prop_multiple_equally_applicable: bool,
 ) -> Option<()> {
-    type CtxFn = Arc<dyn Fn(Expr) -> Model>;
+    type CtxFn = Arc<dyn Fn(Expr) -> SubModel>;
     let mut results: Vec<(RuleResult<'_>, u16, Expr, CtxFn)> = vec![];
 
     // Iterate over rules by priority in descending order.
     'top: for (priority, rules) in rules_grouped.iter() {
         // Using Biplate, rewrite both the expression tree, and any value lettings in the symbol
         // table.
-        for (expr, ctx) in <_ as Biplate<Expr>>::contexts_bi(model) {
+        for (expr, ctx) in submodel_ctx(submodel.clone()) {
             // Clone expr and ctx so they can be reused
             let expr = expr.clone();
             let ctx = ctx.clone();
             for rd in rules {
-                match (rd.rule.application)(&expr, &model.as_submodel().symbols()) {
+                match (rd.rule.application)(&expr, &submodel.symbols()) {
                     Ok(red) => {
                         // Collect applicable rules
                         results.push((
@@ -99,15 +120,15 @@ fn try_rewrite_model(
             }
 
             // Extract the single applicable rule and apply it
-            log_rule_application(result, expr, model);
+            log_rule_application(result, expr, submodel);
 
             // Replace expr with new_expression
-            *model = ctx(result.reduction.new_expression.clone());
+            *submodel = ctx(result.reduction.new_expression.clone());
 
             // Apply new symbols and top level
-            result.reduction.clone().apply(model);
+            result.reduction.clone().apply(submodel);
 
-            println!("{}", &model);
+            println!("{}", &submodel);
         }
     }
 

crates/conjure_core/src/rule_engine/rewriter_common.rs

@@ -3,12 +3,9 @@ use super::{
     resolve_rules::{ResolveRulesError, RuleData},
     Reduction,
 };
-use crate::{
-    ast::{
-        pretty::{pretty_variable_declaration, pretty_vec},
-        Expression,
-    },
-    Model,
+use crate::ast::{
+    pretty::{pretty_variable_declaration, pretty_vec},
+    Expression, SubModel,
 };
 
 use itertools::Itertools;
@@ -28,7 +25,7 @@ pub struct RuleResult<'a> {
 pub fn log_rule_application(
     result: &RuleResult,
     initial_expression: &Expression,
-    initial_model: &Model,
+    initial_model: &SubModel,
 ) {
     let red = &result.reduction;
     let rule = result.rule_data.rule;
@@ -63,7 +60,7 @@ pub fn log_rule_application(
     let new_variables_str = {
         let mut vars: Vec<String> = vec![];
 
-        for var_name in red.added_symbols(&initial_model.as_submodel().symbols()) {
+        for var_name in red.added_symbols(&initial_model.symbols()) {
             #[allow(clippy::unwrap_used)]
             vars.push(format!(
                 "  {}",

crates/conjure_core/src/rule_engine/rule.rs

@@ -6,7 +6,7 @@ use std::rc::Rc;
 use thiserror::Error;
 
 use crate::ast::Declaration;
-use crate::ast::{Expression, Model, Name, SymbolTable};
+use crate::ast::{Expression, Name, SubModel, SymbolTable};
 
 #[derive(Debug, Error)]
 pub enum ApplicationError {
@@ -100,10 +100,9 @@ impl Reduction {
     }
 
     /// Applies side-effects (e.g. symbol table updates)
-    pub fn apply(self, model: &mut Model) {
-        let submodel = model.as_submodel_mut();
-        submodel.symbols_mut().extend(self.symbols); // Add new assignments to the symbol table
-        submodel.add_constraints(self.new_top.clone());
+    pub fn apply(self, model: &mut SubModel) {
+        model.symbols_mut().extend(self.symbols); // Add new assignments to the symbol table
+        model.add_constraints(self.new_top.clone());
     }
 
     /// Gets symbols added by this reduction