1
use merc_unsafety::ConcurrentIndexedSet;
2

            
3
use crate::SequenceForest;
4
use crate::SequenceForestContext;
5
use crate::Slot;
6
use crate::Tree;
7

            
8
/// A stable handle to a state stored in a [`DiscoveredSet`].
9
///
10
/// Handles are stable and unique per state, but the underlying indices may be
11
/// sparse rather than a contiguous `0..n`.
12
#[repr(transparent)]
13
#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
14
pub struct StateRef(usize);
15

            
16
impl StateRef {
17
    /// Returns the underlying index as a `usize`.
18
72440
    pub fn index(self) -> usize {
19
72440
        self.0
20
72440
    }
21
}
22

            
23
/// A set of `T` state vectors that deduplicates equal vectors and assigns each
24
/// a stable [`StateRef`].
25
///
26
/// The set is thread-safe: states can be inserted and looked up concurrently
27
/// through `&self` from multiple threads.
28
pub struct DiscoveredSet<T = usize>
29
where
30
    T: Copy,
31
{
32
    /// Hash-consed forest backing every stored state sequence.
33
    forest: SequenceForest<T, 2>,
34
    /// Deduplicates canonical roots and assigns each a stable index into the
35
    /// forest.
36
    states: ConcurrentIndexedSet<Tree>,
37
}
38

            
39
impl<T> DiscoveredSet<T>
40
where
41
    T: Slot,
42
{
43
    /// Creates a new empty discovered set.
44
195
    pub fn new() -> DiscoveredSet<T> {
45
195
        DiscoveredSet {
46
195
            forest: SequenceForest::new(),
47
195
            states: ConcurrentIndexedSet::new(),
48
195
        }
49
195
    }
50

            
51
    /// Creates a new empty discovered set with room for at least `capacity`
52
    /// states before reallocating.
53
    pub fn with_capacity(capacity: usize) -> DiscoveredSet<T> {
54
        DiscoveredSet {
55
            forest: SequenceForest::new(),
56
            states: ConcurrentIndexedSet::with_capacity(capacity),
57
        }
58
    }
59

            
60
    /// Inserts `state` and returns its handle together with a boolean that is
61
    /// true when the state was newly inserted and false when it was already
62
    /// present.
63
21694
    pub fn insert(&self, state: &[T]) -> (StateRef, bool) {
64
21694
        self.insert_with(state, &mut SequenceForestContext::new())
65
21694
    }
66

            
67
    /// Inserts `state` like [`DiscoveredSet::insert`], reusing the scratch
68
    /// buffers in `context`.
69
25016
    pub fn insert_with(&self, state: &[T], context: &mut SequenceForestContext) -> (StateRef, bool) {
70
25016
        let root = self.forest.insert_with(state, context);
71
25016
        let (index, is_new) = self.states.insert(root);
72
25016
        (StateRef(index), is_new)
73
25016
    }
74

            
75
    /// Returns the handle of `state` if it is present, or `None` otherwise.
76
    ///
77
    /// Note this is *not* a read-only operation. To obtain the canonical handle
78
    /// to compare against, `state` is interned into the backing forest, which
79
    /// is append-only. As a result, failures also grow the forest.
80
    pub fn index(&self, state: &[T]) -> Option<StateRef> {
81
        self.index_with(state, &mut SequenceForestContext::new())
82
    }
83

            
84
    /// Looks up `state` like [`DiscoveredSet::index`], reusing the scratch buffers in `context`.
85
    pub fn index_with(&self, state: &[T], context: &mut SequenceForestContext) -> Option<StateRef> {
86
        let root = self.forest.insert_with(state, context);
87
        self.states.index(&root).map(StateRef)
88
    }
89

            
90
    /// Returns true if `state` is present in the set. Like
91
    /// [`DiscoveredSet::index`], this interns `state` into the forest.
92
    pub fn contains(&self, state: &[T]) -> bool {
93
        self.index(state).is_some()
94
    }
95

            
96
    /// Reconstructs the state vector for `reference` into the freshly cleared
97
    /// `out` buffer.
98
22977
    pub fn get_into(&self, reference: StateRef, out: &mut Vec<T>) -> bool {
99
22977
        out.clear();
100
22977
        let Some(&tree) = self.states.get_by_index(reference.index()) else {
101
            return false;
102
        };
103
22977
        out.extend(self.forest.iter(tree));
104
22977
        true
105
22977
    }
106

            
107
    /// Returns the state vector for `reference`, allocating a fresh [`Vec`].
108
    ///
109
    /// Prefer [`DiscoveredSet::get_into`] on hot paths to reuse a buffer.
110
    pub fn get(&self, reference: StateRef) -> Option<Vec<T>> {
111
        let &tree = self.states.get_by_index(reference.index())?;
112
        Some(self.forest.iter(tree).collect())
113
    }
114

            
115
    /// Removes all states, invalidating every previously returned
116
    /// [`StateRef`].
117
    pub fn clear(&mut self) {
118
        self.forest.clear();
119
        self.states.clear();
120
    }
121
}
122

            
123
impl<T> DiscoveredSet<T>
124
where
125
    T: Copy,
