1
use std::fmt;
2
use std::sync::Arc;
3

            
4
use merc_unsafety::ShardedHashMap;
5
use rustc_hash::FxBuildHasher;
6

            
7
use merc_utilities::MercError;
8
use merc_utilities::ShardedCounter;
9

            
10
use crate::LPS;
11
use crate::SequenceForest;
12
use crate::SequenceForestContext;
13
use crate::Summand;
14
use crate::Tree;
15

            
16
/// Controls the enumeration caching behaviour.
17
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
18
#[cfg_attr(feature = "clap", derive(clap::ValueEnum))]
19
pub enum CachingStrategy {
20
    /// No caching; every enumeration is performed from scratch.
21
    None,
22
    /// Each summand maintains its own independent cache.
23
    Local,
24
}
25

            
26
/// Number of shards in each summand's local cache.
27
const CACHE_SHARDS: usize = 16;
28

            
29
/// A wrapper around an [`LPS`] that caches summand enumeration results.
30
///
31
/// The cache key for each summand is formed by projecting the state vector onto
32
/// the summand's read positions. The cache is thread-safe, and local context is
33
/// used for the scratch buffers required to project the state vector and
34
/// reconstruct next states from cached values.
35
pub struct CacheLPS<P: LPS> {
36
    inner: Arc<P>,
37
    summands: Vec<CacheSummandWrapper<P>>,
38
}
39

            
40
/// A single cached enumeration result keyed by the read-position projection.
41
///
42
/// Stored as a bare element in a summand's cache table. The results indicate the
43
/// transition label and the captured write-position values for each next state.
44
struct CacheEntry<L> {
45
    /// Hash-consed projection of the source state onto the read positions.
46
    key: Tree,
47
    /// Captured `(label, write-position values tree)` pairs for the key.
48
    results: Vec<(L, Tree)>,
49
}
50

            
51
/// Per-thread reusable scratch buffers for enumeration.
52
pub struct CacheContext<P: LPS> {
53
    /// Buffer projecting the state vector onto read (then write) positions.
54
    key_buf: Vec<P::Value>,
55
    /// Buffer reconstructing a next-state from a cached write-position tree.
56
    replay_buf: Vec<P::Value>,
57
    /// Context reused when interning into the forest.
58
    forest_context: SequenceForestContext,
59
    /// Enumeration context for the wrapped inner summand (cache misses).
60
    inner: <P::Summand as Summand>::Context,
61
}
62

            
63
/// Thin metadata wrapper for a single cached summand.
64
pub struct CacheSummandWrapper<P: LPS> {
65
    /// Index of the summand in the LPS, used for cache lookups.
66
    index: usize,
67

            
68
    /// Positions in the state vector that are read by this summand; these form the cache key.
69
    read_positions: Vec<usize>,
70
    /// Positions in the state vector that are written by this summand; these are stored in the cache values.
71
    write_positions: Vec<usize>,
72

            
73
    /// Caching strategy in effect for this summand.
74
    strategy: CachingStrategy,
75

            
76
    /// This summand's local enumeration cache, keyed by the read-position
77
    /// projection tree. One independent cache per summand, as intended by the
78
    /// `Local` strategy.
79
    cache: ShardedHashMap<CacheEntry<P::Label>, FxBuildHasher>,
80

            
81
    /// Hash-consed forest holding the keys and captured values. Shared across
82
    /// all summands so equal sequences are stored once.
83
    forest: Arc<SequenceForest<P::Value, 2>>,
84

            
85
    /// Shared reference to the inner LPS for delegating cache misses.
86
    inner: Arc<P>,
87

            
88
    /// Number of enumerations served from the cache. Only updated when the
89
    /// `metrics` feature is enabled; otherwise it stays at zero.
90
    hits: ShardedCounter,
91
    /// Number of enumerations delegated to the inner summand. Only updated when
92
    /// the `metrics` feature is enabled; otherwise it stays at zero.
93
    misses: ShardedCounter,
94
}
95

            
96
impl<P: LPS> CacheLPS<P> {
97
18
    pub fn new(inner: P, strategy: CachingStrategy) -> Self {
98
18
        let inner = Arc::new(inner);
99
18
        let forest = Arc::new(SequenceForest::new());
100

            
101
18
        let summands: Vec<_> = inner
102
18
            .summands()
103
18
            .iter()
104
18
            .enumerate()
105
18
            .map(|(i, s)| CacheSummandWrapper {
106
40
                index: i,
107
40
                read_positions: s.read_positions().to_vec(),
108
40
                write_positions: s.write_positions().to_vec(),
109
40
                strategy,
110
40
                cache: ShardedHashMap::with_shards(CACHE_SHARDS),
111
40
                forest: Arc::clone(&forest),
112
40
                inner: Arc::clone(&inner),
113
40
                hits: ShardedCounter::new(),
114
40
                misses: ShardedCounter::new(),
115
40
            })
116
18
            .collect();
117

            
118
18
        CacheLPS { inner, summands }
119
18
    }
120

            
121
    /// Collects per-summand cache metrics for this [`CacheLPS`].
122
    ///
123
    /// The returned [`CacheMetrics`] implements [`fmt::Display`] for a
124
    /// human-readable summary of cache hits, misses and occupancy.
125
    pub fn metrics(&self) -> CacheMetrics {
126
        let summands = self
127
            .summands
128
            .iter()
129
            .map(|s| {
130
                let entries = match s.strategy {
131
                    CachingStrategy::None => 0,
132
                    CachingStrategy::Local => s.cache.len(),
133
                };
134

            
135
                SummandCacheMetrics {
136
                    index: s.index,
137
                    strategy: s.strategy,
138
                    hits: s.hits.get(),
139
                    misses: s.misses.get(),
140
                    entries,
141
                }
142
            })
143
            .collect();
144

            
145
        CacheMetrics { summands }
146
    }
147
}
148

            
149
/// Cache metrics for a single summand.
150
#[derive(Clone, Copy, Debug)]
151
pub struct SummandCacheMetrics {
152
    /// Index of the summand in the LPS.
153
    pub index: usize,
154
    /// Caching strategy in effect for this summand.
155
    pub strategy: CachingStrategy,
156
    /// Number of enumerations served from the cache.
157
    pub hits: u64,
158
    /// Number of enumerations delegated to the inner summand.
159
    pub misses: u64,
160
    /// Number of distinct keys currently stored for this summand.
161
    pub entries: usize,
162
}
163

            
164
impl SummandCacheMetrics {
165
    /// Total number of enumeration lookups.
166
    pub fn lookups(&self) -> u64 {
167
        self.hits + self.misses
168
    }
169

            
170
    /// Fraction of lookups served from the cache, in `[0.0, 1.0]`.
171
    ///
172
    /// Returns `0.0` when there were no lookups.
173
    pub fn hit_rate(&self) -> f64 {
174
        let lookups = self.lookups();
175
        if lookups == 0 {
176
            0.0
177
        } else {
178
            self.hits as f64 / lookups as f64
179
        }
180
    }
181
}
182

            
183
/// Aggregated cache metrics for every summand of a [`CacheLPS`].
184
#[derive(Clone, Debug)]
185
pub struct CacheMetrics {
186
    /// Per-summand metrics, ordered by summand index.
187
    pub summands: Vec<SummandCacheMetrics>,
188
}
189

            
190
impl CacheMetrics {
191
    /// Total number of cache hits across all summands.
192
    pub fn total_hits(&self) -> u64 {
193
        self.summands.iter().map(|s| s.hits).sum()
194
    }
195

            
196
    /// Total number of cache misses across all summands.
197
    pub fn total_misses(&self) -> u64 {
198
        self.summands.iter().map(|s| s.misses).sum()
199
    }
200

            
201
    /// Total number of cached entries across all summands.
202
    pub fn total_entries(&self) -> usize {
203
        self.summands.iter().map(|s| s.entries).sum()
204
    }
205

            
206
    /// Fraction of lookups served from the cache across all summands.
207
    pub fn hit_rate(&self) -> f64 {
208
        let hits = self.total_hits();
209
        let lookups = hits + self.total_misses();
210
        if lookups == 0 {
211
            0.0
212
        } else {
213
            hits as f64 / lookups as f64
214
        }
215
    }
216
}
217

            
218
impl fmt::Display for CacheMetrics {
219
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
220
        if cfg!(not(feature = "metrics")) {
221
            return write!(f, "enable the 'metrics' feature to see cache metrics");
222
        }
223

            
224
        writeln!(f, "summand cache metrics:")?;
225
        writeln!(
226
            f,
227
            "  {:>7}  {:>8}  {:>10}  {:>10}  {:>8}  {:>8}",
228
            "summand", "strategy", "hits", "misses", "hit%", "entries"
229
        )?;
230
        for s in &self.summands {
231
            let strategy = match s.strategy {
232
                CachingStrategy::None => "none",
233
                CachingStrategy::Local => "local",
234
            };
235
            writeln!(
236
                f,
237
                "  {:>7}  {:>8}  {:>10}  {:>10}  {:>7.1}%  {:>8}",
238
                s.index,
239
                strategy,
240
                s.hits,
241
                s.misses,
242
                s.hit_rate() * 100.0,
243
                s.entries
244
            )?;
245
        }
246
        write!(
247
            f,
248
            "  total: {} hits, {} misses ({:.1}% hit rate), {} entries",
249
            self.total_hits(),
250
            self.total_misses(),
251
            self.hit_rate() * 100.0,
252
            self.total_entries()
253
        )
254
    }
