refactor/attestation-target (#79)

* moved attestationTargets function to array extension * renaming * updated * Update Types.swift * fix indent * cleaned up votes function. added test * added tests * fixed
yeeth · Mar 24, 2019 · 1f95768 · 1f95768
1 parent d6e6418
commit 1f95768
Show file tree

Hide file tree

Showing 8 changed files with 93 additions and 28 deletions.
diff --git a/Sources/BeaconChain/DataStructures/Blocks/BeaconBlock.swift b/Sources/BeaconChain/DataStructures/Blocks/BeaconBlock.swift
@@ -1,6 +1,6 @@
 import Foundation
 
-struct BeaconBlock: Equatable {
+public struct BeaconBlock: Equatable {
     let slot: UInt64
     let parentRoot: Data
     let stateRoot: Data

diff --git a/Sources/BeaconChain/DataStructures/Transactions/Attestations/Attestation.swift b/Sources/BeaconChain/DataStructures/Transactions/Attestations/Attestation.swift
@@ -1,6 +1,6 @@
 import Foundation
 
-struct Attestation: Equatable {
+public struct Attestation: Equatable {
     let aggregationBitfield: Data
     let data: AttestationData
     let custodyBitfield: Data

diff --git a/Sources/BeaconChain/Extensions/Array+AttestationTarget.swift b/Sources/BeaconChain/Extensions/Array+AttestationTarget.swift
@@ -0,0 +1,17 @@
+import Foundation
+
+extension Array where Element == AttestationTarget {
+
+    func votes(store: Store, state: BeaconState, block: BeaconBlock) -> UInt64 {
+        return compactMap {
+            (index, target) in
+
+            guard store.ancestor(block: target, slot: target.slot) == block else {
+                return nil
+            }
+
+            return BeaconChain.getEffectiveBalance(state: state, index: index) / FORK_CHOICE_BALANCE_INCREMENT
+        }
+        .reduce(0, +)
+    }
+}
diff --git a/Sources/BeaconChain/ForkChoice/LMDGhost.swift b/Sources/BeaconChain/ForkChoice/LMDGhost.swift
@@ -3,7 +3,7 @@ import Foundation
 class LMDGhost: ForkChoice {
 
     func execute(store: Store, startState: BeaconState, startBlock: BeaconBlock) -> BeaconBlock {
-        let attestationTargets = startState.validatorRegistry.activeIndices(epoch: startState.slot.toEpoch()).map {
+        let targets = startState.validatorRegistry.activeIndices(epoch: startState.slot.toEpoch()).map {
             ($0, store.latestAttestationTarget(validator: $0))
         }
 
@@ -15,29 +15,8 @@ class LMDGhost: ForkChoice {
             }
 
             head = children.max {
-                voteCount(store: store, state: startState, block: $0, attestationTargets: attestationTargets) < voteCount(store: store, state: startState, block: $1, attestationTargets: attestationTargets)
+                targets.votes(store: store, state: startState, block: $0) < targets.votes(store: store, state: startState, block: $1)
             }!
         }
     }
-
-    private func voteCount(
-        store: Store,
-        state: BeaconState,
-        block: BeaconBlock,
-        attestationTargets: [(ValidatorIndex, BeaconBlock)]
-    ) -> UInt64 {
-        return attestationTargets.compactMap {
-            (index, target) in
-            guard let ancestor = store.ancestor(block: target, slot: index) else {
-                return nil
-            }
-
-            if ancestor == block {
-                return BeaconChain.getEffectiveBalance(state: state, index: index) / FORK_CHOICE_BALANCE_INCREMENT
-            }
-
-            return nil
-        }
-        .reduce(0, +)
-    }
 }
diff --git a/Sources/BeaconChain/Store/Store.swift b/Sources/BeaconChain/Store/Store.swift
@@ -1,6 +1,6 @@
 import Foundation
 
-protocol Store {
+public protocol Store {
 
     func parent(_ block: BeaconBlock) -> BeaconBlock
     func children(_ block: BeaconBlock) -> [BeaconBlock]

diff --git a/Sources/BeaconChain/Types.swift b/Sources/BeaconChain/Types.swift
@@ -1,10 +1,12 @@
 import Foundation
 
-typealias Slot = UInt64
+public typealias Slot = UInt64
 typealias Epoch = UInt64
 typealias Shard = UInt64
-typealias ValidatorIndex = UInt64
+public typealias ValidatorIndex = UInt64
 typealias Gwei = UInt64
 typealias Bytes32 = Data // @todo needs to be 32 fixed length data
 typealias BLSPubkey = Data // @todo needs to be 48 fixed length data
 typealias BLSSignature = Data // @todo needs to be 96 fixed length data
+
+typealias AttestationTarget = (ValidatorIndex, BeaconBlock)
diff --git a/Tests/BeaconChainTests/Extensions/Array+AttestationTargetTests.swift b/Tests/BeaconChainTests/Extensions/Array+AttestationTargetTests.swift
@@ -0,0 +1,47 @@
+import XCTest
+@testable import BeaconChain
+
+final class ArrayAttestationTargetTests: XCTestCase {
+
+    func testVotes() {
+        var state = BeaconChain.genesisState(
+            genesisTime: 10,
+            latestEth1Data: Eth1Data(depositRoot: Data(count: 32), blockHash: Data(count: 32)),
+            depositLength: 0
+        )
+
+        state.validatorBalances.append(32000000000)
+        state.validatorBalances.append(32000000000)
+        state.validatorBalances.append(32000000000)
+
+        let block = BeaconBlock(
+            slot: GENESIS_SLOT,
+            parentRoot: ZERO_HASH,
+            stateRoot: ZERO_HASH,
+            randaoReveal: ZERO_HASH,
+            eth1Data: state.latestEth1Data,
+            body: BeaconBlockBody(
+                proposerSlashings: [ProposerSlashing](),
+                attesterSlashings: [AttesterSlashing](),
+                attestations: [Attestation](),
+                deposits: [Deposit](),
+                voluntaryExits: [VoluntaryExit](),
+                transfers: [Transfer]()
+            ),
+            signature: ZERO_HASH
+        )
+
+        let store = MockStore()
+
+        var badBlock = block
+        badBlock.signature = EMPTY_SIGNATURE
+
+        var targets = [AttestationTarget]()
+        targets.append((0, block))
+        targets.append((1, badBlock))
+        targets.append((2, block))
+
+        XCTAssertEqual(64, targets.votes(store: store, state: state, block: block))
+    }
+
+}
diff --git a/Tests/BeaconChainTests/Mocks/MockStore.swift b/Tests/BeaconChainTests/Mocks/MockStore.swift
@@ -0,0 +1,20 @@
+import BeaconChain
+
+class MockStore: Store {
+
+    func parent(_ block: BeaconBlock) -> BeaconBlock {
+        return block
+    }
+
+    func children(_ block: BeaconBlock) -> [BeaconBlock] {
+        return [block]
+    }
+
+    func latestAttestation(validator: ValidatorIndex) -> Attestation {
+        fatalError()
+    }
+
+    func latestAttestationTarget(validator: ValidatorIndex) -> BeaconBlock {
+        fatalError()
+    }
+}