126
{
127
    /// Returns the number of distinct states in the set.
128
20000
    pub fn len(&self) -> usize {
129
20000
        self.states.len()
130
20000
    }
131

            
132
    /// Returns true if the set is empty.
133
    pub fn is_empty(&self) -> bool {
134
        self.states.is_empty()
135
    }
136
}
137

            
138
impl<T> Default for DiscoveredSet<T>
139
where
140
    T: Slot,
141
{
142
    fn default() -> DiscoveredSet<T> {
143
        DiscoveredSet::new()
144
    }
145
}
146

            
147
#[cfg(test)]
148
mod tests {
149
    use std::collections::HashMap;
150
    use std::sync::Arc;
151
    use std::sync::Mutex;
152

            
153
    use rand::RngExt;
154

            
155
    use merc_utilities::random_test;
156
    use merc_utilities::random_test_threads;
157

            
158
    use super::DiscoveredSet;
159
    use super::StateRef;
160

            
161
    #[test]
162
    #[cfg_attr(miri, ignore)]
163
1
    fn random_insert_dedup() {
164
100
        random_test(100, |rng| {
165
            // Inserts random states and checks deduplication, handle stability and
166
            // reconstruction against a HashMap reference.
167
100
            let set: DiscoveredSet<usize> = DiscoveredSet::new();
168
100
            let mut seen: HashMap<Vec<usize>, StateRef> = HashMap::new();
169
100
            let mut out = Vec::new();
170

            
171
100
            for _ in 0..200 {
172
20000
                let len = rng.random_range(0..20);
173
189588
                let state: Vec<usize> = (0..len).map(|_| rng.random_range(0..10)).collect();
174

            
175
20000
                let (reference, is_new) = set.insert(&state);
176
20000
                assert_eq!(is_new, !seen.contains_key(&state), "states are new exactly once");
177
20000
                seen.entry(state.clone()).or_insert(reference);
178

            
179
20000
                assert_eq!(seen[&state], reference, "equal states share a handle");
180
20000
                assert!(reference.index() < seen.len(), "indices stay dense");
181
20000
                assert_eq!(set.len(), seen.len());
182

            
183
20000
                assert!(set.get_into(reference, &mut out));
184
20000
                assert_eq!(out, state);
185
            }
186
100
        });
187
1
    }
188

            
189
    #[test]
190
    #[cfg_attr(miri, ignore)]
191
1
    fn random_concurrent_insert() {
192
        // Inserts random states concurrently and checks that equal states share a
193
        // handle across threads.
194
1
        let set: Arc<DiscoveredSet<usize>> = Arc::new(DiscoveredSet::new());
195
1
        let seen: Arc<Mutex<HashMap<Vec<usize>, StateRef>>> = Arc::new(Mutex::new(HashMap::new()));
196

            
197
1
        random_test_threads(
198
            200,
199
            8,
200
8
            || (Arc::clone(&set), Arc::clone(&seen)),
201
1600
            move |rng, (set, seen)| {
202
1600
                let len = rng.random_range(0..20);
203
14900
                let state: Vec<usize> = (0..len).map(|_| rng.random_range(0..30)).collect();
204

            
205
1600
                let (reference, _) = set.insert(&state);
206
1600
                let mut out = Vec::new();
207
1600
                assert!(set.get_into(reference, &mut out));
208
1600
                assert_eq!(out, state);
209

            
210
1600
                let mut seen = seen.lock().unwrap();
211
1600
                match seen.get(&state) {
212
159
                    Some(&prev) => assert_eq!(prev, reference, "equal states interned to different handles"),
213
1441
                    None => {
214
1441
                        seen.insert(state, reference);
215
1441
                    }
216
                }
217
1600
            },
218
        );
219
1
    }
220
}