primitives - BalancesMap - simplify and make it work with newtypes

Author: elpiel

primitives/src/balances_map.rs

@@ -1,19 +1,29 @@
 use std::collections::BTreeMap;
 
 use crate::BigNum;
+use serde_hex::{SerHexSeq, StrictPfx};
 use std::collections::btree_map::{Entry, Iter, Values};
 
-use serde::{ser::SerializeMap, Deserialize, Serialize, Serializer};
+use serde::{Deserialize, Serialize};
 use std::iter::FromIterator;
+use std::ops::Deref;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialOrd, Ord, PartialEq, Eq)]
+#[serde(transparent)]
+pub struct BalancesKey(#[serde(with = "SerHexSeq::<StrictPfx>")] Vec<u8>);
+impl Deref for BalancesKey {
+    type Target = Vec<u8>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
 
-type BalancesKey = String;
 type BalancesValue = BigNum;
 
+#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
 #[serde(transparent)]
-#[derive(Clone, Debug, Deserialize, Default)]
-pub struct BalancesMap(
-    #[serde(serialize_with = "serialize_balances_map")] BTreeMap<BalancesKey, BalancesValue>,
-);
+pub struct BalancesMap(BTreeMap<BalancesKey, BalancesValue>);
 
 impl BalancesMap {
     pub fn iter(&self) -> Iter<'_, BalancesKey, BalancesValue> {
@@ -37,24 +47,43 @@ impl BalancesMap {
     }
 }
 
-impl FromIterator<(BalancesKey, BalancesValue)> for BalancesMap {
-    fn from_iter<I: IntoIterator<Item = (BalancesKey, BalancesValue)>>(iter: I) -> Self {
+impl<K: AsRef<[u8]>> FromIterator<(K, BalancesValue)> for BalancesMap {
+    fn from_iter<I: IntoIterator<Item = (K, BalancesValue)>>(iter: I) -> Self {
         // @TODO: Is there better way to do this?
-        let btree_map: BTreeMap<BalancesKey, BalancesValue> = iter.into_iter().collect();
+        let btree_map: BTreeMap<BalancesKey, BalancesValue> = iter
+            .into_iter()
+            .map(|(k, v)| (BalancesKey(k.as_ref().to_vec()), v))
+            .collect();
 
         BalancesMap(btree_map)
     }
 }
 
-impl Serialize for BalancesMap {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        let mut map = serializer.serialize_map(Some(self.0.len()))?;
-        for (key, big_num) in self.0.iter() {
-            map.serialize_entry(&key.to_lowercase(), big_num)?;
-        }
-        map.end()
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::BigNum;
+
+    #[test]
+    fn test_map() {
+        let data = vec![
+            (
+                "0xce07CbB7e054514D590a0262C93070D838bFBA2e".to_string(),
+                BigNum::from(50_u64),
+            ),
+            (
+                "0x061d5e2a67d0a9a10f1c732bca12a676d83f79663a396f7d87b3e30b9b411088".to_string(),
+                BigNum::from(100_u64),
+            ),
+        ];
+
+        let balances_map: BalancesMap = data.into_iter().collect();
+
+        let actual_json = serde_json::to_string(&balances_map).expect("Should serialize it");
+
+        let balances_map_from_json: BalancesMap =
+            serde_json::from_str(&string).expect("Should deserialize it");
+
+        assert_eq!(balances_map, balances_map_from_json);
     }
 }