From ff3795dbf0d23a2e7d3c6b639083964ca1e08f25 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Wed, 11 Jul 2018 18:16:16 +1000
Subject: [PATCH] Update shuffling function

---
 src/state/config.rs           |  6 ++---
 src/state/state_transition.rs | 45 +++++++++++++++++++----------------
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/src/state/config.rs b/src/state/config.rs
index d55f435660..5cb378286e 100644
--- a/src/state/config.rs
+++ b/src/state/config.rs
@@ -1,13 +1,13 @@
 pub struct Config {
-    pub attester_count: u32,
-    pub max_validators: u32
+    pub attester_count: u64,
+    pub max_validators: u64
 }
 
 impl Config {
     pub fn standard() -> Self {
         Self {
             attester_count: 32,
-            max_validators: 2u32.pow(24)
+            max_validators: 2u64.pow(24)
         }
     }
 }
diff --git a/src/state/state_transition.rs b/src/state/state_transition.rs
index 6d2c9e0e8d..4c31175e86 100644
--- a/src/state/state_transition.rs
+++ b/src/state/state_transition.rs
@@ -8,52 +8,57 @@ use super::config::Config;
 const AGG_VOTE_MSG_SIZE: i32 = 2 + 32 + 32 + 8 + 8;
 
 // 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)
+fn get_shift_from_source(source: &[u8], offset: usize) -> usize {
+    (source[offset + 2] as usize) |
+        ((source[offset + 1] as usize) << 8) |
+        ((source[offset    ] as usize) << 16)
 }
 
 // Given entropy in the form of `seed`, return a shuffled list of validators
-// of size `validator_count` or `sample`.
+// indicies of size `validator_count` or `sample`.
 pub fn get_shuffling(
     seed: &Sha256Digest,
-    validator_count: &u32,
-    sample: &Option<u32>,
+    validator_count: &usize,
+    sample_size: &Option<usize>,
     config: &Config) 
-    -> Vec<u32>
+    -> Vec<usize>
 {
-    let max_validators = config.max_validators;
-    assert!(*validator_count <= max_validators);
+    assert!(*validator_count > 0);
+    let mut output: Vec<usize> = (0..*validator_count).collect();
+
+    assert!(*validator_count <= (config.max_validators as usize));
     
-    // 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,
+    // Use a reduced "sample_size" output range if specified
+    let output_range: &usize = match sample_size {
+        Some(x) => {
+            assert!(x <= validator_count,
+                    "sample_size should be <= validator_count");
+            x
+        },
         None => validator_count
     };
-    let mut output: Vec<u32> = (0..*validator_range).collect();
+
+    // Do the first blake hash round
     let mut source = Blake2s::new();
     source.input(&seed);
     
     let mut v = 0;
-    while v < *validator_range {
+    while v < *output_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;
+            let shuffled_position: usize = (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
+        // Re-hash the source (TODO: this does one extra hash, can be optimised)
         source = Blake2s::new();
         source.input(&current_source);
     }
-    output
+    output[0..*output_range].to_vec()
 }
 
 // Given an aggregate_vote and a crystallized_state,