Skip to main content

riddle/
env.rs

1use crate::{
2    RiddleError,
3    core::Core,
4    scope::{Class, Predicate, Scope, Type},
5};
6use core::fmt;
7use std::{
8    any::Any,
9    cell::RefCell,
10    collections::HashMap,
11    ops::Deref,
12    rc::{Rc, Weak},
13};
14
15#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
16pub struct ObjectId(pub(super) usize);
17
18impl From<usize> for ObjectId {
19    fn from(val: usize) -> Self {
20        ObjectId(val)
21    }
22}
23
24impl Deref for ObjectId {
25    type Target = usize;
26
27    fn deref(&self) -> &Self::Target {
28        &self.0
29    }
30}
31
32impl fmt::Display for ObjectId {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        write!(f, "obj-{}", self.0)
35    }
36}
37
38#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
39pub struct AtomId(pub(super) usize);
40
41impl Deref for AtomId {
42    type Target = usize;
43
44    fn deref(&self) -> &Self::Target {
45        &self.0
46    }
47}
48
49impl fmt::Display for AtomId {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
51        write!(f, "atm-{}", self.0)
52    }
53}
54
55#[derive(Clone)]
56pub enum Slot {
57    Primitive(Rc<dyn Var>),
58    ObjectRef(ObjectId),
59    AtomRef(AtomId),
60}
61
62impl fmt::Display for Slot {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            Slot::Primitive(var) => write!(f, "{}", var.var_type().name()),
66            Slot::ObjectRef(obj_id) => write!(f, "Object({})", *obj_id),
67            Slot::AtomRef(atom_id) => write!(f, "Atom({})", *atom_id),
68        }
69    }
70}
71
72pub trait Var {
73    fn var_type(&self) -> Rc<dyn Type>;
74    fn as_any(self: Rc<Self>) -> Rc<dyn Any>;
75    fn as_env(self: Rc<Self>) -> Option<Rc<dyn Env>> {
76        None
77    }
78}
79
80pub trait Env {
81    fn parent(&self) -> Option<Rc<dyn Env>>;
82    fn get(&self, name: &str) -> Option<Slot>;
83    fn set(&self, name: String, value: Slot);
84}
85
86pub struct CommonEnv {
87    parent: Option<Rc<dyn Env>>,
88    variables: RefCell<HashMap<String, Slot>>,
89}
90
91impl CommonEnv {
92    pub fn new(parent: Option<Rc<dyn Env>>) -> Self {
93        Self { parent, variables: RefCell::new(HashMap::new()) }
94    }
95}
96
97impl Env for CommonEnv {
98    fn parent(&self) -> Option<Rc<dyn Env>> {
99        self.parent.clone()
100    }
101
102    fn get(&self, name: &str) -> Option<Slot> {
103        self.variables.borrow().get(name).cloned().or_else(|| self.parent.as_ref()?.get(name))
104    }
105
106    fn set(&self, name: String, value: Slot) {
107        self.variables.borrow_mut().insert(name, value);
108    }
109}
110
111pub enum BoolExpr {
112    Term { var_type: Weak<dyn Type>, term: Slot },
113    Not { var_type: Weak<dyn Type>, term: Rc<BoolExpr> },
114    Eq { var_type: Weak<dyn Type>, left: Slot, right: Slot },
115    Lt { var_type: Weak<dyn Type>, left: Slot, right: Slot },
116    Leq { var_type: Weak<dyn Type>, left: Slot, right: Slot },
117    Or { var_type: Weak<dyn Type>, terms: Vec<Rc<BoolExpr>> },
118    And { var_type: Weak<dyn Type>, terms: Vec<Rc<BoolExpr>> },
119}
120
121impl Var for BoolExpr {
122    fn var_type(&self) -> Rc<dyn Type> {
123        match self {
124            BoolExpr::Term { var_type: var_tp, .. } | BoolExpr::Not { var_type: var_tp, .. } | BoolExpr::Eq { var_type: var_tp, .. } | BoolExpr::Lt { var_type: var_tp, .. } | BoolExpr::Leq { var_type: var_tp, .. } | BoolExpr::Or { var_type: var_tp, .. } | BoolExpr::And { var_type: var_tp, .. } => var_tp.upgrade().unwrap(),
125        }
126    }
127
128    fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
129        self
130    }
131}
132
133pub struct Object {
134    id: ObjectId,
135    class: Weak<dyn Class>,
136    env: CommonEnv,
137}
138
139impl Object {
140    pub(super) fn new(id: ObjectId, class: Rc<dyn Class>) -> Self {
141        Self { id, class: Rc::downgrade(&class), env: CommonEnv::new(None) }
142    }
143
144    pub fn id(&self) -> ObjectId {
145        self.id
146    }
147
148    pub fn class(&self) -> Rc<dyn Class> {
149        self.class.upgrade().unwrap()
150    }
151}
152
153impl Var for Object {
154    fn var_type(&self) -> Rc<dyn Type> {
155        self.class.upgrade().unwrap()
156    }
157
158    fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
159        self
160    }
161
162    fn as_env(self: Rc<Self>) -> Option<Rc<dyn Env>> {
163        Some(self.clone())
164    }
165}
166
167impl Env for Object {
168    fn parent(&self) -> Option<Rc<dyn Env>> {
169        self.env.parent.clone()
170    }
171
172    fn get(&self, name: &str) -> Option<Slot> {
173        self.env.get(name)
174    }
175
176    fn set(&self, name: String, value: Slot) {
177        self.env.set(name, value);
178    }
179}
180
181pub struct Atom {
182    id: AtomId,
183    predicate: Weak<Predicate>,
184    fact: bool,
185    env: CommonEnv,
186}
187
188impl Atom {
189    pub fn new(id: AtomId, predicate: Rc<Predicate>, fact: bool, args: HashMap<String, Slot>) -> Self {
190        let env = match args.get("tau") {
191            Some(tau) => match tau {
192                Slot::Primitive(var) => var.clone().as_env().expect("Tau variable does not have an environment").clone(),
193                Slot::ObjectRef(obj_id) => predicate.clone().core().get_object(*obj_id).expect("Object ID in tau does not exist").as_env().expect("Object in tau does not have an environment").clone(),
194                Slot::AtomRef(atom_id) => predicate.clone().core().get_atom(*atom_id).expect("Atom ID in tau does not exist").as_env().expect("Atom in tau does not have an environment").clone(),
195            },
196            None => predicate.clone().core(),
197        };
198        let env = CommonEnv::new(Some(env));
199        for (name, value) in args {
200            env.set(name, value);
201        }
202        Self { id, predicate: Rc::downgrade(&predicate), fact, env }
203    }
204
205    pub fn id(&self) -> AtomId {
206        self.id
207    }
208
209    pub fn predicate(&self) -> Rc<Predicate> {
210        self.predicate.upgrade().unwrap()
211    }
212
213    pub fn is_fact(&self) -> bool {
214        self.fact
215    }
216}
217
218impl Var for Atom {
219    fn var_type(&self) -> Rc<dyn Type> {
220        self.predicate.upgrade().unwrap()
221    }
222
223    fn as_any(self: Rc<Self>) -> Rc<dyn Any> {
224        self
225    }
226
227    fn as_env(self: Rc<Self>) -> Option<Rc<dyn Env>> {
228        Some(self.clone())
229    }
230}
231
232impl Env for Atom {
233    fn parent(&self) -> Option<Rc<dyn Env>> {
234        self.env.parent.clone()
235    }
236
237    fn get(&self, name: &str) -> Option<Slot> {
238        self.env.get(name)
239    }
240
241    fn set(&self, name: String, value: Slot) {
242        self.env.set(name, value);
243    }
244}
245
246fn push_negations(expr: Rc<BoolExpr>) -> Rc<BoolExpr> {
247    match expr.as_ref() {
248        BoolExpr::Not { term, .. } => push_inverted(term.clone()),
249        BoolExpr::And { var_type, terms } => Rc::new(BoolExpr::And {
250            var_type: var_type.clone(),
251            terms: terms.iter().map(|t| push_negations(t.clone())).collect(),
252        }),
253        BoolExpr::Or { var_type, terms } => Rc::new(BoolExpr::Or {
254            var_type: var_type.clone(),
255            terms: terms.iter().map(|t| push_negations(t.clone())).collect(),
256        }),
257        _ => expr,
258    }
259}
260
261/// Processes an expression as if a `Not` wrapper were applied to it.
262fn push_inverted(expr: Rc<BoolExpr>) -> Rc<BoolExpr> {
263    match expr.as_ref() {
264        // Double Negation: Not(Not(term)) => term
265        BoolExpr::Not { term, .. } => push_negations(term.clone()),
266
267        // De Morgan: Not(And(A, B)) => Or(Not(A), Not(B))
268        BoolExpr::And { var_type, terms } => Rc::new(BoolExpr::Or {
269            var_type: var_type.clone(),
270            terms: terms.iter().map(|t| push_inverted(t.clone())).collect(),
271        }),
272
273        // De Morgan: Not(Or(A, B)) => And(Not(A), Not(B))
274        BoolExpr::Or { var_type, terms } => Rc::new(BoolExpr::And {
275            var_type: var_type.clone(),
276            terms: terms.iter().map(|t| push_inverted(t.clone())).collect(),
277        }),
278
279        BoolExpr::Leq { var_type, left, right } => Rc::new(BoolExpr::Lt { var_type: var_type.clone(), left: right.clone(), right: left.clone() }),
280        BoolExpr::Lt { var_type, left, right } => Rc::new(BoolExpr::Leq { var_type: var_type.clone(), left: right.clone(), right: left.clone() }),
281
282        BoolExpr::Term { var_type: var_tp, .. } | BoolExpr::Eq { var_type: var_tp, .. } => Rc::new(BoolExpr::Not { var_type: var_tp.clone(), term: expr }),
283    }
284}
285
286fn distribute(expr: Rc<BoolExpr>) -> Rc<BoolExpr> {
287    match expr.as_ref() {
288        BoolExpr::Or { var_type, terms } => {
289            // Step 1: Recursively distribute child nodes and flatten any nested Ors
290            let mut distributed_terms = Vec::new();
291            for t in terms {
292                let dist = distribute(t.clone());
293                if let BoolExpr::Or { terms: inner_terms, .. } = dist.as_ref() {
294                    distributed_terms.extend(inner_terms.clone());
295                } else {
296                    distributed_terms.push(dist);
297                }
298            }
299
300            // Step 2: Build the Cartesian product of terms over And boundaries
301            // Start with a pool containing a single empty clause
302            let mut result_ands: Vec<Vec<Rc<BoolExpr>>> = vec![vec![]];
303
304            for term in distributed_terms {
305                if let BoolExpr::And { terms: and_terms, .. } = term.as_ref() {
306                    // Split all existing combinations across the newly encountered And choices
307                    let mut next_ands = Vec::new();
308                    for existing_and in &result_ands {
309                        for and_term in and_terms {
310                            let mut combo = existing_and.clone();
311                            combo.push(and_term.clone());
312                            next_ands.push(combo);
313                        }
314                    }
315                    result_ands = next_ands;
316                } else {
317                    // Leaf nodes or Or nodes get appended to all current paths
318                    for existing_and in &mut result_ands {
319                        existing_and.push(term.clone());
320                    }
321                }
322            }
323
324            // Step 3: Map our combinations back into Or nodes inside a master And node
325            let cnf_or_nodes: Vec<Rc<BoolExpr>> = result_ands.into_iter().map(|sub_terms| Rc::new(BoolExpr::Or { var_type: var_type.clone(), terms: sub_terms })).collect();
326
327            // Optimization: If no distribution happened, don't wrap in a redundant And
328            if cnf_or_nodes.len() == 1 { cnf_or_nodes[0].clone() } else { Rc::new(BoolExpr::And { var_type: var_type.clone(), terms: cnf_or_nodes }) }
329        }
330
331        BoolExpr::And { var_type, terms } => {
332            // Flatten nested Ands to keep the AST compact
333            let mut distributed_terms = Vec::new();
334            for t in terms {
335                let dist = distribute(t.clone());
336                if let BoolExpr::And { terms: inner_terms, .. } = dist.as_ref() {
337                    distributed_terms.extend(inner_terms.clone());
338                } else {
339                    distributed_terms.push(dist);
340                }
341            }
342            Rc::new(BoolExpr::And { var_type: var_type.clone(), terms: distributed_terms })
343        }
344
345        _ => expr,
346    }
347}
348
349pub fn to_cnf(expr: Rc<BoolExpr>) -> Rc<BoolExpr> {
350    distribute(push_negations(expr))
351}
352
353pub fn get_var_by_path(core: &dyn Core, env: &dyn Env, path: &[String]) -> Result<Slot, RiddleError> {
354    let (first, rest) = path.split_first().ok_or_else(|| RiddleError::RuntimeError("Empty variable path".into()))?;
355    rest.iter().try_fold(env.get(first).ok_or_else(|| RiddleError::NotFound(first.to_string()))?, |acc, id| match acc {
356        Slot::Primitive(var) => var.clone().as_env().ok_or_else(|| RiddleError::NotAnEnvironment(format!("Variable '{}' in path does not have an environment", first)))?.get(id).ok_or_else(|| RiddleError::NotFound(format!("Variable '{}' in path not found in variable '{}'", id, first))),
357        Slot::ObjectRef(obj_id) => {
358            let obj = core.get_object(obj_id).ok_or_else(|| RiddleError::NotFound(format!("Object {} not found", *obj_id)))?;
359            obj.as_env().ok_or_else(|| RiddleError::NotAnEnvironment(format!("Object {} does not have an environment", *obj_id)))?.get(id).ok_or_else(|| RiddleError::NotFound(format!("Variable '{}' in path not found in object {}", id, *obj_id)))
360        }
361        Slot::AtomRef(atom_id) => {
362            let atom = core.get_atom(atom_id).ok_or_else(|| RiddleError::NotFound(format!("Atom {} not found", *atom_id)))?;
363            atom.as_env().ok_or_else(|| RiddleError::NotAnEnvironment(format!("Atom {} does not have an environment", *atom_id)))?.get(id).ok_or_else(|| RiddleError::NotFound(format!("Variable '{}' in path not found in atom {}", id, *atom_id)))
364        }
365    })
366}