From 51bb5f0f709422db4b959c173117ce77da098a95 Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Thu, 15 Mar 2018 12:53:36 -0600
Subject: [PATCH] Implement UInt32::shr() for SHA256.

---
 src/circuit/uint32.rs | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/src/circuit/uint32.rs b/src/circuit/uint32.rs
index 4714724..ad2dc84 100644
--- a/src/circuit/uint32.rs
+++ b/src/circuit/uint32.rs
@@ -155,6 +155,25 @@ impl UInt32 {
         }
     }
 
+    pub fn shr(&self, by: usize) -> Self {
+        let by = by % 32;
+
+        let fill = Boolean::constant(false);
+
+        let new_bits = self.bits
+                           .iter() // The bits are least significant first
+                           .skip(by) // Skip the bits that will be lost during the shift
+                           .chain(Some(&fill).into_iter().cycle()) // Rest will be zeros
+                           .take(32) // Only 32 bits needed!
+                           .cloned()
+                           .collect();
+
+        UInt32 {
+            bits: new_bits,
+            value: self.value.map(|v| v >> by as u32)
+        }
+    }
+
     /// XOR this `UInt32` with another `UInt32`
     pub fn xor<E, CS>(
         &self,
@@ -483,6 +502,7 @@ mod test {
 
         for i in 0..32 {
             let b = a.rotr(i);
+            assert_eq!(a.bits.len(), b.bits.len());
 
             assert!(b.value.unwrap() == num);
 
@@ -501,4 +521,24 @@ mod test {
             num = num.rotate_right(1);
         }
     }
+
+    #[test]
+    fn test_uint32_shr() {
+        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+        for _ in 0..50 {
+            for i in 0..60 {
+                let num = rng.gen();
+                let a = UInt32::constant(num).shr(i);
+                let b = UInt32::constant(num >> i);
+
+                assert_eq!(a.value.unwrap(), num >> i);
+
+                assert_eq!(a.bits.len(), b.bits.len());
+                for (a, b) in a.bits.iter().zip(b.bits.iter()) {
+                    assert_eq!(a.get_value().unwrap(), b.get_value().unwrap());
+                }
+            }
+        }
+    }
 }