255
}
256

            
257
impl<P: LPS> LPS for CacheLPS<P> {
258
    type Value = P::Value;
259
    type Label = P::Label;
260
    type StateInfo = P::StateInfo;
261
    type Summand = CacheSummandWrapper<P>;
262

            
263
36
    fn initial_state(&self) -> Vec<Self::Value> {
264
36
        self.inner.initial_state()
265
36
    }
266

            
267
544
    fn summands(&self) -> &[Self::Summand] {
268
544
        &self.summands
269
544
    }
270

            
271
90
    fn create_context(&self) -> <Self::Summand as Summand>::Context {
272
90
        CacheContext {
273
90
            key_buf: Vec::new(),
274
90
            replay_buf: Vec::new(),
275
90
            forest_context: SequenceForestContext::new(),
276
90
            inner: self.inner.create_context(),
277
90
        }
278
90
    }
279

            
280
544
    fn prepare<'a>(
281
544
        &'a self,
282
544
        context: &mut <Self::Summand as Summand>::Context,
283
544
        state: &'a [Self::Value],
284
544
    ) -> impl Iterator<Item = usize> + 'a {
285
        // The cache wraps each inner summand one-for-one in order, so the inner
286
        // summand indices map directly onto `self.summands`.
287
544
        self.inner.prepare(&mut context.inner, state)
