Skip to main content

sk_store/
index.rs

1use std::collections::HashMap;
2use std::mem::take;
3
4use serde::{
5    Deserialize,
6    Serialize,
7};
8use sk_core::k8s::{
9    GVK,
10    format_gvk_name,
11};
12
13#[derive(Clone, Debug, Default, Deserialize, Serialize)]
14pub struct TraceIndex {
15    #[serde(flatten)]
16    index: HashMap<GVK, HashMap<String, u64>>,
17}
18
19impl TraceIndex {
20    pub fn new() -> TraceIndex {
21        TraceIndex::default()
22    }
23
24    pub fn contains(&self, gvk: &GVK, ns_name: &str) -> bool {
25        self.index.get(gvk).is_some_and(|gvk_hash| gvk_hash.contains_key(ns_name))
26    }
27
28    pub fn flattened_keys(&self) -> Vec<String> {
29        self.index
30            .iter()
31            .flat_map(|(gvk, gvk_hash)| gvk_hash.keys().map(move |ns_name| format_gvk_name(gvk, ns_name)))
32            .collect()
33    }
34
35    pub fn get(&self, gvk: &GVK, ns_name: &str) -> Option<u64> {
36        self.index.get(gvk)?.get(ns_name).cloned()
37    }
38
39    pub fn insert(&mut self, gvk: GVK, ns_name: String, hash: u64) {
40        self.index.entry(gvk).or_default().insert(ns_name, hash);
41    }
42
43    pub fn is_empty(&self) -> bool {
44        self.index.values().all(|gvk_hash| gvk_hash.is_empty())
45    }
46
47    pub fn len(&self) -> usize {
48        self.index.values().map(|gvk_hash| gvk_hash.len()).sum()
49    }
50
51    pub fn remove(&mut self, gvk: GVK, ns_name: &str) {
52        self.index.entry(gvk).and_modify(|gvk_hash| {
53            gvk_hash.remove(ns_name);
54        });
55    }
56
57    pub fn take_gvk_index(&mut self, gvk: &GVK) -> HashMap<String, u64> {
58        take(self.index.get_mut(gvk).unwrap_or(&mut HashMap::new()))
59    }
60}