Skip to main content

steel_parser/
interner.rs

1use alloc::sync::Arc;
2use compact_str::CompactString;
3use core::fmt;
4use lasso::Key;
5use lasso::Spur;
6use once_cell::sync::OnceCell;
7use rustc_hash::FxBuildHasher;
8use serde::{Deserialize, Serialize};
9
10// TODO: Serialize and Deserialize should resolve() -> Otherwise we're in for deep trouble
11// trying to serialize and deserialize this
12#[derive(Copy, Clone, PartialEq, Eq, Ord, PartialOrd, Hash)]
13#[repr(transparent)]
14pub struct InternedString(Spur);
15
16impl Serialize for InternedString {
17    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
18    where
19        S: serde::Serializer,
20    {
21        self.resolve().serialize(serializer)
22    }
23}
24
25impl<'de> Deserialize<'de> for InternedString {
26    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
27    where
28        D: serde::Deserializer<'de>,
29    {
30        let key = <&str>::deserialize(deserializer)?;
31
32        Ok(InternedString::from(key))
33    }
34}
35
36impl InternedString {
37    pub fn from_static(ident: &'static str) -> Self {
38        Self(
39            INTERNER
40                .get_or_init(|| Arc::new(ThreadedRodeo::with_hasher(FxBuildHasher::default())))
41                .get_or_intern_static(ident),
42        )
43    }
44
45    pub fn from_string(ident: String) -> Self {
46        Self(
47            INTERNER
48                .get_or_init(|| Arc::new(ThreadedRodeo::with_hasher(FxBuildHasher::default())))
49                .get_or_intern(ident),
50        )
51    }
52
53    #[allow(clippy::should_implement_trait)]
54    pub fn from_str(ident: &str) -> Self {
55        Self(
56            INTERNER
57                .get_or_init(|| Arc::new(ThreadedRodeo::with_hasher(FxBuildHasher::default())))
58                .get_or_intern(ident),
59        )
60    }
61
62    pub fn new(key: usize) -> Self {
63        Self(Spur::try_from_usize(key).unwrap())
64    }
65
66    pub fn get(self) -> Spur {
67        self.0
68    }
69
70    pub fn try_get(ident: &str) -> Option<InternedString> {
71        INTERNER.get().unwrap().get(ident).map(InternedString)
72    }
73
74    #[doc(hidden)]
75    pub fn as_u32(self) -> u32 {
76        self.get().into_usize() as u32
77    }
78
79    pub fn resolve(&self) -> &str {
80        resolve(&self.0)
81    }
82}
83
84impl From<CompactString> for InternedString {
85    fn from(value: CompactString) -> Self {
86        Self::from_str(&value)
87    }
88}
89
90impl From<&str> for InternedString {
91    fn from(ident: &str) -> Self {
92        Self::from_str(ident)
93    }
94}
95
96impl From<String> for InternedString {
97    fn from(ident: String) -> Self {
98        Self::from_string(ident)
99    }
100}
101
102impl From<Spur> for InternedString {
103    fn from(spur: Spur) -> Self {
104        Self(spur)
105    }
106}
107
108impl fmt::Debug for InternedString {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        write!(f, "{}", self.get().into_usize())
111    }
112}
113
114impl fmt::Display for InternedString {
115    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116        write!(f, "{}", self.resolve())
117    }
118}
119
120// impl Serialize for InternedString {
121//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
122//     where
123//         S: Serializer,
124//     {
125//         serializer.serialize_str(self.resolve())
126//     }
127// }
128
129use lasso::ThreadedRodeo;
130
131static INTERNER: OnceCell<Arc<ThreadedRodeo<Spur, rustc_hash::FxBuildHasher>>> = OnceCell::new();
132
133pub fn interned_current_memory_usage() -> usize {
134    INTERNER.get().unwrap().current_memory_usage()
135}
136
137pub fn take_interner() -> Arc<ThreadedRodeo<Spur, rustc_hash::FxBuildHasher>> {
138    Arc::clone(INTERNER.get().unwrap())
139}
140
141pub fn initialize_with(
142    interner: Arc<ThreadedRodeo<Spur, rustc_hash::FxBuildHasher>>,
143) -> Result<(), Arc<ThreadedRodeo<Spur, rustc_hash::FxBuildHasher>>> {
144    INTERNER.set(interner)
145}
146
147pub fn get_interner() -> Option<&'static Arc<ThreadedRodeo<Spur, rustc_hash::FxBuildHasher>>> {
148    INTERNER.get()
149}
150
151pub fn add_interner(interner: Arc<ThreadedRodeo>) {
152    let guard = INTERNER.get().unwrap();
153
154    for key in interner.strings() {
155        guard.get_or_intern(key);
156    }
157}
158
159#[test]
160fn test_initialization() {
161    INTERNER.get_or_init(|| Arc::new(ThreadedRodeo::with_hasher(FxBuildHasher::default())));
162    let key = INTERNER.get().unwrap().get_or_intern_static("hello world");
163
164    let resolved_string = INTERNER.get().unwrap().resolve(&key);
165
166    println!("resolved string: {resolved_string:?}");
167}
168
169fn resolve(key: &Spur) -> &str {
170    INTERNER.get().unwrap().resolve(key)
171}