feature/epoch-extension (#55)

* added EpochNumber extension * removes getEntryExitEpoch * added test
yeeth · Feb 22, 2019 · bad5295 · bad5295
1 parent 25f194d
commit bad5295
Show file tree

Hide file tree

Showing 7 changed files with 41 additions and 21 deletions.
diff --git a/Sources/BeaconChain/BeaconChain.swift b/Sources/BeaconChain/BeaconChain.swift
@@ -31,14 +31,6 @@ extension BeaconChain {
     static func getCurrentEpoch(state: BeaconState) -> EpochNumber {
         return state.slot.toEpoch()
     }
-
-    static func getEpochStartSlot(_ epoch: EpochNumber) -> SlotNumber {
-        return epoch * EPOCH_LENGTH
-    }
-
-    static func getEntryExitEpoch(_ epoch: EpochNumber) -> EpochNumber {
-        return epoch + 1 + ENTRY_EXIT_DELAY
-    }
 }
 
 extension BeaconChain {
@@ -490,7 +482,7 @@ extension BeaconChain {
 extension BeaconChain {
 
     static func activateValidator(state: inout BeaconState, index: ValidatorIndex, genesis: Bool) {
-        state.validatorRegistry[Int(index)].activationEpoch = genesis ? GENESIS_EPOCH : getEntryExitEpoch(getCurrentEpoch(state: state))
+        state.validatorRegistry[Int(index)].activationEpoch = genesis ? GENESIS_EPOCH : getCurrentEpoch(state: state).entryExitEpoch()
     }
 
     static func initiateValidatorExit(state: inout BeaconState, index: ValidatorIndex) {
@@ -499,11 +491,11 @@ extension BeaconChain {
 
     static func exitValidator(state: inout BeaconState, index: ValidatorIndex) {
         var validator = state.validatorRegistry[Int(index)]
-        if validator.exitEpoch <= getEntryExitEpoch(getCurrentEpoch(state: state)) {
+        if validator.exitEpoch <= getCurrentEpoch(state: state).entryExitEpoch() {
             return
         }
 
-        validator.exitEpoch = getEntryExitEpoch(getCurrentEpoch(state: state))
+        validator.exitEpoch = getCurrentEpoch(state: state).entryExitEpoch()
         state.validatorRegistry[Int(index)] = validator
     }
 

diff --git a/Sources/BeaconChain/DataStructures/Operations/Attestations/AttestationData.swift b/Sources/BeaconChain/DataStructures/Operations/Attestations/AttestationData.swift
@@ -7,6 +7,6 @@ struct AttestationData: Equatable {
     let epochBoundaryRoot: Data
     let shardBlockRoot: Data
     let latestCrosslink: Crosslink
-    let justifiedEpoch: UInt64
+    let justifiedEpoch: EpochNumber
     let justifiedBlockRoot: Data
 }
diff --git a/Sources/BeaconChain/Extensions/EpochNumber.swift b/Sources/BeaconChain/Extensions/EpochNumber.swift
@@ -0,0 +1,14 @@
+import Foundation
+
+extension EpochNumber {
+
+
+    func startSlot() -> SlotNumber {
+        return self * EPOCH_LENGTH
+    }
+
+    func entryExitEpoch() -> EpochNumber {
+        return self + 1 + ENTRY_EXIT_DELAY
+    }
+
+}
diff --git a/Sources/BeaconChain/Extensions/SlotNumber.swift b/Sources/BeaconChain/Extensions/SlotNumber.swift
@@ -6,4 +6,4 @@ extension SlotNumber {
         return self / EPOCH_LENGTH
     }
 
-}
+}
diff --git a/Sources/BeaconChain/StateTransition.swift b/Sources/BeaconChain/StateTransition.swift
@@ -153,7 +153,7 @@ extension StateTransition {
 
             let e = attestation.data.justifiedEpoch >= BeaconChain.getCurrentEpoch(state: state) ? state.justifiedEpoch : state.previousJustifiedEpoch
             assert(attestation.data.justifiedEpoch == e)
-            assert(attestation.data.justifiedBlockRoot == BeaconChain.getBlockRoot(state: state, slot: BeaconChain.getEpochStartSlot(attestation.data.justifiedEpoch)))
+            assert(attestation.data.justifiedBlockRoot == BeaconChain.getBlockRoot(state: state, slot: attestation.data.justifiedEpoch.startSlot()))
 
             assert(
                 state.latestCrosslinks[Int(attestation.data.shard)] == attestation.data.latestCrosslink ||
@@ -260,7 +260,7 @@ extension StateTransition {
             let validator = state.validatorRegistry[Int(exit.validatorIndex)]
 
             let epoch = BeaconChain.getCurrentEpoch(state: state)
-            assert(validator.exitEpoch > BeaconChain.getEntryExitEpoch(epoch))
+            assert(validator.exitEpoch > epoch.entryExitEpoch())
             assert(epoch >= exit.epoch)
 
             let exitMessage = BeaconChain.hashTreeRoot(
@@ -310,7 +310,7 @@ extension StateTransition {
         }
 
         let currentEpochBoundryAttestations = currentEpochAttestations.filter {
-            $0.data.epochBoundaryRoot == BeaconChain.getBlockRoot(state: state, slot: BeaconChain.getEpochStartSlot(currentEpoch))
+            $0.data.epochBoundaryRoot == BeaconChain.getBlockRoot(state: state, slot: currentEpoch.startSlot())
                 && $0.data.justifiedEpoch == state.justifiedEpoch
         }
 
@@ -341,7 +341,7 @@ extension StateTransition {
         let previousEpochJustifiedAttestingBalance = (previousEpochJustifiedAttesterIndices as [ValidatorIndex]).totalBalance(state: state)
 
         let previousEpochBoundaryAttestations = previousEpochJustifiedAttestations.filter {
-            $0.data.epochBoundaryRoot == BeaconChain.getBlockRoot(state: state, slot: BeaconChain.getEpochStartSlot(previousEpoch))
+            $0.data.epochBoundaryRoot == BeaconChain.getBlockRoot(state: state, slot: previousEpoch.startSlot())
         }
 
         let previousEpochBoundaryAttesterIndices = previousEpochBoundaryAttestations.flatMap {
@@ -501,7 +501,7 @@ extension StateTransition {
         previousEpochAttestations: [PendingAttestation]
     ) {
 
-        for slot in BeaconChain.getEpochStartSlot(previousEpoch)..<BeaconChain.getEpochStartSlot(nextEpoch) {
+        for slot in previousEpoch.startSlot()..<nextEpoch.startSlot() {
             let crosslinkCommitteesAtSlot = BeaconChain.getCrosslinkCommitteesAtSlot(state: state, slot: slot)
 
             for (_, (crosslinkCommittee, shard)) in crosslinkCommitteesAtSlot.enumerated() {
@@ -622,7 +622,7 @@ extension StateTransition {
             state.validatorBalances[Int(proposer)] += baseReward(state: state, index: index, baseRewardQuotient: baseRewardQuotient) / INCLUDER_REWARD_QUOTIENT
         }
 
-        for slot in BeaconChain.getEpochStartSlot(previousEpoch)..<BeaconChain.getEpochStartSlot(currentEpoch) {
+        for slot in previousEpoch.startSlot()..<currentEpoch.startSlot() {
             let crosslinkCommitteesAtSlot = BeaconChain.getCrosslinkCommitteesAtSlot(state: state, slot: slot)
 
             for (_, (crosslinkCommittee, shard)) in crosslinkCommitteesAtSlot.enumerated() {
@@ -701,7 +701,7 @@ extension StateTransition {
 
         var balanceChurn = UInt64(0)
         for (i, v) in state.validatorRegistry.enumerated() {
-            if v.activationEpoch > BeaconChain.getEntryExitEpoch(currentEpoch) && state.validatorBalances[Int(i)] >= MAX_DEPOSIT_AMOUNT {
+            if v.activationEpoch > currentEpoch.entryExitEpoch() && state.validatorBalances[Int(i)] >= MAX_DEPOSIT_AMOUNT {
                 balanceChurn += BeaconChain.getEffectiveBalance(state: state, index: ValidatorIndex(i))
                 if balanceChurn > maxBalanceChurn {
                     break
@@ -713,7 +713,7 @@ extension StateTransition {
 
         balanceChurn = 0
         for (i, v) in state.validatorRegistry.enumerated() {
-            if v.exitEpoch > BeaconChain.getEntryExitEpoch(currentEpoch) && v.statusFlags & StatusFlag.INITIATED_EXIT.rawValue == 1 {
+            if v.exitEpoch > currentEpoch.entryExitEpoch() && v.statusFlags & StatusFlag.INITIATED_EXIT.rawValue == 1 {
                 balanceChurn += BeaconChain.getEffectiveBalance(state: state, index: ValidatorIndex(i))
                 if balanceChurn > maxBalanceChurn {
                     break

diff --git a/Tests/BeaconChainTests/Extensions/EpochNumberTests.swift b/Tests/BeaconChainTests/Extensions/EpochNumberTests.swift
@@ -0,0 +1,13 @@
+import XCTest
+@testable import BeaconChain
+
+final class EpochNumberTests: XCTestCase {
+
+    func testStartSlot() {
+        XCTAssertEqual(EpochNumber(1).startSlot(), 64)
+    }
+
+    func testEntryExitEpoch() {
+        XCTAssertEqual(EpochNumber(1).entryExitEpoch(), 6)
+    }
+}
diff --git a/Tests/BeaconChainTests/XCTestManifests.swift b/Tests/BeaconChainTests/XCTestManifests.swift
@@ -8,6 +8,7 @@ public func allTests() -> [XCTestCaseEntry] {
         testCase(ArrayValidatorIndexTests.allTests),
         testCase(ArrayValidatorTests.allTests),
         testCase(SlotNumberTests.allTests),
+        testCase(EpochNumberTests.allTests),
 //        testCase(StateTransitionTests.allTests)
     ]
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -6,4 +6,4 @@ extension SlotNumber { @@
             return self / EPOCH_LENGTH
         }
-    }
+    }