Skip to main content

slotted_egraphs/
lang.rs

1use crate::*;
2
3#[derive(Debug, Clone)]
4pub enum SyntaxElem {
5    String(String), // used for identitifers and payloads
6    AppliedId(AppliedId),
7    Slot(Slot),
8}
9
10pub trait LanguageChildren: Debug + Clone + Hash + Eq {
11    // TODO: add private_slot_occurrences aswell!
12    fn all_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot>;
13    fn public_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot>;
14    fn applied_id_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut AppliedId>;
15
16    fn all_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot>;
17    fn public_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot>;
18    fn applied_id_occurrences_iter(&self) -> impl Iterator<Item = &AppliedId>;
19
20    fn to_syntax(&self) -> Vec<SyntaxElem>;
21    fn from_syntax(_: &[SyntaxElem]) -> Option<Self>;
22
23    fn weak_shape_impl(&mut self, _m: &mut (SlotMap, u32)) {
24        todo!()
25    }
26}
27
28fn on_see_slot(s: &mut Slot, m: &mut (SlotMap, u32)) {
29    if let Some(s2) = m.0.get(*s) {
30        *s = s2;
31    } else {
32        add_slot(s, m);
33    }
34}
35
36fn add_slot(s: &mut Slot, m: &mut (SlotMap, u32)) {
37    let s2 = Slot::numeric(m.1);
38    m.1 += 1;
39    m.0.insert(*s, s2);
40    *s = s2;
41}
42
43impl LanguageChildren for AppliedId {
44    fn all_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
45        self.m.values_mut()
46    }
47    fn public_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
48        self.m.values_mut()
49    }
50    fn applied_id_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut AppliedId> {
51        std::iter::once(self)
52    }
53
54    fn all_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
55        self.m.values_immut()
56    }
57    fn public_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
58        self.m.values_immut()
59    }
60    fn applied_id_occurrences_iter(&self) -> impl Iterator<Item = &AppliedId> {
61        std::iter::once(self)
62    }
63
64    fn to_syntax(&self) -> Vec<SyntaxElem> {
65        vec![SyntaxElem::AppliedId(self.clone())]
66    }
67    fn from_syntax(elems: &[SyntaxElem]) -> Option<Self> {
68        match elems {
69            [SyntaxElem::AppliedId(x)] => Some(x.clone()),
70            _ => None,
71        }
72    }
73
74    fn weak_shape_impl(&mut self, m: &mut (SlotMap, u32)) {
75        for x in self.m.values_mut() {
76            on_see_slot(x, m);
77        }
78    }
79}
80
81impl LanguageChildren for Slot {
82    fn all_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
83        std::iter::once(self)
84    }
85    fn public_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
86        std::iter::once(self)
87    }
88    fn applied_id_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut AppliedId> {
89        std::iter::empty()
90    }
91
92    fn all_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
93        std::iter::once(self)
94    }
95    fn public_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
96        std::iter::once(self)
97    }
98    fn applied_id_occurrences_iter(&self) -> impl Iterator<Item = &AppliedId> {
99        std::iter::empty()
100    }
101
102    fn to_syntax(&self) -> Vec<SyntaxElem> {
103        vec![SyntaxElem::Slot(*self)]
104    }
105    fn from_syntax(elems: &[SyntaxElem]) -> Option<Self> {
106        match elems {
107            [SyntaxElem::Slot(x)] => Some(x.clone()),
108            _ => None,
109        }
110    }
111
112    fn weak_shape_impl(&mut self, m: &mut (SlotMap, u32)) {
113        on_see_slot(self, m);
114    }
115}
116
117/// Implements [LanguageChildren] for payload types that are independent of Slots. For example u32, String etc.
118#[macro_export]
119macro_rules! bare_language_child {
120    ($($id:ident),*) => {
121        $(
122        impl LanguageChildren for $id {
123            fn all_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item=&mut Slot> { std::iter::empty() }
124            fn public_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item=&mut Slot> { std::iter::empty() }
125            fn applied_id_occurrences_iter_mut(&mut self) -> impl Iterator<Item=&mut AppliedId> { std::iter::empty() }
126
127            fn all_slot_occurrences_iter(&self) -> impl Iterator<Item=&Slot> { std::iter::empty() }
128            fn public_slot_occurrences_iter(&self) -> impl Iterator<Item=&Slot> { std::iter::empty() }
129            fn applied_id_occurrences_iter(&self) -> impl Iterator<Item=&AppliedId> { std::iter::empty() }
130
131            fn to_syntax(&self) -> Vec<SyntaxElem> { vec![SyntaxElem::String(self.to_string())] }
132            fn from_syntax(elems: &[SyntaxElem]) -> Option<Self> {
133                match elems {
134                    [SyntaxElem::String(x)] => x.parse().ok(),
135                    _ => None,
136                }
137            }
138
139            fn weak_shape_impl(&mut self, _m: &mut (SlotMap, u32)) {}
140        }
141        )*
142    }
143}
144
145bare_language_child!(
146    u128, u64, u32, u16, u8, i128, i64, i32, i16, i8, usize, isize, bool, char, Symbol
147);
148
149#[derive(Hash, PartialEq, Eq, Debug, Clone, PartialOrd, Ord)]
150pub struct Bind<T> {
151    pub slot: Slot,
152    pub elem: T,
153}
154
155impl<L: LanguageChildren> LanguageChildren for Bind<L> {
156    // mut:
157    fn all_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
158        std::iter::once(&mut self.slot).chain(self.elem.all_slot_occurrences_iter_mut())
159    }
160
161    fn public_slot_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut Slot> {
162        self.elem
163            .public_slot_occurrences_iter_mut()
164            .filter(|x| **x != self.slot)
165    }
166
167    fn applied_id_occurrences_iter_mut(&mut self) -> impl Iterator<Item = &mut AppliedId> {
168        self.elem.applied_id_occurrences_iter_mut()
169    }
170
171    // immut:
172    fn all_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
173        std::iter::once(&self.slot).chain(self.elem.all_slot_occurrences_iter())
174    }
175
176    fn public_slot_occurrences_iter(&self) -> impl Iterator<Item = &Slot> {
177        self.elem
178            .public_slot_occurrences_iter()
179            .filter(|x| **x != self.slot)
180    }
181
182    fn applied_id_occurrences_iter(&self) -> impl Iterator<Item = &AppliedId> {
183        self.elem.applied_id_occurrences_iter()
184    }
185
186    // syntax:
187    fn to_syntax(&self) -> Vec<SyntaxElem> {
188        let mut v = vec![SyntaxElem::Slot(self.slot)];
189        v.extend(self.elem.to_syntax());
190
191        v
192    }
193
194    fn from_syntax(elems: &[SyntaxElem]) -> Option<Self> {
195        let SyntaxElem::Slot(slot) = elems.get(0)? else {
196            return None;
197        };
198        let elem = L::from_syntax(&elems[1..])?;
199
200        Some(Bind { slot: *slot, elem })
201    }
202
203    fn weak_shape_impl(&mut self, m: &mut (SlotMap, u32)) {
204        let s = self.slot;
205        add_slot(&mut self.slot, m);
206        self.elem.weak_shape_impl(m);
207        m.0.remove(s);
208    }
209}
210
211// TODO: add LanguageChildren definition for tuples.
212
213/// A trait to define your Language (i.e. your E-Node type).
214pub trait Language: Debug + Clone + Hash + Eq + Ord {
215    /// List the mutable references of all child [Slot]s in your E-Node, in order of occurrence.
216    fn all_slot_occurrences_mut(&mut self) -> Vec<&mut Slot>;
217
218    /// List the mutable references to all *public* child [Slot]s in your E-Node, in order of occurrence.
219    ///
220    /// Public Slots are those, which are visible from the outside of that e-node.
221    /// * A typical example would be a `(var $x)` e-node, which has a *public* slot `$x`.
222    /// * A typical counter-example would be the `(lam $x body)` e-node, which has a *private* slot `$x`.
223    fn public_slot_occurrences_mut(&mut self) -> Vec<&mut Slot>;
224
225    /// List the mutable references to all child [AppliedId]s in your E-Node, in the order of occurrence.
226    fn applied_id_occurrences_mut(&mut self) -> Vec<&mut AppliedId>;
227
228    fn all_slot_occurrences(&self) -> Vec<Slot>;
229    fn public_slot_occurrences(&self) -> Vec<Slot>;
230    fn applied_id_occurrences(&self) -> Vec<&AppliedId>;
231
232    /// This function will be used to display your E-Node.
233    fn to_syntax(&self) -> Vec<SyntaxElem>;
234
235    /// This function will be used to parse your E-Node.
236    fn from_syntax(_: &[SyntaxElem]) -> Option<Self>;
237
238    fn slots(&self) -> SmallHashSet<Slot>;
239    fn weak_shape_inplace(&mut self) -> Bijection;
240
241    #[track_caller]
242    #[doc(hidden)]
243    fn check(&self) {
244        let mut c = self.clone();
245        let all: HashSet<*mut Slot> = c
246            .all_slot_occurrences_mut()
247            .into_iter()
248            .map(|x| x as *mut Slot)
249            .collect();
250        let public: HashSet<*mut Slot> = c
251            .public_slot_occurrences_mut()
252            .into_iter()
253            .map(|x| x as *mut Slot)
254            .collect();
255        let private: HashSet<*mut Slot> = c
256            .private_slot_occurrences_mut()
257            .into_iter()
258            .map(|x| x as *mut Slot)
259            .collect();
260
261        assert!(public.is_disjoint(&private));
262
263        // This also catches errors, where different Slot-addresses have the same slot names. This also counts as a collision!
264        let f = |x: Vec<Slot>| x.into_iter().collect::<HashSet<_>>();
265        assert!(f(c.public_slot_occurrences()).is_disjoint(&f(c.private_slot_occurrences())));
266
267        let all2: HashSet<*mut Slot> = public.union(&private).copied().collect();
268        assert_eq!(all2, all);
269    }
270
271    // generated methods:
272
273    #[doc(hidden)]
274    fn private_slot_occurrences_mut(&mut self) -> Vec<&mut Slot> {
275        let public = self.public_slot_occurrences();
276        let mut out = self.all_slot_occurrences_mut();
277        out.retain(|x| !public.contains(x));
278        out
279    }
280
281    #[doc(hidden)]
282    fn private_slot_occurrences(&self) -> Vec<Slot> {
283        let public = self.public_slot_occurrences();
284        let mut out = self.all_slot_occurrences();
285        out.retain(|x| !public.contains(x));
286        out
287    }
288
289    #[doc(hidden)]
290    fn private_slots(&self) -> SmallHashSet<Slot> {
291        self.private_slot_occurrences().into_iter().collect()
292    }
293
294    #[doc(hidden)]
295    fn map_applied_ids(&self, mut f: impl FnMut(AppliedId) -> AppliedId) -> Self {
296        let mut c = self.clone();
297        for x in c.applied_id_occurrences_mut() {
298            *x = f(x.clone());
299        }
300        c
301    }
302
303    // TODO m.values() might collide with your private slot names.
304    // Should we rename our private slots to be safe?
305    #[doc(hidden)]
306    fn apply_slotmap_partial(&self, m: &SlotMap) -> Self {
307        let mut prv = vec![].into();
308        if CHECKS {
309            prv = self.private_slots();
310        }
311
312        let mut c = self.clone();
313        for x in c.public_slot_occurrences_mut() {
314            let y = m[*x];
315
316            // If y collides with a private slot, we have a problem.
317            if CHECKS {
318                assert!(!prv.contains(&y));
319            }
320
321            *x = y;
322        }
323        c
324    }
325
326    #[track_caller]
327    #[doc(hidden)]
328    fn apply_slotmap(&self, m: &SlotMap) -> Self {
329        if CHECKS {
330            assert!(
331                m.keys().is_superset(&self.slots()),
332                "Language::apply_slotmap: The SlotMap doesn't map all free slots!"
333            );
334        }
335        self.apply_slotmap_partial(m)
336    }
337
338    #[doc(hidden)]
339    fn apply_slotmap_fresh(&self, m: &SlotMap) -> Self {
340        let mut prv = vec![].into();
341        if CHECKS {
342            prv = self.private_slots();
343        }
344
345        let mut c = self.clone();
346        for x in c.public_slot_occurrences_mut() {
347            let y = m.get(*x).unwrap_or_else(Slot::fresh);
348
349            // If y collides with a private slot, we have a problem.
350            if CHECKS {
351                assert!(!prv.contains(&y));
352            }
353
354            *x = y;
355        }
356        c
357    }
358
359    #[doc(hidden)]
360    fn ids(&self) -> Vec<Id> {
361        self.applied_id_occurrences()
362            .into_iter()
363            .map(|x| x.id)
364            .collect()
365    }
366
367    // let n.weak_shape() = (sh, bij); then
368    // - sh.apply_slotmap(bij) is equivalent to n (excluding lambda variable renames)
369    // - bij.slots() == n.slots(). Note that these would also include redundant slots.
370    // - sh is the lexicographically lowest equivalent version of n, reachable by bijective renaming of slots (including redundant ones).
371    #[doc(hidden)]
372    fn weak_shape(&self) -> (Self, Bijection) {
373        let mut c = self.clone();
374        let bij = c.weak_shape_inplace();
375        (c, bij)
376    }
377
378    #[doc(hidden)]
379    fn refresh_private(&self) -> Self {
380        let mut c = self.clone();
381        let prv: SmallHashSet<Slot> = c.private_slot_occurrences().into_iter().collect();
382        let fresh = SlotMap::bijection_from_fresh_to(&prv).inverse();
383        for x in c.private_slot_occurrences_mut() {
384            *x = fresh[*x];
385        }
386        c
387    }
388
389    #[doc(hidden)]
390    fn refresh_slots(&self, set: SmallHashSet<Slot>) -> Self {
391        let mut c = self.clone();
392        let fresh = SlotMap::bijection_from_fresh_to(&set).inverse();
393        for x in c.all_slot_occurrences_mut() {
394            if set.contains(x) {
395                *x = fresh[*x];
396            }
397        }
398        c
399    }
400
401    // refreshes private and redundant slots.
402    // The public slots are given by `public`.
403    #[doc(hidden)]
404    fn refresh_internals(&self, public: SmallHashSet<Slot>) -> Self {
405        let mut c = self.clone();
406        let internals = &c
407            .all_slot_occurrences()
408            .into_iter()
409            .collect::<SmallHashSet<_>>()
410            - &public;
411        let fresh = SlotMap::bijection_from_fresh_to(&internals).inverse();
412        for x in c.all_slot_occurrences_mut() {
413            if internals.contains(x) {
414                *x = fresh[*x];
415            }
416        }
417        c
418    }
419}