288
544
    }
289

            
290
544
    fn state_info(&self, state: &[Self::Value]) -> Self::StateInfo {
291
544
        self.inner.state_info(state)
292
544
    }
293
}
294

            
295
impl<P: LPS> CacheSummandWrapper<P> {
296
730
    fn replay_cached(
297
730
        &self,
298
730
        context: &mut CacheContext<P>,
299
730
        state: &[P::Value],
300
730
        results: &[(P::Label, Tree)],
301
730
        report: &mut impl FnMut(&P::Label, &[P::Value]) -> Result<(), MercError>,
302
730
    ) -> Result<(), MercError> {
303
        // Each cached tree holds only the values at the write positions.
304
        // The remaining positions are passed through unchanged, so they
305
        // must be taken from the *live* source state rather than the state
306
        // for which the entry was originally computed.
307
730
        let replay_buf = &mut context.replay_buf;
308
730
        for (label, write_tree) in results {
309
585
            replay_buf.clear();
310
585
            replay_buf.extend_from_slice(state);
311
636
            for (&pos, value) in self.write_positions.iter().zip(self.forest.iter(*write_tree)) {
312
636
                replay_buf[pos] = value;
313
636
            }
314
585
            report(label, replay_buf)?;
315
        }
316
730
        Ok(())
317
730
    }
318
}
319

            
320
impl<P: LPS> Summand for CacheSummandWrapper<P> {
321
    type Value = P::Value;
322
    type Label = P::Label;
323
    type Context = CacheContext<P>;
324

            
325
    /// # Concurrency
326
    ///
327
    /// On a cache hit the captured results are replayed in place while this
328
    /// summand's cache shard is **read-locked**, so `report` runs under that
329
    /// lock. As a consequence `report` must not re-enter this same summand's
330
    /// `enumerate`, because that will deadlock.
331
1636
    fn enumerate<F>(&self, context: &mut Self::Context, state: &[Self::Value], mut report: F) -> Result<(), MercError>
332
1636
    where
333
1636
        F: FnMut(&Self::Label, &[Self::Value]) -> Result<(), MercError>,
334
    {
335
1636
        if self.strategy == CachingStrategy::None {
336
818
            return self.inner.summands()[self.index].enumerate(&mut context.inner, state, report);
337
818
        }
338

            
339
        // Project the state vector onto the read positions to form the cache
340
        // key, interning it into the shared forest.
341
818
        let key_tree = {
342
            let CacheContext {
343
818
                key_buf,
344
818
                forest_context,
345
                ..
346
818
            } = &mut *context;
347
818
            key_buf.clear();
348
938
            for &pos in &self.read_positions {
349
938
                key_buf.push(state[pos]);
350
938
            }
351
818
            self.forest.insert_with(key_buf, forest_context)
352
        };
353

            
354
        // This summand owns its cache, so the key is the projection tree alone.
355
818
        let hash = self.cache.hash(&key_tree);
356
818
        let eq = |entry: &CacheEntry<P::Label>| entry.key == key_tree;
357

            
358
        // Fast path: a present entry is replayed in place under the shard read
359
        // lock, without cloning the entry's captured results.
360
818
        let hit = self.cache.find_with(hash, eq, |entry| {
361
730
            self.replay_cached(context, state, &entry.results, &mut report)
362
730
        });
363
818
        if let Some(result) = hit {
364
            #[cfg(feature = "metrics")]
365
            self.hits.increment();
366
730
            result?;
367
730
            return Ok(());
368
88
        }
369

            
370
        #[cfg(feature = "metrics")]
371
        self.misses.increment();
372
        // Cache MISS: delegate to inner summand, capture results. Only the
373
        // values at the write positions are stored; on replay they are
374
        // scattered back onto the live source state (see the hit branch).
375
88
        let mut captured: Vec<(P::Label, Tree)> = Vec::new();
376
        {
377
88
            let inner_summand = &self.inner.summands()[self.index];
378
88
            let forest = &self.forest;
379
88
            let write_positions = &self.write_positions;
380
            let CacheContext {
381
88
                key_buf,
382
88
                forest_context,
383
88
                inner,
384
                ..
385
88
            } = &mut *context;
386

            
387
88
            inner_summand.enumerate(inner, state, |label, next_state| {
388
75
                key_buf.clear();
389
90
                for &pos in write_positions {
390
90
                    key_buf.push(next_state[pos]);
391
90
                }
392
75
                let tree = forest.insert_with(key_buf, forest_context);
393
75
                captured.push((label.clone(), tree));
394
75
                report(label, next_state)
395
75
            })?;
396
        }
397

            
398
        // Publish the captured results. A concurrent thread may have inserted
399
        // the same key meanwhile; `insert_with` keeps the resident entry and our
400
        // copy is dropped. Either way the transitions above were already
401
        // reported.
402
88
        self.cache.insert_with(
403
88
            hash,
404
88
            eq,
405
            |entry| self.cache.hash(&entry.key),
406
            || CacheEntry {
407
88
                key: key_tree,
408
88
                results: captured,
409
88
            },
410
        );
411

            
412
88
        Ok(())
413
1636
    }
414

            
415
    fn read_positions(&self) -> &[usize] {
416
        &self.read_positions
417
    }
418

            
419
    fn write_positions(&self) -> &[usize] {
420
        &self.write_positions
421
    }
422
}