egglog/
util.rs

1use crate::{
2    ast::ResolvedVar,
3    core::{ResolvedCall, SpecializedPrimitive},
4};
5use std::borrow::Cow;
6
7pub const INTERNAL_SYMBOL_PREFIX: &str = "@";
8
9/// Gets rid of internal symbol prefixes for printing.
10/// This allows us to test parsing of desugared programs.
11pub fn sanitize_internal_name(name: &str) -> Cow<'_, str> {
12    if let Some(stripped) = name.strip_prefix(INTERNAL_SYMBOL_PREFIX) {
13        Cow::Owned(format!("_{}", stripped))
14    } else {
15        Cow::Borrowed(name)
16    }
17}
18
19pub(crate) type BuildHasher = std::hash::BuildHasherDefault<rustc_hash::FxHasher>;
20pub(crate) type HashMap<K, V> = hashbrown::HashMap<K, V, BuildHasher>;
21pub(crate) type HashSet<K> = hashbrown::HashSet<K, BuildHasher>;
22pub(crate) type HEntry<'a, A, B> = hashbrown::hash_map::Entry<'a, A, B, BuildHasher>;
23pub type IndexMap<K, V> = indexmap::IndexMap<K, V, BuildHasher>;
24pub type IndexSet<K> = indexmap::IndexSet<K, BuildHasher>;
25
26/// Generates fresh symbols for internal use during typechecking and flattening.
27/// These are guaranteed not to collide with the
28/// user's symbols because they use a reserved prefix.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct SymbolGen {
31    count: usize,
32    reserved_string: String,
33}
34
35impl SymbolGen {
36    pub fn new(reserved_string: String) -> Self {
37        Self {
38            count: 0,
39            reserved_string,
40        }
41    }
42
43    pub fn has_been_used(&self) -> bool {
44        self.count > 0
45    }
46
47    pub fn reserved_prefix(&self) -> &str {
48        &self.reserved_string
49    }
50
51    pub fn is_reserved(&self, symbol: &str) -> bool {
52        !self.reserved_string.is_empty() && symbol.starts_with(&self.reserved_string)
53    }
54}
55
56/// This trait lets us statically dispatch between `fresh` methods for generic structs.
57pub trait FreshGen<Head: ?Sized, Leaf> {
58    fn fresh(&mut self, name_hint: &Head) -> Leaf;
59}
60
61impl FreshGen<str, String> for SymbolGen {
62    fn fresh(&mut self, name_hint: &str) -> String {
63        let s = format!("{}{}{}", self.reserved_string, name_hint, self.count);
64        self.count += 1;
65        s
66    }
67}
68
69impl FreshGen<String, String> for SymbolGen {
70    fn fresh(&mut self, name_hint: &String) -> String {
71        self.fresh(name_hint.as_str())
72    }
73}
74
75impl FreshGen<ResolvedCall, ResolvedVar> for SymbolGen {
76    fn fresh(&mut self, name_hint: &ResolvedCall) -> ResolvedVar {
77        let name = format!("{}{}{}", self.reserved_string, name_hint, self.count);
78        self.count += 1;
79        let sort = match name_hint {
80            ResolvedCall::Func(f) => f.output.clone(),
81            ResolvedCall::Primitive(SpecializedPrimitive { output, .. }) => output.clone(),
82        };
83        ResolvedVar {
84            name,
85            sort,
86            // fresh variables are never global references, since globals
87            // are desugared away by `remove_globals`
88            is_global_ref: false,
89        }
90    }
91}