egglog/
util.rs

1use crate::{ast::ResolvedVar, core::ResolvedCall};
2
3pub(crate) type BuildHasher = std::hash::BuildHasherDefault<rustc_hash::FxHasher>;
4pub(crate) type HashMap<K, V> = hashbrown::HashMap<K, V, BuildHasher>;
5pub(crate) type HashSet<K> = hashbrown::HashSet<K, BuildHasher>;
6pub(crate) type HEntry<'a, A, B> = hashbrown::hash_map::Entry<'a, A, B, BuildHasher>;
7pub type IndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasher>;
8pub type IndexSet<K> = indexmap::IndexSet<K, BuildHasher>;
9
10pub use egglog_ast::generic_ast_helpers::INTERNAL_SYMBOL_PREFIX;
11
12/// Generates fresh symbols for internal use during typechecking and flattening.
13/// These are guaranteed not to collide with the
14/// user's symbols because they use a reserved prefix.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct SymbolGen {
17    hint_to_count: HashMap<String, usize>,
18    reserved_string: String,
19    leave_off_zero: bool,
20}
21
22impl SymbolGen {
23    /// Create a new symbol generator with the given reserved prefix.
24    pub fn new(reserved_string: String) -> Self {
25        Self {
26            hint_to_count: HashMap::default(),
27            reserved_string,
28            leave_off_zero: true,
29        }
30    }
31
32    /// By default, the first symbol generated with a given hint
33    /// does not have a numeric suffix (e.g., "var" instead of "var0").
34    /// This method changes that behavior.
35    pub fn include_zero(&mut self, include: bool) {
36        self.leave_off_zero = !include;
37    }
38
39    /// Check if this symbol generator has been used to generate any symbols.
40    pub fn has_been_used(&self) -> bool {
41        !self.hint_to_count.is_empty()
42    }
43
44    /// Get the reserved prefix used by this symbol generator.
45    pub fn reserved_prefix(&self) -> &str {
46        &self.reserved_string
47    }
48
49    /// Check if the given symbol is reserved (i.e., starts with the reserved prefix).
50    pub fn is_reserved(&self, symbol: &str) -> bool {
51        !self.reserved_string.is_empty() && symbol.starts_with(&self.reserved_string)
52    }
53}
54
55/// This trait lets us statically dispatch between `fresh` methods for generic structs.
56pub trait FreshGen<Head: ?Sized, Leaf> {
57    fn fresh(&mut self, name_hint: &Head) -> Leaf;
58}
59
60impl FreshGen<str, String> for SymbolGen {
61    fn fresh(&mut self, name_hint: &str) -> String {
62        let entry = self.hint_to_count.entry(name_hint.to_string()).or_insert(0);
63        let count_before = *entry;
64        *entry += 1;
65        format!(
66            "{}{}{}",
67            self.reserved_string,
68            name_hint,
69            if self.leave_off_zero && count_before == 0 {
70                "".to_string()
71            } else {
72                count_before.to_string()
73            }
74        )
75    }
76}
77
78impl FreshGen<String, String> for SymbolGen {
79    fn fresh(&mut self, name_hint: &String) -> String {
80        self.fresh(name_hint.as_str())
81    }
82}
83
84impl FreshGen<ResolvedCall, ResolvedVar> for SymbolGen {
85    fn fresh(&mut self, name_hint: &ResolvedCall) -> ResolvedVar {
86        let entry = self
87            .hint_to_count
88            .entry(format!("{name_hint}"))
89            .or_insert(0);
90        let count = *entry;
91        *entry += 1;
92        let name = format!(
93            "{}{}{}",
94            self.reserved_string,
95            name_hint,
96            if self.leave_off_zero && count == 0 {
97                "".to_string()
98            } else {
99                count.to_string()
100            }
101        );
102        let sort = match name_hint {
103            ResolvedCall::Func(f) => f.output.clone(),
104            ResolvedCall::Primitive(prim) => prim.output().clone(),
105        };
106        ResolvedVar {
107            name,
108            sort,
109            // fresh variables are never global references, since globals
110            // are desugared away by `remove_globals`
111            is_global_ref: false,
112        }
113    }
114}
egglog/util.rs

egglog/
util.rs
