From e73290dd5623f5257b6b9dd150cf2dba6e4ada88 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Tue, 10 Jul 2018 19:08:36 +1000
Subject: [PATCH] Add shuffling fn for state_transition

---
 src/state/config.rs           | 13 +++++
 src/state/mod.rs              |  2 +
 src/state/state_transition.rs | 94 +++++++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+)
 create mode 100644 src/state/config.rs
 create mode 100644 src/state/state_transition.rs

diff --git a/src/state/config.rs b/src/state/config.rs
new file mode 100644
index 000000000..d55f43566
--- /dev/null
+++ b/src/state/config.rs
@@ -0,0 +1,13 @@
+pub struct Config {
+    pub attester_count: u32,
+    pub max_validators: u32
+}
+
+impl Config {
+    pub fn standard() -> Self {
+        Self {
+            attester_count: 32,
+            max_validators: 2u32.pow(24)
+        }
+    }
+}
diff --git a/src/state/mod.rs b/src/state/mod.rs
index b4162b7f2..e10ef90d2 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -6,9 +6,11 @@ use super::utils;
 
 pub mod active_state;
 pub mod crystallized_state;
+pub mod config;
 pub mod aggregate_vote;
 pub mod block;
 pub mod crosslink_record;
 pub mod partial_crosslink_record;
 pub mod recent_proposer_record;
+pub mod state_transition;
 pub mod validator_record;
diff --git a/src/state/state_transition.rs b/src/state/state_transition.rs
new file mode 100644
index 000000000..c9b352940
--- /dev/null
+++ b/src/state/state_transition.rs
@@ -0,0 +1,94 @@
+use super::utils::types::{ Sha256Digest };
+use super::blake2::{ Blake2s, Digest };
+use super::config::Config;
+
+// Interprets a 3-byte slice from a [u8] as an integer.
+fn get_shift_from_source(source: &[u8], offset: usize) -> u32 {
+    (source[offset + 2] as u32) |
+        ((source[offset + 1] as u32) << 8) |
+        ((source[offset    ] as u32) << 16)
+}
+
+// Given entropy in the form of `seed`, return a shuffled list of validators
+// of size `validator_count` or `sample`.
+pub fn get_shuffling(
+    seed: Sha256Digest,
+    validator_count: u32,
+    sample: Option<u32>,
+    config: Config) 
+    -> Vec<u32>
+{
+    let max_validators = config.max_validators;
+    assert!(validator_count <= max_validators);
+    
+    // TODO: figure out why the Python implementation uses
+    // this `rand_max` var.
+    // let rand_max = max_validators - (max_validators % validator_count);
+    let validator_range = match sample {
+        Some(x) => x,
+        None => validator_count
+    };
+    let mut output: Vec<u32> = (0..validator_range).collect();
+    let mut source = Blake2s::new();
+    source.input(&seed);
+    
+    let mut v = 0;
+    while v < validator_range {
+        let current_source = source.result();
+        let mut source_offset = 0;
+        while source_offset < 30 {
+            let m = get_shift_from_source(&current_source, source_offset);
+            let shuffled_position = (m % (validator_count - v)) + v;
+            output.swap(v as usize, shuffled_position as usize);
+            v += 1;
+            if v >= validator_count { break; }
+            source_offset += 3;
+        }
+        // Re-hash the source
+        source = Blake2s::new();
+        source.input(&current_source);
+    }
+    output
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_shuffling_shift_fn() {
+        let mut x = get_shift_from_source(
+            &vec![0_u8, 0, 1],
+            0);
+        assert_eq!((x as u32), 1);
+
+        x = get_shift_from_source(
+            &vec![0_u8, 1, 1],
+            0);
+        assert_eq!(x, 257);
+        
+        x = get_shift_from_source(
+            &vec![1_u8, 1, 1],
+            0);
+        assert_eq!(x, 65793);
+        
+        x = get_shift_from_source(
+            &vec![255_u8, 1, 1],
+            0);
+        assert_eq!(x, 16711937);
+    }
+
+
+    #[test]
+    fn test_shuffling() {
+        let s = get_shuffling(
+            Sha256Digest::zero(),
+            10,
+            None,
+            Config::standard());
+        assert_eq!(s, 
+                   vec!(0, 9, 7, 6, 4, 1, 8, 5, 2, 3),
+                   "10 validator shuffle was not as expected");
+    }
+}