From a322e0f1b3d2484b4728f6495d9dc09edd789b08 Mon Sep 17 00:00:00 2001
From: Victor Koenders <gpg@trangar.com>
Date: Sat, 16 Oct 2021 12:21:52 +0200
Subject: [PATCH] Added support for BinaryHeap, BTreeMap, BTreeSet and VecDeque

---
 src/features/impl_alloc.rs | 118 ++++++++++++++++++++++++++++++++++++-
 tests/alloc.rs             |  17 ++++++
 2 files changed, 134 insertions(+), 1 deletion(-)

diff --git a/src/features/impl_alloc.rs b/src/features/impl_alloc.rs
index 53fa85b..30e707a 100644
--- a/src/features/impl_alloc.rs
+++ b/src/features/impl_alloc.rs
@@ -5,7 +5,7 @@ use crate::{
     error::{DecodeError, EncodeError},
     Config,
 };
-use alloc::{borrow::Cow, boxed::Box, rc::Rc, sync::Arc, vec::Vec};
+use alloc::{borrow::Cow, boxed::Box, collections::*, rc::Rc, sync::Arc, vec::Vec};
 
 #[derive(Default)]
 struct VecWriter {
@@ -35,6 +35,122 @@ pub fn encode_to_vec_with_config<E: enc::Encodeable, C: Config>(
     Ok(encoder.into_writer().inner)
 }
 
+impl<T> Decodable for BinaryHeap<T>
+where
+    T: Decodable + Ord,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let len = usize::decode(&mut decoder)?;
+        let mut map = BinaryHeap::with_capacity(len);
+        for _ in 0..len {
+            let key = T::decode(&mut decoder)?;
+            map.push(key);
+        }
+        Ok(map)
+    }
+}
+
+impl<T> Encodeable for BinaryHeap<T>
+where
+    T: Encodeable + Ord,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        self.len().encode(&mut encoder)?;
+        for val in self.iter() {
+            val.encode(&mut encoder)?;
+        }
+        Ok(())
+    }
+}
+
+impl<K, V> Decodable for BTreeMap<K, V>
+where
+    K: Decodable + Ord,
+    V: Decodable,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let len = usize::decode(&mut decoder)?;
+        let mut map = BTreeMap::new();
+        for _ in 0..len {
+            let key = K::decode(&mut decoder)?;
+            let value = V::decode(&mut decoder)?;
+            map.insert(key, value);
+        }
+        Ok(map)
+    }
+}
+
+impl<K, V> Encodeable for BTreeMap<K, V>
+where
+    K: Encodeable + Ord,
+    V: Encodeable,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        self.len().encode(&mut encoder)?;
+        for (key, val) in self.iter() {
+            key.encode(&mut encoder)?;
+            val.encode(&mut encoder)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T> Decodable for BTreeSet<T>
+where
+    T: Decodable + Ord,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let len = usize::decode(&mut decoder)?;
+        let mut map = BTreeSet::new();
+        for _ in 0..len {
+            let key = T::decode(&mut decoder)?;
+            map.insert(key);
+        }
+        Ok(map)
+    }
+}
+
+impl<T> Encodeable for BTreeSet<T>
+where
+    T: Encodeable + Ord,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        self.len().encode(&mut encoder)?;
+        for item in self.iter() {
+            item.encode(&mut encoder)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T> Decodable for VecDeque<T>
+where
+    T: Decodable,
+{
+    fn decode<D: Decode>(mut decoder: D) -> Result<Self, DecodeError> {
+        let len = usize::decode(&mut decoder)?;
+        let mut map = VecDeque::with_capacity(len);
+        for _ in 0..len {
+            let key = T::decode(&mut decoder)?;
+            map.push_back(key);
+        }
+        Ok(map)
+    }
+}
+
+impl<T> Encodeable for VecDeque<T>
+where
+    T: Encodeable,
+{
+    fn encode<E: Encode>(&self, mut encoder: E) -> Result<(), EncodeError> {
+        self.len().encode(&mut encoder)?;
+        for item in self.iter() {
+            item.encode(&mut encoder)?;
+        }
+        Ok(())
+    }
+}
+
 impl<T> Decodable for Vec<T>
 where
     T: Decodable,
diff --git a/tests/alloc.rs b/tests/alloc.rs
index 2118b65..f096b5f 100644
--- a/tests/alloc.rs
+++ b/tests/alloc.rs
@@ -5,6 +5,7 @@ extern crate alloc;
 mod utils;
 
 use alloc::borrow::Cow;
+use alloc::collections::*;
 use alloc::rc::Rc;
 use alloc::sync::Arc;
 use utils::the_same;
@@ -53,4 +54,20 @@ fn test_alloc_commons() {
     the_same(Cow::<u32>::Borrowed(&5));
     the_same(Rc::<u32>::new(5));
     the_same(Arc::<u32>::new(5));
+    the_same({
+        let mut map = BTreeMap::<u32, i32>::new();
+        map.insert(5, -5);
+        map
+    });
+    the_same({
+        let mut set = BTreeSet::<u32>::new();
+        set.insert(5);
+        set
+    });
+    the_same({
+        let mut set = VecDeque::<u32>::new();
+        set.push_back(15);
+        set.push_front(5);
+        set
+    });
 }