From 47048f15dc35e0e5d0b24fc4e9bca50472cbc292 Mon Sep 17 00:00:00 2001
From: Hugo Rosenkranz-Costa <hugo.rosenkranz@cosmian.com>
Date: Fri, 10 Nov 2023 17:45:44 +0100
Subject: [PATCH] feat: add `LinkedVec` data struct

---
 src/data_struct/linked_hashmap.rs | 20 ++++---
 src/data_struct/linked_vec.rs     | 97 +++++++++++++++++++++++++++++++
 src/data_struct/mod.rs            |  2 +
 3 files changed, 111 insertions(+), 8 deletions(-)
 create mode 100644 src/data_struct/linked_vec.rs

diff --git a/src/data_struct/linked_hashmap.rs b/src/data_struct/linked_hashmap.rs
index 01cc32f7..75fc20bf 100644
--- a/src/data_struct/linked_hashmap.rs
+++ b/src/data_struct/linked_hashmap.rs
@@ -46,13 +46,13 @@ where
     V: Clone,
     K: Hash + Eq + PartialEq + Clone + PartialOrd + Ord,
 {
-    type Item = &'a V;
+    type Item = (&'a K, &'a V);
 
     fn next(&mut self) -> Option<Self::Item> {
         self.current_key.and_then(|key| {
             self.lhm.get_link_entry(key).map(|entry| {
                 self.current_key = entry.next_key.as_ref();
-                entry.get_value()
+                (key, entry.get_value())
             })
         })
     }
@@ -88,10 +88,7 @@ where
     }
 
     /// Internal pattern matching
-    fn expected_entry(
-        &'_ mut self,
-        key: K,
-    ) -> Result<OccupiedEntry<'_, K, LinkedEntry<K, V>>, Error> {
+    fn expected_entry(&mut self, key: K) -> Result<OccupiedEntry<K, LinkedEntry<K, V>>, Error> {
         match self.map.entry(key) {
             Entry::Occupied(e) => Ok(e),
             Entry::Vacant(_) => Err(Error::KeyError("Key not found".to_string())),
@@ -145,7 +142,7 @@ where
         self.map.get(key).map(LinkedEntry::get_value)
     }
 
-    pub fn iter_link<'a>(&'a self, key: &'a K) -> impl Iterator<Item = &V> + 'a {
+    pub fn iter_link<'a>(&'a self, key: &'a K) -> impl Iterator<Item = (&K, &V)> + 'a {
         LinkedHashMapIterator::<'a, K, V> {
             lhm: self,
             current_key: Some(key),
@@ -189,7 +186,14 @@ fn test_linked_hashmap() -> Result<(), Error> {
     assert_eq!(res.len(), 5);
 
     let res: Vec<_> = lhm.iter_link(&1).collect();
-    assert_eq!(res, vec!["key1", "key11", "key111"]);
+    assert_eq!(
+        res,
+        vec![
+            (&1, &"key1".to_string()),
+            (&11, &"key11".to_string()),
+            (&111, &"key111".to_string())
+        ]
+    );
 
     lhm.pop_chain(1)?;
     let res: Vec<_> = lhm.iter().collect();
diff --git a/src/data_struct/linked_vec.rs b/src/data_struct/linked_vec.rs
new file mode 100644
index 00000000..22344a94
--- /dev/null
+++ b/src/data_struct/linked_vec.rs
@@ -0,0 +1,97 @@
+use std::collections::LinkedList;
+
+use crate::Error;
+
+#[derive(Default)]
+pub struct LinkedVec<K, V>
+where
+    K: PartialEq + Eq,
+{
+    data: Vec<LinkedList<(K, V)>>,
+}
+
+impl<K, V> LinkedVec<K, V>
+where
+    K: PartialEq + Eq,
+{
+    pub fn new() -> Self {
+        Self { data: Vec::new() }
+    }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            data: Vec::with_capacity(capacity),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.data.iter().map(|list| list.len()).sum()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    pub fn insert_new_chain(
+        &mut self,
+        iterator: impl Iterator<Item = (K, V)>,
+    ) -> Result<(), Error> {
+        let mut new_list = LinkedList::new();
+        for (key, value) in iterator {
+            new_list.push_front((key, value));
+        }
+
+        self.data.push(new_list);
+
+        Ok(())
+    }
+
+    pub fn pop_chain(&mut self, chain_index: usize, stop_key: &K) -> Result<(), Error> {
+        let list = &mut self.data[chain_index];
+
+        while let Some((key, _)) = list.back() {
+            if key == stop_key {
+                return Ok(());
+            }
+            list.pop_back();
+        }
+        Err(Error::KeyError("Stop key was not found".to_string()))
+    }
+
+    pub fn iter_chain(&self, chain_index: usize) -> impl Iterator<Item = &(K, V)> {
+        self.data[chain_index].iter()
+    }
+}
+
+#[test]
+fn test_linked_vec() -> Result<(), Error> {
+    let mut lv: LinkedVec<u32, String> = LinkedVec::new();
+    assert_eq!(lv.len(), 0);
+
+    lv.insert_new_chain(
+        vec![
+            (1, "key1".to_string()),
+            (11, "key11".to_string()),
+            (11, "key111".to_string()),
+        ]
+        .into_iter(),
+    )?;
+    lv.insert_new_chain(vec![(2, "key2".to_string()), (22, "key22".to_string())].into_iter())?;
+    lv.insert_new_chain(vec![(3, "key3".to_string())].into_iter())?;
+
+    assert_eq!(lv.data.len(), 3);
+    assert_eq!(lv.len(), 6);
+
+    lv.pop_chain(0, &11)?;
+    assert_eq!(lv.len(), 5);
+
+    lv.pop_chain(1, &22)?;
+    assert_eq!(lv.len(), 4);
+
+    assert!(lv.pop_chain(2, &0).is_err()); // key 0 was never reached
+    assert_eq!(lv.len(), 3); // the chain 3 was completely removed
+
+    assert_eq!(lv.iter_chain(0).collect::<Vec<_>>().len(), 2);
+
+    Ok(())
+}
diff --git a/src/data_struct/mod.rs b/src/data_struct/mod.rs
index 54bd8bbb..804455ba 100644
--- a/src/data_struct/mod.rs
+++ b/src/data_struct/mod.rs
@@ -1,3 +1,5 @@
 pub mod linked_hashmap;
+pub mod linked_vec;
 
 use linked_hashmap::LinkedHashMap;
+use linked_vec::LinkedVec;