foundations/key_index.rs at main · Berylsoft/foundations · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
#![allow(clippy::type_complexity)]

use std::collections::{BTreeMap, BTreeSet};

pub struct KeyIndex<K1, K2, K3, VI> {
    idx1: BTreeMap<K1, BTreeSet<VI>>,
    idx2: BTreeMap<K2, BTreeSet<VI>>,
    idx3: BTreeMap<K3, BTreeSet<VI>>,
}

pub struct InsertError<K1, K2, K3, VI> {
    pub ty: InsertErrorType,
    pub val: (K1, K2, K3, VI),
}

pub enum InsertErrorType {
    K1,
    K2,
    K3,
}

pub struct Query<K1, K2, K3> {
    pub k1: Option<K1>,
    pub k2: Option<K2>,
    pub k3: Option<K3>,
}

impl<K1, K2, K3, VI> Default for KeyIndex<K1, K2, K3, VI> {
    fn default() -> Self {
        KeyIndex {
            idx1: BTreeMap::new(),
            idx2: BTreeMap::new(),
            idx3: BTreeMap::new(),
        }
    }
}

impl<K1, K2, K3, VI> KeyIndex<K1, K2, K3, VI>
where
    K1: Ord + Copy,
    K2: Ord + Copy,
    K3: Ord + Copy,
    VI: Ord + Copy,
{
    pub fn insert(&mut self, k1: K1, k2: K2, k3: K3, vi: VI) -> Result<(), InsertError<K1, K2, K3, VI>> {
        macro_rules! kn_impl {
            ($($kn:ident, $idxn:ident, $Kn:ident;)*) => {
                $(if !self.$idxn.entry($kn).or_default().insert(vi) {
                    return Err(InsertError {
                        ty: InsertErrorType::$Kn,
                        val: (k1, k2, k3, vi),
                    });
                })*
            };
        }
        kn_impl!(k1, idx1, K1; k2, idx2, K2; k3, idx3, K3;);
        Ok(())
    }

    pub fn query(&self, query: Query<K1, K2, K3>) -> BTreeSet<VI> {
        let Query { k1, k2, k3 } = query;
        let mut res = BTreeSet::new();
        macro_rules! kn_impl {
            ($($kn:ident, $idxn:ident;)*) => {
                $(if let Some(k) = $kn {
                    if let Some(s) = self.$idxn.get(&k) {
                        for vi in s {
                            assert!(res.insert(*vi));
                        }
                    }
                })*
            };
        }
        kn_impl!(k1, idx1; k2, idx2; k3, idx3;);
        res
    }
}

impl<K1, K2, K3> KeyIndex<K1, K2, K3, usize>
where
    K1: Ord + Copy,
    K2: Ord + Copy,
    K3: Ord + Copy,
{
    pub fn from_seq<V, S: AsRef<[(K1, K2, K3, V)]>>(seq: S) -> Result<KeyIndex<K1, K2, K3, usize>, InsertError<K1, K2, K3, usize>> {
        let mut res = KeyIndex::default();
        for (vi, (k1, k2, k3, _)) in seq.as_ref().iter().enumerate() {
            res.insert(*k1, *k2, *k3, vi)?;
        }
        Ok(res)
    }
}