From 0b9375688e2427a6e9521b2e81af8be367dc5db4 Mon Sep 17 00:00:00 2001
From: Georgios Konstantopoulos <me@gakonst.com>
Date: Sun, 17 Jan 2021 14:10:30 +0200
Subject: [PATCH] fix: switch between units correctly (#170)

---
 ethers-core/src/utils/mod.rs   | 36 ++++++++++++++++++++----
 ethers-core/src/utils/units.rs | 51 ++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 6 deletions(-)
 create mode 100644 ethers-core/src/utils/units.rs

diff --git a/ethers-core/src/utils/mod.rs b/ethers-core/src/utils/mod.rs
index 4c613e62..5c784526 100644
--- a/ethers-core/src/utils/mod.rs
+++ b/ethers-core/src/utils/mod.rs
@@ -20,6 +20,9 @@ pub use solc::{CompiledContract, Solc};
 mod hash;
 pub use hash::{hash_message, id, keccak256, serialize};
 
+mod units;
+pub use units::Units;
+
 /// Re-export RLP
 pub use rlp;
 
@@ -38,9 +41,11 @@ pub fn format_ether<T: Into<U256>>(amount: T) -> U256 {
     amount.into() / WEI_IN_ETHER
 }
 
-/// Divides with the number of decimals
-pub fn format_units<T: Into<U256>>(amount: T, decimals: usize) -> U256 {
-    amount.into() / decimals
+/// Divides the provided amount with 10^{units} provided.
+pub fn format_units<T: Into<U256>, K: Into<Units>>(amount: T, units: K) -> U256 {
+    let units = units.into();
+    let amount = amount.into();
+    amount / 10u64.pow(units.as_num())
 }
 
 /// Converts the input to a U256 and converts from Ether to Wei.
@@ -59,12 +64,13 @@ where
     Ok(eth.try_into()? * WEI_IN_ETHER)
 }
 
-/// Multiplies with the number of decimals
-pub fn parse_units<S>(eth: S, decimals: usize) -> Result<U256, S::Error>
+/// Multiplies the provided amount with 10^{units} provided.
+pub fn parse_units<S, K>(amount: S, units: K) -> Result<U256, S::Error>
 where
     S: TryInto<U256>,
+    K: Into<Units>,
 {
-    Ok(eth.try_into()? * decimals)
+    Ok(amount.try_into()? * 10u64.pow(units.into().as_num()))
 }
 
 /// The address for an Ethereum contract is deterministically computed from the
@@ -166,6 +172,24 @@ mod tests {
         assert_eq!(WEI_IN_ETHER.as_u64(), 1e18 as u64);
     }
 
+    #[test]
+    fn test_format_units() {
+        let gwei_in_ether = format_units(WEI_IN_ETHER, 9);
+        assert_eq!(gwei_in_ether.as_u64(), 1e9 as u64);
+
+        let eth = format_units(WEI_IN_ETHER, "ether");
+        assert_eq!(eth.as_u64(), 1);
+    }
+
+    #[test]
+    fn test_parse_units() {
+        let gwei = parse_units(1, 9).unwrap();
+        assert_eq!(gwei.as_u64(), 1e9 as u64);
+
+        let eth = parse_units(1, "ether").unwrap();
+        assert_eq!(eth, WEI_IN_ETHER);
+    }
+
     #[test]
     fn addr_checksum() {
         let addr_list = vec![
diff --git a/ethers-core/src/utils/units.rs b/ethers-core/src/utils/units.rs
new file mode 100644
index 00000000..1102d338
--- /dev/null
+++ b/ethers-core/src/utils/units.rs
@@ -0,0 +1,51 @@
+/// Common Ethereum unit types.
+pub enum Units {
+    /// Ether corresponds to 1e18 Wei
+    Ether,
+    /// Gwei corresponds to 1e9 Wei
+    Gwei,
+    /// Wei corresponds to 1 Wei
+    Wei,
+    /// Use this for other less frequent unit sizes
+    Other(u32),
+}
+
+impl Units {
+    pub fn as_num(&self) -> u32 {
+        match self {
+            Units::Ether => 18,
+            Units::Gwei => 9,
+            Units::Wei => 1,
+            Units::Other(inner) => *inner,
+        }
+    }
+}
+
+impl From<u32> for Units {
+    fn from(src: u32) -> Self {
+        Units::Other(src)
+    }
+}
+
+impl From<i32> for Units {
+    fn from(src: i32) -> Self {
+        Units::Other(src as u32)
+    }
+}
+
+impl From<usize> for Units {
+    fn from(src: usize) -> Self {
+        Units::Other(src as u32)
+    }
+}
+
+impl From<&str> for Units {
+    fn from(src: &str) -> Self {
+        match src.to_lowercase().as_str() {
+            "ether" => Units::Ether,
+            "gwei" => Units::Gwei,
+            "wei" => Units::Wei,
+            _ => panic!("unrecognized units"),
+        }
+    }
+}