foundations/key_index_map.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
#![allow(clippy::type_complexity)]

use alloc::vec::Vec;
use std::collections::VecDeque;
use crate::key_index::*;

pub struct KeyIndexedMap<K1, K2, K3, V> {
    idx: KeyIndex<K1, K2, K3, usize>,
    seq: Vec<(K1, K2, K3, V)>,
}

pub struct QueryIter<'a, K1, K2, K3, V> {
    vi: VecDeque<usize>,
    seq: &'a Vec<(K1, K2, K3, V)>,
}

impl<K1, K2, K3, V> Default for KeyIndexedMap<K1, K2, K3, V> {
    fn default() -> Self {
        KeyIndexedMap {
            idx: KeyIndex::default(),
            seq: Vec::new(),
        }
    }
}

impl<K1, K2, K3, V> KeyIndexedMap<K1, K2, K3, V>
where
    K1: Ord + Copy,
    K2: Ord + Copy,
    K3: Ord + Copy,
{
    pub fn get_idx(&self) -> &KeyIndex<K1, K2, K3, usize> {
        &self.idx
    }

    pub fn get_seq(&self) -> &Vec<(K1, K2, K3, V)> {
        &self.seq
    }

    pub fn from_seq<S: IntoIterator<Item = (K1, K2, K3, V)>>(seq: S) -> Result<KeyIndexedMap<K1, K2, K3, V>, InsertError<K1, K2, K3, usize>> {
        let seq = seq.into_iter().collect();
        let idx = KeyIndex::from_seq(&seq)?;
        Ok(KeyIndexedMap { idx, seq })
    }

    pub fn into_parts(self) -> (KeyIndex<K1, K2, K3, usize>, Vec<(K1, K2, K3, V)>) {
        let KeyIndexedMap { idx, seq } = self;
        (idx, seq)
    }

    pub fn insert(&mut self, k1: K1, k2: K2, k3: K3, v: V) -> Result<(), InsertError<K1, K2, K3, usize>> {
        let vi = self.seq.len();
        self.seq.push((k1, k2, k3, v));
        self.idx.insert(k1, k2, k3, vi)
    }

    pub fn query(&mut self, query: Query<K1, K2, K3>) -> QueryIter<K1, K2, K3, V> {
        QueryIter {
            vi: self.idx.query(query).into_iter().collect(),
            seq: &self.seq,
        }
    }

    pub fn get(&self, k1: K1, k2: K2, k3: K3) -> Option<&(K1, K2, K3, V)> {
        let query: Vec<_> = self.idx.query(Query { k1: Some(k1), k2: Some(k2), k3: Some(k3) }).into_iter().collect();
        if query.len() == 1 {
            let [res]: [_; 1] = query.try_into().unwrap();
            let res = self.seq.get(res).unwrap();
            Some(res)
        } else if query.is_empty() {
            None
        } else {
            unreachable!()
        }
    }
}

impl<'a, K1, K2, K3, V> Iterator for QueryIter<'a, K1, K2, K3, V> {
    type Item = &'a (K1, K2, K3, V);

    fn next(&mut self) -> Option<Self::Item> {
        let vi = self.vi.pop_front()?;
        Some(self.seq.get(vi).unwrap())
    }
}