From 2277273caa49ac03cda2fe2e9fe6591ebcc8e713 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Wed, 29 Aug 2018 16:57:25 +1000
Subject: [PATCH] Fix get_new_shuffling, add functioning test.

---
 .../state/transition/validator_allocation.rs  | 125 ++++++++++++++----
 1 file changed, 102 insertions(+), 23 deletions(-)

diff --git a/lighthouse/state/transition/validator_allocation.rs b/lighthouse/state/transition/validator_allocation.rs
index eb1d49413..f425eded6 100644
--- a/lighthouse/state/transition/validator_allocation.rs
+++ b/lighthouse/state/transition/validator_allocation.rs
@@ -70,7 +70,7 @@ fn generate_cycle(
 
     let (committees_per_slot, slots_per_committee) = {
         if validator_count >= cycle_length * min_committee_size {
-            let committees_per_slot = validator_count / cycle_length / (min_committee_size * 2);
+            let committees_per_slot = validator_count / cycle_length / (min_committee_size * 2) + 1;
             let slots_per_committee = 1;
             (committees_per_slot, slots_per_committee)
         } else {
@@ -84,23 +84,24 @@ fn generate_cycle(
         }
     };
 
-    let mut cycle: DelegatedCycle = vec![];
-    let split_iter = validator_indices.split(|i| i % cycle_length == 0);
-    for (i, slot_indices) in split_iter.enumerate() {
-        let shard_id_start = crosslinking_shard_start * i * committees_per_slot / slots_per_committee;
+    let cycle = validator_indices
+        .chunks(validator_indices.len() / *cycle_length)
+        .enumerate()
+        .map(|(i, slot_indices)| {
+            let shard_id_start = crosslinking_shard_start + i * committees_per_slot / slots_per_committee;
+            return slot_indices
+                .chunks(slot_indices.len() / committees_per_slot)
+                .enumerate()
+                .map(|(j, shard_indices)| {
+                    return ShardAndCommittee{
+                        shard_id: ((shard_id_start + j) % shard_count) as u16,
+                        committee: shard_indices.to_vec(),
+                    }
+                })
+                .collect()
+        })
+        .collect();
 
-        let shard_iter = slot_indices.split(|i| i % committees_per_slot == 0);
-        let slot: DelegatedSlot = shard_iter
-            .enumerate()
-            .map(|(j, shard_indices)| {
-                ShardAndCommittee{
-                    shard_id: ((shard_id_start + j) % shard_count) as u16,
-                    committee: shard_indices.to_vec(),
-                }
-            })
-            .collect();
-        cycle.push(slot);
-    };
     Ok(cycle)
 }
 
@@ -114,18 +115,76 @@ mod tests {
         crosslinking_shard_start: &usize,
         cycle_length: &usize,
         min_committee_size: &usize)
-        -> Result<DelegatedCycle, TransitionError>
+        -> (Vec<usize>, Vec<usize>, Result<DelegatedCycle, TransitionError>)
     {
         let validator_indices = (0_usize..*validator_count).into_iter().collect();
         let shard_indices = (0_usize..*shard_count).into_iter().collect();
-        generate_cycle(
+        let result = generate_cycle(
             &validator_indices,
             &shard_indices,
             &crosslinking_shard_start,
             &cycle_length,
-            &min_committee_size)
+            &min_committee_size);
+        (validator_indices, shard_indices, result)
     }
 
+    #[allow(dead_code)]
+    fn print_cycle(cycle: &DelegatedCycle) {
+        cycle.iter()
+            .enumerate()
+            .for_each(|(i, slot)| {
+                println!("slot {:?}", &i);
+                slot.iter()
+                    .enumerate()
+                    .for_each(|(i, sac)| {
+                        println!("#{:?}\tshard_id={}\tcommittee.len()={}",
+                            &i, &sac.shard_id, &sac.committee.len())
+                    })
+            });
+    }
+
+    fn flatten_validators(cycle: &DelegatedCycle)
+        -> Vec<usize>
+    {
+        let mut flattened = vec![];
+        for slot in cycle.iter() {
+            for sac in slot.iter() {
+                for validator in sac.committee.iter() {
+                    flattened.push(*validator);
+                }
+            }
+        }
+        flattened
+    }
+
+    fn flatten_and_dedup_shards(cycle: &DelegatedCycle)
+        -> Vec<usize>
+    {
+        let mut flattened = vec![];
+        for slot in cycle.iter() {
+            for sac in slot.iter() {
+                flattened.push(sac.shard_id as usize);
+            }
+        }
+        flattened.dedup();
+        flattened
+    }
+
+    fn flatten_shards_in_slots(cycle: &DelegatedCycle)
+        -> Vec<Vec<usize>>
+    {
+        let mut shards_in_slots: Vec<Vec<usize>> = vec![];
+        for slot in cycle.iter() {
+            let mut shards: Vec<usize> = vec![];
+            for sac in slot.iter() {
+                shards.push(sac.shard_id as usize);
+            }
+            shards_in_slots.push(shards);
+        }
+        shards_in_slots
+    }
+
+
     #[test]
     fn test_generate_cycle() {
         let validator_count: usize = 100;
@@ -133,13 +192,33 @@ mod tests {
         let crosslinking_shard_start: usize = 0;
         let cycle_length: usize = 20;
         let min_committee_size: usize = 10;
-        let result = generate_cycle_helper(
+        let (validators, shards, result) = generate_cycle_helper(
             &validator_count,
             &shard_count,
             &crosslinking_shard_start,
             &cycle_length,
             &min_committee_size);
-        println!("{:?}", result);
-    }
+        let cycle = result.unwrap();
 
+        let assigned_validators = flatten_validators(&cycle);
+        let assigned_shards = flatten_and_dedup_shards(&cycle);
+        let shards_in_slots = flatten_shards_in_slots(&cycle);
+        assert_eq!(assigned_validators, validators, "Validator assignment incorrect");
+        assert_eq!(assigned_shards, shards, "Shard assignment incorrect");
+
+        let expected_shards_in_slots: Vec<Vec<usize>> = vec![
+            vec![0], vec![0],   // Each line is 2 slots..
+            vec![1], vec![1],
+            vec![2], vec![2],
+            vec![3], vec![3],
+            vec![4], vec![4],
+            vec![5], vec![5],
+            vec![6], vec![6],
+            vec![7], vec![7],
+            vec![8], vec![8],
+            vec![9], vec![9],
+        ];
+        // assert!(compare_shards_in_slots(&cycle, &expected_shards_in_slots));
+        assert_eq!(expected_shards_in_slots, shards_in_slots, "Shard assignment incorrect.")
+    }
 }