wen-restart: Find heaviest fork (solana-labs#183)

* Pass the final result of LastVotedForkSlots aggregation to next
stage and find the heaviest fork we will Gossip to others.

* Change comments.

* Small fixes to address PR comments.

* Move correctness proof to SIMD.

* Fix a broken merge.

* Use blockstore to check parent slot of any block in FindHeaviestFork

* Change error message.

* Add special message when first slot in the list doesn't link to root.

wen-restart/proto/wen_restart.proto

@@ -20,10 +20,23 @@ message LastVotedForkSlotsRecord {
 
 message LastVotedForkSlotsAggregateRecord {
     map<string, LastVotedForkSlotsRecord> received = 1;
+    optional LastVotedForkSlotsAggregateFinal final_result = 2;
+}
+
+message LastVotedForkSlotsAggregateFinal {
+    map<uint64, uint64> slots_stake_map = 1;
+    uint64 total_active_stake = 2;
+}
+
+message HeaviestFork {
+    uint64 slot = 1;
+    string bankhash = 2;
+    uint64 total_active_stake = 3;
 }
 
 message WenRestartProgress {
     State state = 1;
     optional LastVotedForkSlotsRecord my_last_voted_fork_slots = 2;
     optional LastVotedForkSlotsAggregateRecord last_voted_fork_slots_aggregate = 3;
+    optional HeaviestFork my_heaviest_fork = 4;
 }

wen-restart/src/last_voted_fork_slots_aggregate.rs

@@ -22,6 +22,12 @@ pub struct LastVotedForkSlotsAggregate {
     slots_to_repair: HashSet<Slot>,
 }
 
+#[derive(Clone, Debug, PartialEq)]
+pub struct LastVotedForkSlotsFinalResult {
+    pub slots_stake_map: HashMap<Slot, u64>,
+    pub total_active_stake: u64,
+}
+
 impl LastVotedForkSlotsAggregate {
     pub(crate) fn new(
         root_slot: Slot,
@@ -35,7 +41,7 @@ impl LastVotedForkSlotsAggregate {
         active_peers.insert(*my_pubkey);
         let mut slots_stake_map = HashMap::new();
         for slot in last_voted_fork_slots {
-            if slot > &root_slot {
+            if slot >= &root_slot {
                 slots_stake_map.insert(*slot, sender_stake);
             }
         }
@@ -137,6 +143,21 @@ impl LastVotedForkSlotsAggregate {
     pub(crate) fn slots_to_repair_iter(&self) -> impl Iterator<Item = &Slot> {
         self.slots_to_repair.iter()
     }
+
+    // TODO(wen): use better epoch stake and add a test later.
+    fn total_active_stake(&self) -> u64 {
+        self.active_peers.iter().fold(0, |sum: u64, pubkey| {
+            sum.saturating_add(Self::validator_stake(&self.epoch_stakes, pubkey))
+        })
+    }
+
+    pub(crate) fn get_final_result(self) -> LastVotedForkSlotsFinalResult {
+        let total_active_stake = self.total_active_stake();
+        LastVotedForkSlotsFinalResult {
+            slots_stake_map: self.slots_stake_map,
+            total_active_stake,
+        }
+    }
 }
 
 #[cfg(test)]