Implement (de)serialization for BTreeMap same as HashMap

[rust-lightning] / lightning / src / util / ser.rs

diff --git a/lightning/src/util/ser.rs b/lightning/src/util/ser.rs
index a656604d458c492c17e8fc9a206c20afe4faa5ef..454c4990e9e5db9f8a17bbba04dbb7d0bfb1f191 100644 (file)
--- a/lightning/src/util/ser.rs
+++ b/lightning/src/util/ser.rs
@@ -22,6 +22,8 @@ use core::cmp;
 use core::convert::TryFrom;
 use core::ops::Deref;
 
+use alloc::collections::BTreeMap;
+
 use bitcoin::secp256k1::{PublicKey, SecretKey};
 use bitcoin::secp256k1::constants::{PUBLIC_KEY_SIZE, SECRET_KEY_SIZE, COMPACT_SIGNATURE_SIZE, SCHNORR_SIGNATURE_SIZE};
 use bitcoin::secp256k1::ecdsa;
@@ -588,43 +590,47 @@ impl<'a, T> From<&'a Vec<T>> for WithoutLength<&'a Vec<T>> {
        fn from(v: &'a Vec<T>) -> Self { Self(v) }
 }
 
-// HashMap
-impl<K, V> Writeable for HashMap<K, V>
-       where K: Writeable + Eq + Hash,
-             V: Writeable
-{
-       #[inline]
-       fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
-       (self.len() as u16).write(w)?;
-               for (key, value) in self.iter() {
-                       key.write(w)?;
-                       value.write(w)?;
+macro_rules! impl_for_map {
+       ($ty: ident, $keybound: ident, $constr: expr) => {
+               impl<K, V> Writeable for $ty<K, V>
+                       where K: Writeable + Eq + $keybound, V: Writeable
+               {
+                       #[inline]
+                       fn write<W: Writer>(&self, w: &mut W) -> Result<(), io::Error> {
+                               (self.len() as u16).write(w)?;
+                               for (key, value) in self.iter() {
+                                       key.write(w)?;
+                                       value.write(w)?;
+                               }
+                               Ok(())
+                       }
                }
-               Ok(())
-       }
-}
 
-impl<K, V> Readable for HashMap<K, V>
-       where K: Readable + Eq + Hash,
-             V: MaybeReadable
-{
-       #[inline]
-       fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
-               let len: u16 = Readable::read(r)?;
-               let mut ret = HashMap::with_capacity(len as usize);
-               for _ in 0..len {
-                       let k = K::read(r)?;
-                       let v_opt = V::read(r)?;
-                       if let Some(v) = v_opt {
-                               if ret.insert(k, v).is_some() {
-                                       return Err(DecodeError::InvalidValue);
+               impl<K, V> Readable for $ty<K, V>
+                       where K: Readable + Eq + $keybound, V: MaybeReadable
+               {
+                       #[inline]
+                       fn read<R: Read>(r: &mut R) -> Result<Self, DecodeError> {
+                               let len: u16 = Readable::read(r)?;
+                               let mut ret = $constr(len as usize);
+                               for _ in 0..len {
+                                       let k = K::read(r)?;
+                                       let v_opt = V::read(r)?;
+                                       if let Some(v) = v_opt {
+                                               if ret.insert(k, v).is_some() {
+                                                       return Err(DecodeError::InvalidValue);
+                                               }
+                                       }
                                }
+                               Ok(ret)
                        }
                }
-               Ok(ret)
        }
 }
 
+impl_for_map!(BTreeMap, Ord, |_| BTreeMap::new());
+impl_for_map!(HashMap, Hash, |len| HashMap::with_capacity(len));
+
 // HashSet
 impl<T> Writeable for HashSet<T>
 where T: Writeable + Eq + Hash