From 400a22ebe5541f64ce68ad61ba2d2dad768a7e8d Mon Sep 17 00:00:00 2001
From: Wei Tang <accounts@that.world>
Date: Wed, 19 Sep 2018 16:00:50 +0800
Subject: [PATCH] Implement codec for U256 (#761)

* Implement codec for U256

* Use little endian and add tests
---
 substrate/core/primitives/src/hasher.rs |  2 +-
 substrate/core/primitives/src/uint.rs   | 46 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/substrate/core/primitives/src/hasher.rs b/substrate/core/primitives/src/hasher.rs
index d4dd39b855..7b5a725b77 100644
--- a/substrate/core/primitives/src/hasher.rs
+++ b/substrate/core/primitives/src/hasher.rs
@@ -45,7 +45,7 @@ pub mod blake2 {
 	impl Hasher for Blake2Hasher {
 		type Out = H256;
 		type StdHasher = PlainHasher;
-		const LENGTH:usize = 32;
+		const LENGTH: usize = 32;
 		fn hash(x: &[u8]) -> Self::Out {
 			blake2_256(x).into()
 		}
diff --git a/substrate/core/primitives/src/uint.rs b/substrate/core/primitives/src/uint.rs
index af3278ef01..32ba192cf0 100644
--- a/substrate/core/primitives/src/uint.rs
+++ b/substrate/core/primitives/src/uint.rs
@@ -43,12 +43,33 @@ macro_rules! impl_serde {
 	}
 }
 
+macro_rules! impl_codec {
+	($name: ident, $len: expr) => {
+		impl ::codec::Encode for $name {
+			fn using_encoded<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
+				let mut bytes = [0u8; $len * 8];
+				self.to_little_endian(&mut bytes);
+				bytes.using_encoded(f)
+			}
+		}
+
+		impl ::codec::Decode for $name {
+			fn decode<I: ::codec::Input>(input: &mut I) -> Option<Self> {
+				<[u8; $len * 8] as ::codec::Decode>::decode(input)
+					.map(|b| $name::from_little_endian(&b))
+			}
+		}
+	}
+}
+
 construct_uint!(U256, 4);
 impl_serde!(U256, 4);
+impl_codec!(U256, 4);
 
 #[cfg(test)]
 mod tests {
 	use super::*;
+	use codec::{Encode, Decode};
 	use substrate_serializer as ser;
 
 	macro_rules! test {
@@ -86,6 +107,31 @@ mod tests {
 
 	test!(U256, test_u256);
 
+	#[test]
+	fn test_u256_codec() {
+		let res1 = vec![120, 0, 0, 0, 0, 0, 0, 0,
+						0, 0, 0, 0, 0, 0, 0, 0,
+						0, 0, 0, 0, 0, 0, 0, 0,
+						0, 0, 0, 0, 0, 0, 0, 0];
+		let res2 = vec![0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+						0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+						0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+						0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff];
+
+		assert_eq!(
+			U256::from(120).encode(),
+			res1);
+		assert_eq!(
+			U256::max_value().encode(),
+			res2);
+		assert_eq!(
+			U256::decode(&mut &res1[..]),
+			Some(U256::from(120)));
+		assert_eq!(
+			U256::decode(&mut &res2[..]),
+			Some(U256::max_value()));
+	}
+
 	#[test]
 	fn test_large_values() {
 		assert_eq!(