Skip to main content

riddle/
env.rs

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