1
use std::hash::Hash;
2
use std::hash::Hasher;
3
use std::sync::Arc;
4
use std::sync::atomic::AtomicUsize;
5
use std::sync::atomic::Ordering;
6

            
7
use dashmap::DashMap;
8
use equivalent::Equivalent;
9
use log::debug;
10
use rustc_hash::FxBuildHasher;
11

            
12
use merc_unsafety::AllocBlock;
13
use merc_unsafety::BlockAllocatorSafe;
14
use merc_unsafety::StablePointer;
15
use merc_unsafety::StablePointerSet;
16

            
17
use crate::Symb;
18
use crate::SymbolIndex;
19
use crate::SymbolRef;
20

            
21
/// Pool for maximal sharing of function symbols, see [SymbolRef]. Ensures that function symbols
22
/// with the same name and arity point to the same [SharedSymbol] object.
23
/// Returns [crate::Symbol] that can be used to refer to the shared symbol, avoiding
24
/// garbage collection of the underlying shared symbol.
25
pub struct SymbolPool {
26
    /// Unique table of all function symbols
27
    symbols: StablePointerSet<SharedSymbol, FxBuildHasher, AllocBlock<SharedSymbol, 1024>>,
28

            
29
    /// A map from prefixes to counters that track the next available index for function symbols
30
    prefix_to_register_function_map: DashMap<String, Arc<AtomicUsize>, FxBuildHasher>,
31
}
32

            
33
impl SymbolPool {
34
    /// Creates a new empty symbol pool.
35
959
    pub(crate) fn new() -> Self {
36
959
        Self {
37
959
            symbols: StablePointerSet::with_hasher_in(FxBuildHasher, AllocBlock::new()),
38
959
            prefix_to_register_function_map: DashMap::with_hasher(FxBuildHasher),
39
959
        }
40
959
    }
41

            
42
    /// Creates or retrieves a function symbol with the given name and arity.
43
4984649
    pub fn create<N>(&self, name: N, arity: usize) -> StablePointer<SharedSymbol>
44
4984649
    where
45
4984649
        N: Into<String> + AsRef<str>,
46
    {
47
        // Get or create symbol index
48
4984649
        let (shared_symbol, inserted) = self.symbols.insert_equiv(&SharedSymbolLookup { name, arity });
49

            
50
4984649
        if inserted {
51
80554
            // If the symbol was newly created, register its prefix.
52
80554
            self.update_prefix(shared_symbol.name());
53
4904095
        }
54

            
55
        // Return cloned symbol
56
4984649
        shared_symbol
57
4984649
    }
58

            
59
    /// Return the symbol of the SharedTerm for the given ATermRef
60
    pub fn symbol_name<'a>(&self, symbol: &'a SymbolRef<'a>) -> &'a str {
61
        symbol.shared().name()
62
    }
63

            
64
    /// Returns the arity of the function symbol
65
    pub fn symbol_arity<'a, 'b, S: Symb<'a, 'b>>(&self, symbol: &'b S) -> usize {
66
        symbol.shared().arity()
67
    }
68

            
69
    /// Returns the number of symbols in the pool.
70
388407
    pub fn len(&self) -> usize {
71
388407
        self.symbols.len()
72
388407
    }
73

            
74
    /// Returns true if the pool is empty.
75
    pub fn is_empty(&self) -> bool {
76
        self.symbols.is_empty()
77
    }
78

            
79
    /// Returns the capacity of the pool.
80
    pub fn capacity(&self) -> usize {
81
        self.symbols.capacity()
82
    }
83

            
84
    /// Retain only symbols satisfying the given predicate.
85
388407
    pub fn retain<F>(&mut self, mut f: F)
86
388407
    where
87
388407
        F: FnMut(&SymbolIndex) -> bool,
88
    {
89
12847785
        self.symbols.retain(|element| f(element));
90

            
91
388407
        let removed_blocks = self.symbols.allocator_mut().remove_free_blocks();
92
388407
        debug!("Removed {} blocks from the symbol pool", removed_blocks);
93
388407
    }
94

            
95
    /// Creates a new prefix counter for the given prefix.
96
1
    pub fn create_prefix(&self, prefix: &str) -> Arc<AtomicUsize> {
97
        // Create a new counter for the prefix if it does not exist
98
1
        let result = match self.prefix_to_register_function_map.get(prefix) {
99
            Some(result) => result.clone(),
100
            None => {
101
1
                let result = Arc::new(AtomicUsize::new(0));
102
1
                assert!(
103
1
                    self.prefix_to_register_function_map
104
1
                        .insert(prefix.to_string(), result.clone())
105
1
                        .is_none(),
106
                    "This key should not yet exist"
107
                );
108
1
                result
109
            }
110
        };
111

            
112
        // Ensure the counter starts at a sufficiently large index
113
1
        self.get_sufficiently_large_postfix_index(prefix, &result);
114
1
        result
115
1
    }
116

            
117
    /// Removes a prefix counter from the pool.
118
    pub fn remove_prefix(&self, prefix: &str) {
119
        // Remove the prefix counter if it exists
120
        self.prefix_to_register_function_map.remove(prefix);
121
    }
122

            
123
    /// Updates the counter for a registered prefix for the newly created symbol.
124
87026
    fn update_prefix(&self, name: &str) {
125
        // Check whether there is a registered prefix p such that name equal pn where n is a number.
126
        // In that case prevent that pn will be generated as a fresh function name.
127
87026
        let start_of_index = name
128
93591
            .rfind(|c: char| !c.is_ascii_digit())
129
87026
            .map(|pos| pos + 1)
130
87026
            .unwrap_or(0);
131

            
132
87026
        if start_of_index < name.len() {
133
11506
            let potential_number = &name[start_of_index..];
134
11506
            let prefix = &name[..start_of_index];
135

            
136
11506
            if let Some(counter) = self.prefix_to_register_function_map.get(prefix)
137
1
                && let Ok(number) = potential_number.parse::<usize>()
138
1
            {
139
1
                counter.fetch_max(number + 1, Ordering::Relaxed);
140
11505
            }
141
75520
        }
142
87026
    }
143

            
144
    /// Traverse all symbols to find the maximum numeric suffix for this prefix
145
1
    fn get_sufficiently_large_postfix_index(&self, prefix: &str, counter: &Arc<AtomicUsize>) {
146
5
        for symbol in self.symbols.iter() {
147
5
            let name = symbol.name();
148
5
            if name.starts_with(prefix) {
149
                // Symbol name starts with the prefix, check for numeric suffix
150
2
                let suffix_start = prefix.len();
151
2
                if suffix_start < name.len() {
152
2
                    let suffix = &name[suffix_start..];
153
2
                    if let Ok(number) = suffix.parse::<usize>() {
154
1
                        // There is a numeric suffix, update the counter if it's larger
155
1
                        counter.fetch_max(number + 1, Ordering::Relaxed);
156
1
                    }
157
                }
158
3
            }
159
        }
160
1
    }
161
}
162

            
163
/// Represents a function symbol with a name and arity.
164
#[derive(Debug, Clone, Eq, PartialEq)]
165
pub struct SharedSymbol {
166
    /// Name of the function
167
    name: String,
168
    /// Number of arguments
169
    arity: usize,
170
}
171

            
172
/// SAFETY: The `SharedSymbol` is never equal to the sentinel value.
173
unsafe impl BlockAllocatorSafe for SharedSymbol {}
174

            
175
impl SharedSymbol {
176
    /// Creates a new function symbol.
177
87026
    pub fn new<N: Into<String>>(name: N, arity: usize) -> Self {
178
87026
        Self {
179
87026
            name: name.into(),
180
87026
            arity,
181
87026
        }
182
87026
    }
183

            
184
    /// Returns the name of the function symbol
185
2325669
    pub fn name(&self) -> &str {
186
2325669
        &self.name
187
2325669
    }
188

            
189
    /// Returns the arity of the function symbol
190
2848501290
    pub fn arity(&self) -> usize {
191
2848501290
        self.arity
192
2848501290
    }
193

            
194
    /// Returns a unique index for this shared symbol
195
    pub fn index(&self) -> usize {
196
        self as *const Self as *const u8 as usize
197
    }
198
}
199

            
200
/// A cheap way to look up SharedSymbol
201
struct SharedSymbolLookup<T: Into<String> + AsRef<str>> {
202
    name: T,
203
    arity: usize,
204
}
205

            
206
impl<T: Into<String> + AsRef<str>> From<&SharedSymbolLookup<T>> for SharedSymbol {
207
80554
    fn from(lookup: &SharedSymbolLookup<T>) -> Self {
208
        // TODO: Not optimal
209
80554
        let string = lookup.name.as_ref().to_string();
210
80554
        Self::new(string, lookup.arity)
211
80554
    }
212
}
213

            
214
impl<T: Into<String> + AsRef<str>> Equivalent<SharedSymbol> for SharedSymbolLookup<T> {
215
4911991
    fn equivalent(&self, other: &SharedSymbol) -> bool {
216
4911991
        self.name.as_ref() == other.name && self.arity == other.arity
217
4911991
    }
218
}
219

            
220
/// These hash implementations should be the same as `SharedSymbol`.
221
impl<T: Into<String> + AsRef<str>> Hash for SharedSymbolLookup<T> {
222
4984649
    fn hash<H: Hasher>(&self, state: &mut H) {
223
4984649
        self.name.as_ref().hash(state);
224
4984649
        self.arity.hash(state);
225
4984649
    }
226
}
227

            
228
impl Hash for SharedSymbol {
229
112034
    fn hash<H: Hasher>(&self, state: &mut H) {
230
112034
        self.name.hash(state);
231
112034
        self.arity.hash(state);
232
112034
    }
233
}
234

            
235
#[cfg(test)]
236
mod tests {
237
    use std::sync::atomic::Ordering;
238

            
239
    use crate::Symbol;
240
    use crate::storage::THREAD_TERM_POOL;
241

            
242
    #[test]
243
1
    fn test_symbol_sharing() {
244
1
        let _ = merc_utilities::test_logger();
245

            
246
1
        let f1 = Symbol::new("f", 2);
247
1
        let f2 = Symbol::new("f", 2);
248

            
249
        // Should be the same object
250
1
        assert_eq!(f1, f2);
251
1
    }
252

            
253
    #[test]
254
1
    fn test_prefix_counter() {
255
1
        let _ = merc_utilities::test_logger();
256

            
257
1
        let _symbol = Symbol::new("x69", 0);
258
1
        let _symbol2 = Symbol::new("x_y", 0);
259

            
260
1
        let value =
261
1
            THREAD_TERM_POOL.with_borrow(|tp| tp.term_pool().write().expect("Lock poisoned!").register_prefix("x"));
262

            
263
1
        assert_eq!(value.load(Ordering::Relaxed), 70);
264

            
265
1
        let _symbol3 = Symbol::new("x_no_effect", 0);
266
1
        let _symbol4 = Symbol::new("x130", 0);
267

            
268
1
        assert_eq!(value.load(Ordering::Relaxed), 131);
269
1
    }
270
}