egglog_ast/
generic_ast_helpers.rs

1use std::borrow::Cow;
2use std::fmt::{Display, Formatter};
3use std::hash::Hash;
4
5use ordered_float::OrderedFloat;
6
7use super::util::ListDisplay;
8use crate::generic_ast::*;
9use crate::span::Span;
10
11// Macro to implement From conversions for Literal types
12macro_rules! impl_from {
13    ($ctor:ident($t:ty)) => {
14        impl From<Literal> for $t {
15            fn from(literal: Literal) -> Self {
16                match literal {
17                    Literal::$ctor(t) => t,
18                    #[allow(unreachable_patterns)]
19                    _ => panic!("Expected {}, got {literal}", stringify!($ctor)),
20                }
21            }
22        }
23
24        impl From<$t> for Literal {
25            fn from(t: $t) -> Self {
26                Literal::$ctor(t)
27            }
28        }
29    };
30}
31
32pub const INTERNAL_SYMBOL_PREFIX: &str = "@";
33
34/// Gets rid of internal symbol prefixes for printing.
35/// This allows us to test parsing of desugared programs.
36pub fn sanitize_internal_name(name: &str) -> Cow<'_, str> {
37    if let Some(stripped) = name.strip_prefix(INTERNAL_SYMBOL_PREFIX) {
38        Cow::Owned(format!("_{}", stripped))
39    } else {
40        Cow::Borrowed(name)
41    }
42}
43
44impl<Head: Display, Leaf: Display> Display for GenericRule<Head, Leaf>
45where
46    Head: Clone + Display,
47    Leaf: Clone + PartialEq + Eq + Display + Hash,
48{
49    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
50        let indent = " ".repeat(7);
51        write!(f, "(rule (")?;
52        for (i, fact) in self.body.iter().enumerate() {
53            if i > 0 {
54                write!(f, "{}", indent)?;
55            }
56
57            if i != self.body.len() - 1 {
58                writeln!(f, "{}", fact)?;
59            } else {
60                write!(f, "{}", fact)?;
61            }
62        }
63        write!(f, ")\n      (")?;
64        for (i, action) in self.head.0.iter().enumerate() {
65            if i > 0 {
66                write!(f, "{}", indent)?;
67            }
68            if i != self.head.0.len() - 1 {
69                writeln!(f, "{}", action)?;
70            } else {
71                write!(f, "{}", action)?;
72            }
73        }
74        let ruleset = if !self.ruleset.is_empty() {
75            format!(":ruleset {}", sanitize_internal_name(&self.ruleset))
76        } else {
77            "".into()
78        };
79        let name = if !self.name.is_empty() {
80            format!(":name \"{}\"", sanitize_internal_name(&self.name))
81        } else {
82            "".into()
83        };
84        write!(f, ")\n{} {} {})", indent, ruleset, name)
85    }
86}
87
88// Use the macro for Int, Float, and String conversions
89impl_from!(Int(i64));
90impl_from!(Float(OrderedFloat<f64>));
91impl_from!(String(String));
92
93impl<Head: Display, Leaf: Display> Display for GenericFact<Head, Leaf> {
94    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
95        match self {
96            GenericFact::Eq(_, e1, e2) => write!(f, "(= {e1} {e2})"),
97            GenericFact::Fact(expr) => write!(f, "{expr}"),
98        }
99    }
100}
101
102// Implement Display for GenericAction
103impl<Head: Display, Leaf: Display> Display for GenericAction<Head, Leaf>
104where
105    Head: Clone + Display,
106    Leaf: Clone + PartialEq + Eq + Display + Hash,
107{
108    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
109        match self {
110            GenericAction::Let(_, lhs, rhs) => write!(f, "(let {} {})", lhs, rhs),
111            GenericAction::Set(_, lhs, args, rhs) => {
112                if args.is_empty() {
113                    write!(f, "(set ({}) {})", lhs, rhs)
114                } else {
115                    write!(
116                        f,
117                        "(set ({} {}) {})",
118                        lhs,
119                        args.iter()
120                            .map(|a| format!("{}", a))
121                            .collect::<Vec<_>>()
122                            .join(" "),
123                        rhs
124                    )
125                }
126            }
127            GenericAction::Union(_, lhs, rhs) => write!(f, "(union {} {})", lhs, rhs),
128            GenericAction::Change(_, change, lhs, args) => {
129                let change_str = match change {
130                    Change::Delete => "delete",
131                    Change::Subsume => "subsume",
132                };
133                write!(
134                    f,
135                    "({} ({} {}))",
136                    change_str,
137                    lhs,
138                    args.iter()
139                        .map(|a| format!("{}", a))
140                        .collect::<Vec<_>>()
141                        .join(" ")
142                )
143            }
144            GenericAction::Panic(_, msg) => write!(f, "(panic \"{}\")", msg),
145            GenericAction::Expr(_, e) => write!(f, "{}", e),
146        }
147    }
148}
149
150impl<Head, Leaf> Display for GenericExpr<Head, Leaf>
151where
152    Head: Display,
153    Leaf: Display,
154{
155    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
156        match self {
157            GenericExpr::Lit(_ann, lit) => write!(f, "{lit}"),
158            GenericExpr::Var(_ann, var) => write!(f, "{var}"),
159            GenericExpr::Call(_ann, op, children) => {
160                write!(f, "({} {})", op, ListDisplay(children, " "))
161            }
162        }
163    }
164}
165
166impl<Head, Leaf> Default for GenericActions<Head, Leaf>
167where
168    Head: Clone + Display,
169    Leaf: Clone + PartialEq + Eq + Display + Hash,
170{
171    fn default() -> Self {
172        Self(vec![])
173    }
174}
175
176impl<Head, Leaf> GenericRule<Head, Leaf>
177where
178    Head: Clone + Display,
179    Leaf: Clone + PartialEq + Eq + Display + Hash,
180{
181    /// Applies `f` to every expression that appears in the rule body or head.
182    pub fn visit_exprs(
183        self,
184        f: &mut impl FnMut(GenericExpr<Head, Leaf>) -> GenericExpr<Head, Leaf>,
185    ) -> Self {
186        Self {
187            span: self.span,
188            head: self.head.visit_exprs(f),
189            body: self
190                .body
191                .into_iter()
192                .map(|bexpr| bexpr.visit_exprs(f))
193                .collect(),
194            name: self.name.clone(),
195            ruleset: self.ruleset.clone(),
196        }
197    }
198
199    /// Applies `f` to each action in the rule head, leaving the body unchanged.
200    pub fn visit_actions(
201        self,
202        f: &mut impl FnMut(GenericAction<Head, Leaf>) -> GenericAction<Head, Leaf>,
203    ) -> Self {
204        Self {
205            span: self.span,
206            head: self.head.visit_actions(f),
207            body: self.body,
208            name: self.name,
209            ruleset: self.ruleset,
210        }
211    }
212
213    /// Applies the provided `head` and `leaf` mappings to every symbol that appears in the rule.
214    pub fn map_symbols<Head2, Leaf2>(
215        self,
216        head: &mut impl FnMut(Head) -> Head2,
217        leaf: &mut impl FnMut(Leaf) -> Leaf2,
218    ) -> GenericRule<Head2, Leaf2>
219    where
220        Head2: Clone + Display,
221        Leaf2: Clone + PartialEq + Eq + Display + Hash,
222    {
223        GenericRule {
224            span: self.span,
225            head: self.head.map_symbols(head, leaf),
226            body: self
227                .body
228                .into_iter()
229                .map(|fact| fact.map_symbols(head, leaf))
230                .collect(),
231            name: self.name,
232            ruleset: self.ruleset,
233        }
234    }
235
236    /// Converts the rule into its unresolved representation by formatting heads and leaves.
237    pub fn make_unresolved(self) -> GenericRule<String, String> {
238        let mut map_head = |h: Head| h.to_string();
239        let mut map_leaf = |l: Leaf| l.to_string();
240        self.map_symbols(&mut map_head, &mut map_leaf)
241    }
242}
243
244impl<Head, Leaf> GenericActions<Head, Leaf>
245where
246    Head: Clone + Display,
247    Leaf: Clone + PartialEq + Eq + Display + Hash,
248{
249    pub fn len(&self) -> usize {
250        self.0.len()
251    }
252
253    pub fn is_empty(&self) -> bool {
254        self.0.is_empty()
255    }
256
257    pub fn iter(&self) -> impl Iterator<Item = &GenericAction<Head, Leaf>> {
258        self.0.iter()
259    }
260
261    pub fn visit_vars(&self, f: &mut impl FnMut(&Span, &Leaf)) {
262        for action in &self.0 {
263            action.visit_vars(f);
264        }
265    }
266
267    /// Transforms every expression appearing in the action list using `f`.
268    pub fn visit_exprs(
269        self,
270        f: &mut impl FnMut(GenericExpr<Head, Leaf>) -> GenericExpr<Head, Leaf>,
271    ) -> Self {
272        Self(self.0.into_iter().map(|a| a.visit_exprs(f)).collect())
273    }
274
275    /// Rewrites each action in the collection with the provided closure.
276    pub fn visit_actions(
277        self,
278        f: &mut impl FnMut(GenericAction<Head, Leaf>) -> GenericAction<Head, Leaf>,
279    ) -> Self {
280        Self(self.0.into_iter().map(f).collect())
281    }
282
283    pub fn new(actions: Vec<GenericAction<Head, Leaf>>) -> Self {
284        Self(actions)
285    }
286
287    pub fn singleton(action: GenericAction<Head, Leaf>) -> Self {
288        Self(vec![action])
289    }
290
291    /// Applies the provided `head` and `leaf` mappings to each action.
292    pub fn map_symbols<Head2, Leaf2>(
293        self,
294        head: &mut impl FnMut(Head) -> Head2,
295        leaf: &mut impl FnMut(Leaf) -> Leaf2,
296    ) -> GenericActions<Head2, Leaf2>
297    where
298        Head2: Clone + Display,
299        Leaf2: Clone + PartialEq + Eq + Display + Hash,
300    {
301        GenericActions(
302            self.0
303                .into_iter()
304                .map(|action| action.map_symbols(head, leaf))
305                .collect(),
306        )
307    }
308
309    /// Converts the actions into their unresolved representation by formatting heads and leaves.
310    pub fn make_unresolved(self) -> GenericActions<String, String> {
311        let mut map_head = |h: Head| h.to_string();
312        let mut map_leaf = |l: Leaf| l.to_string();
313        self.map_symbols(&mut map_head, &mut map_leaf)
314    }
315}
316
317impl<Head, Leaf> GenericAction<Head, Leaf>
318where
319    Head: Clone + Display,
320    Leaf: Clone + Eq + Display + Hash,
321{
322    pub fn visit_vars(&self, f: &mut impl FnMut(&Span, &Leaf)) {
323        if let GenericAction::Let(span, lhs, _) = self {
324            f(span, lhs);
325        }
326        let mut visit = |expr: GenericExpr<Head, Leaf>| match expr {
327            GenericExpr::Var(span, var) => {
328                f(&span, &var);
329                GenericExpr::Var(span, var)
330            }
331            other => other,
332        };
333        let _ = self.clone().visit_exprs(&mut visit);
334    }
335
336    // Applys `f` to all expressions in the action.
337    pub fn map_exprs(
338        &self,
339        f: &mut impl FnMut(&GenericExpr<Head, Leaf>) -> GenericExpr<Head, Leaf>,
340    ) -> Self {
341        match self {
342            GenericAction::Let(span, lhs, rhs) => {
343                GenericAction::Let(span.clone(), lhs.clone(), f(rhs))
344            }
345            GenericAction::Set(span, lhs, args, rhs) => {
346                let right = f(rhs);
347                GenericAction::Set(
348                    span.clone(),
349                    lhs.clone(),
350                    args.iter().map(f).collect(),
351                    right,
352                )
353            }
354            GenericAction::Change(span, change, lhs, args) => GenericAction::Change(
355                span.clone(),
356                *change,
357                lhs.clone(),
358                args.iter().map(f).collect(),
359            ),
360            GenericAction::Union(span, lhs, rhs) => {
361                GenericAction::Union(span.clone(), f(lhs), f(rhs))
362            }
363            GenericAction::Panic(span, msg) => GenericAction::Panic(span.clone(), msg.clone()),
364            GenericAction::Expr(span, e) => GenericAction::Expr(span.clone(), f(e)),
365        }
366    }
367
368    /// Applys `f` to all sub-expressions (including `self`)
369    /// bottom-up, collecting the results.
370    pub fn visit_exprs(
371        self,
372        f: &mut impl FnMut(GenericExpr<Head, Leaf>) -> GenericExpr<Head, Leaf>,
373    ) -> Self {
374        match self {
375            GenericAction::Let(span, lhs, rhs) => {
376                GenericAction::Let(span, lhs.clone(), rhs.visit_exprs(f))
377            }
378            // TODO should we refactor `Set` so that we can map over Expr::Call(lhs, args)?
379            // This seems more natural to oflatt
380            // Currently, visit_exprs does not apply f to the first argument of Set.
381            GenericAction::Set(span, lhs, args, rhs) => {
382                let args = args.into_iter().map(|e| e.visit_exprs(f)).collect();
383                GenericAction::Set(span, lhs.clone(), args, rhs.visit_exprs(f))
384            }
385            GenericAction::Change(span, change, lhs, args) => {
386                let args = args.into_iter().map(|e| e.visit_exprs(f)).collect();
387                GenericAction::Change(span, change, lhs.clone(), args)
388            }
389            GenericAction::Union(span, lhs, rhs) => {
390                GenericAction::Union(span, lhs.visit_exprs(f), rhs.visit_exprs(f))
391            }
392            GenericAction::Panic(span, msg) => GenericAction::Panic(span, msg.clone()),
393            GenericAction::Expr(span, e) => GenericAction::Expr(span, e.visit_exprs(f)),
394        }
395    }
396
397    pub fn subst(&self, subst: &mut impl FnMut(&Span, &Leaf) -> GenericExpr<Head, Leaf>) -> Self {
398        self.map_exprs(&mut |e| e.subst_leaf(subst))
399    }
400
401    pub fn map_def_use(self, fvar: &mut impl FnMut(Leaf, bool) -> Leaf) -> Self {
402        macro_rules! fvar_expr {
403            () => {
404                |span, s: _| GenericExpr::Var(span.clone(), fvar(s.clone(), false))
405            };
406        }
407        match self {
408            GenericAction::Let(span, lhs, rhs) => {
409                let lhs = fvar(lhs, true);
410                let rhs = rhs.subst_leaf(&mut fvar_expr!());
411                GenericAction::Let(span, lhs, rhs)
412            }
413            GenericAction::Set(span, lhs, args, rhs) => {
414                let args = args
415                    .into_iter()
416                    .map(|e| e.subst_leaf(&mut fvar_expr!()))
417                    .collect();
418                let rhs = rhs.subst_leaf(&mut fvar_expr!());
419                GenericAction::Set(span, lhs.clone(), args, rhs)
420            }
421            GenericAction::Change(span, change, lhs, args) => {
422                let args = args
423                    .into_iter()
424                    .map(|e| e.subst_leaf(&mut fvar_expr!()))
425                    .collect();
426                GenericAction::Change(span, change, lhs.clone(), args)
427            }
428            GenericAction::Union(span, lhs, rhs) => {
429                let lhs = lhs.subst_leaf(&mut fvar_expr!());
430                let rhs = rhs.subst_leaf(&mut fvar_expr!());
431                GenericAction::Union(span, lhs, rhs)
432            }
433            GenericAction::Panic(span, msg) => GenericAction::Panic(span, msg.clone()),
434            GenericAction::Expr(span, e) => {
435                GenericAction::Expr(span, e.subst_leaf(&mut fvar_expr!()))
436            }
437        }
438    }
439
440    /// Applies the provided `head` and `leaf` mappings to the action and all nested expressions.
441    pub fn map_symbols<Head2, Leaf2>(
442        self,
443        head: &mut impl FnMut(Head) -> Head2,
444        leaf: &mut impl FnMut(Leaf) -> Leaf2,
445    ) -> GenericAction<Head2, Leaf2>
446    where
447        Head2: Clone + Display,
448        Leaf2: Clone + Eq + Display + Hash,
449    {
450        match self {
451            GenericAction::Let(span, lhs, rhs) => {
452                GenericAction::Let(span, leaf(lhs), rhs.map_symbols(head, leaf))
453            }
454            GenericAction::Set(span, head_sym, args, rhs) => {
455                let mut mapped_args = Vec::with_capacity(args.len());
456                for arg in args {
457                    mapped_args.push(arg.map_symbols(head, leaf));
458                }
459                GenericAction::Set(
460                    span,
461                    head(head_sym),
462                    mapped_args,
463                    rhs.map_symbols(head, leaf),
464                )
465            }
466            GenericAction::Change(span, change, head_sym, args) => {
467                let mut mapped_args = Vec::with_capacity(args.len());
468                for arg in args {
469                    mapped_args.push(arg.map_symbols(head, leaf));
470                }
471                GenericAction::Change(span, change, head(head_sym), mapped_args)
472            }
473            GenericAction::Union(span, lhs, rhs) => GenericAction::Union(
474                span,
475                lhs.map_symbols(head, leaf),
476                rhs.map_symbols(head, leaf),
477            ),
478            GenericAction::Panic(span, msg) => GenericAction::Panic(span, msg),
479            GenericAction::Expr(span, expr) => {
480                GenericAction::Expr(span, expr.map_symbols(head, leaf))
481            }
482        }
483    }
484
485    /// Converts the action into its unresolved representation using String by
486    /// formatting heads and leaves.
487    pub fn make_unresolved(self) -> GenericAction<String, String> {
488        let mut map_head = |h: Head| h.to_string();
489        let mut map_leaf = |l: Leaf| l.to_string();
490        self.map_symbols(&mut map_head, &mut map_leaf)
491    }
492}
493
494impl<Head, Leaf> GenericFact<Head, Leaf>
495where
496    Head: Clone + Display,
497    Leaf: Clone + PartialEq + Eq + Display + Hash,
498{
499    pub fn visit_vars(&self, f: &mut impl FnMut(&Span, &Leaf)) {
500        let mut visit = |expr: GenericExpr<Head, Leaf>| match expr {
501            GenericExpr::Var(span, var) => {
502                f(&span, &var);
503                GenericExpr::Var(span, var)
504            }
505            other => other,
506        };
507        let _ = self.clone().visit_exprs(&mut visit);
508    }
509
510    pub fn visit_exprs(
511        self,
512        f: &mut impl FnMut(GenericExpr<Head, Leaf>) -> GenericExpr<Head, Leaf>,
513    ) -> GenericFact<Head, Leaf> {
514        match self {
515            GenericFact::Eq(span, e1, e2) => {
516                GenericFact::Eq(span, e1.visit_exprs(f), e2.visit_exprs(f))
517            }
518            GenericFact::Fact(expr) => GenericFact::Fact(expr.visit_exprs(f)),
519        }
520    }
521
522    pub fn map_exprs<Head2, Leaf2>(
523        &self,
524        f: &mut impl FnMut(&GenericExpr<Head, Leaf>) -> GenericExpr<Head2, Leaf2>,
525    ) -> GenericFact<Head2, Leaf2> {
526        match self {
527            GenericFact::Eq(span, e1, e2) => GenericFact::Eq(span.clone(), f(e1), f(e2)),
528            GenericFact::Fact(expr) => GenericFact::Fact(f(expr)),
529        }
530    }
531
532    pub fn subst<Leaf2, Head2>(
533        &self,
534        subst_leaf: &mut impl FnMut(&Span, &Leaf) -> GenericExpr<Head2, Leaf2>,
535        subst_head: &mut impl FnMut(&Head) -> Head2,
536    ) -> GenericFact<Head2, Leaf2> {
537        self.map_exprs(&mut |e| e.subst(subst_leaf, subst_head))
538    }
539}
540
541impl<Head, Leaf> GenericFact<Head, Leaf>
542where
543    Leaf: Clone + PartialEq + Eq + Display + Hash,
544    Head: Clone + Display,
545{
546    /// Applies the provided `head` and `leaf` mappings to the fact.
547    pub fn map_symbols<Head2, Leaf2>(
548        self,
549        head: &mut impl FnMut(Head) -> Head2,
550        leaf: &mut impl FnMut(Leaf) -> Leaf2,
551    ) -> GenericFact<Head2, Leaf2>
552    where
553        Head2: Clone + Display,
554        Leaf2: Clone + PartialEq + Eq + Display + Hash,
555    {
556        match self {
557            GenericFact::Eq(span, e1, e2) => {
558                GenericFact::Eq(span, e1.map_symbols(head, leaf), e2.map_symbols(head, leaf))
559            }
560            GenericFact::Fact(expr) => GenericFact::Fact(expr.map_symbols(head, leaf)),
561        }
562    }
563
564    /// Converts all heads and leaves to strings.
565    pub fn make_unresolved(self) -> GenericFact<String, String> {
566        let mut map_head = |h: Head| h.to_string();
567        let mut map_leaf = |l: Leaf| l.to_string();
568        self.map_symbols(&mut map_head, &mut map_leaf)
569    }
570}
571
572impl<Head: Clone + Display, Leaf: Hash + Clone + Display + Eq> GenericExpr<Head, Leaf> {
573    pub fn visit_vars(&self, f: &mut impl FnMut(&Span, &Leaf)) {
574        let mut visit = |expr: GenericExpr<Head, Leaf>| match expr {
575            GenericExpr::Var(span, var) => {
576                f(&span, &var);
577                GenericExpr::Var(span, var)
578            }
579            other => other,
580        };
581        let _ = self.clone().visit_exprs(&mut visit);
582    }
583
584    pub fn span(&self) -> Span {
585        match self {
586            GenericExpr::Lit(span, _) => span.clone(),
587            GenericExpr::Var(span, _) => span.clone(),
588            GenericExpr::Call(span, _, _) => span.clone(),
589        }
590    }
591
592    pub fn is_var(&self) -> bool {
593        matches!(self, GenericExpr::Var(_, _))
594    }
595
596    pub fn get_var(&self) -> Option<Leaf> {
597        match self {
598            GenericExpr::Var(_ann, v) => Some(v.clone()),
599            _ => None,
600        }
601    }
602
603    fn children(&self) -> &[Self] {
604        match self {
605            GenericExpr::Var(_, _) | GenericExpr::Lit(_, _) => &[],
606            GenericExpr::Call(_, _, children) => children,
607        }
608    }
609
610    pub fn ast_size(&self) -> usize {
611        let mut size = 0;
612        self.walk(&mut |_e| size += 1, &mut |_| {});
613        size
614    }
615
616    pub fn walk(&self, pre: &mut impl FnMut(&Self), post: &mut impl FnMut(&Self)) {
617        pre(self);
618        self.children()
619            .iter()
620            .for_each(|child| child.walk(pre, post));
621        post(self);
622    }
623
624    pub fn fold<Out>(&self, f: &mut impl FnMut(&Self, Vec<Out>) -> Out) -> Out {
625        let ts = self.children().iter().map(|child| child.fold(f)).collect();
626        f(self, ts)
627    }
628
629    /// Applys `f` to all sub-expressions (including `self`)
630    /// bottom-up, collecting the results.
631    pub fn visit_exprs(self, f: &mut impl FnMut(Self) -> Self) -> Self {
632        match self {
633            GenericExpr::Lit(..) => f(self),
634            GenericExpr::Var(..) => f(self),
635            GenericExpr::Call(span, op, children) => {
636                let children = children.into_iter().map(|c| c.visit_exprs(f)).collect();
637                f(GenericExpr::Call(span, op.clone(), children))
638            }
639        }
640    }
641
642    /// `subst` replaces occurrences of variables and head symbols in the expression.
643    pub fn subst<Head2, Leaf2>(
644        &self,
645        subst_leaf: &mut impl FnMut(&Span, &Leaf) -> GenericExpr<Head2, Leaf2>,
646        subst_head: &mut impl FnMut(&Head) -> Head2,
647    ) -> GenericExpr<Head2, Leaf2> {
648        match self {
649            GenericExpr::Lit(span, lit) => GenericExpr::Lit(span.clone(), lit.clone()),
650            GenericExpr::Var(span, v) => subst_leaf(span, v),
651            GenericExpr::Call(span, op, children) => {
652                let children = children
653                    .iter()
654                    .map(|c| c.subst(subst_leaf, subst_head))
655                    .collect();
656                GenericExpr::Call(span.clone(), subst_head(op), children)
657            }
658        }
659    }
660
661    pub fn subst_leaf<Leaf2>(
662        &self,
663        subst_leaf: &mut impl FnMut(&Span, &Leaf) -> GenericExpr<Head, Leaf2>,
664    ) -> GenericExpr<Head, Leaf2> {
665        self.subst(subst_leaf, &mut |x| x.clone())
666    }
667
668    /// Applies the provided `head` and `leaf` mappings to every symbol within the expression.
669    pub fn map_symbols<Head2, Leaf2>(
670        self,
671        head: &mut impl FnMut(Head) -> Head2,
672        leaf: &mut impl FnMut(Leaf) -> Leaf2,
673    ) -> GenericExpr<Head2, Leaf2> {
674        match self {
675            GenericExpr::Lit(span, lit) => GenericExpr::Lit(span, lit),
676            GenericExpr::Var(span, var) => GenericExpr::Var(span, leaf(var)),
677            GenericExpr::Call(span, op, children) => {
678                let mut mapped_children = Vec::with_capacity(children.len());
679                for child in children {
680                    mapped_children.push(child.map_symbols(head, leaf));
681                }
682                GenericExpr::Call(span, head(op), mapped_children)
683            }
684        }
685    }
686
687    /// Converts all heads and leaves to strings.
688    pub fn make_unresolved(self) -> GenericExpr<String, String> {
689        let mut map_head = |h: Head| h.to_string();
690        let mut map_leaf = |l: Leaf| l.to_string();
691        self.map_symbols(&mut map_head, &mut map_leaf)
692    }
693
694    pub fn vars(&self) -> impl Iterator<Item = Leaf> + '_ {
695        let iterator: Box<dyn Iterator<Item = Leaf>> = match self {
696            GenericExpr::Lit(_ann, _l) => Box::new(std::iter::empty()),
697            GenericExpr::Var(_ann, v) => Box::new(std::iter::once(v.clone())),
698            GenericExpr::Call(_ann, _head, exprs) => Box::new(exprs.iter().flat_map(|e| e.vars())),
699        };
700        iterator
701    }
702}
703
704impl Display for Literal {
705    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
706        match &self {
707            Literal::Int(i) => Display::fmt(i, f),
708            Literal::Float(n) => {
709                // need to display with decimal if there is none
710                let str = n.to_string();
711                if let Ok(_num) = str.parse::<i64>() {
712                    write!(f, "{}.0", str)
713                } else {
714                    write!(f, "{}", str)
715                }
716            }
717            Literal::Bool(b) => Display::fmt(b, f),
718            Literal::String(s) => write!(f, "\"{}\"", s),
719            Literal::Unit => write!(f, "()"),
720        }
721    }
722}
egglog_ast/generic_ast_helpers.rs

egglog_ast/
generic_ast_helpers.rs
