feat:Add balance helper

Signed-off-by: Chen Kai <[email protected]>

src/consensus/helpers/balance.zig

@@ -32,6 +32,19 @@ pub fn getTotalBalance(state: *const consensus.BeaconState, indices: std.AutoHas
     return @max(preset.ActivePreset.get().EFFECTIVE_BALANCE_INCREMENT, total);
 }
 
+pub fn getTotalActiveBalance(state: *const consensus.BeaconState, allocator: std.mem.Allocator) !primitives.Gwei {
+    const active_indices = try validator_helper.getActiveValidatorIndices(state, epoch_helper.getCurrentEpoch(state), allocator);
+    defer allocator.free(active_indices);
+    var indices_set = std.AutoHashMap(primitives.ValidatorIndex, void).init(allocator);
+    defer indices_set.deinit();
+
+    for (active_indices) |index| {
+        try indices_set.put(index, {});
+    }
+
+    return getTotalBalance(state, indices_set);
+}
+
 test "test getTotalBalance" {
     preset.ActivePreset.set(preset.Presets.minimal);
     defer preset.ActivePreset.reset();
@@ -111,3 +124,77 @@ test "test getTotalBalance" {
         total,
     );
 }
+
+test "test getTotalActiveBalance" {
+    preset.ActivePreset.set(preset.Presets.minimal);
+    defer preset.ActivePreset.reset();
+    var finalized_checkpoint = consensus.Checkpoint{
+        .epoch = 5,
+        .root = .{0} ** 32,
+    };
+
+    var validators = std.ArrayList(consensus.Validator).init(std.testing.allocator);
+    defer validators.deinit();
+
+    const validator1 = consensus.Validator{
+        .pubkey = undefined,
+        .withdrawal_credentials = undefined,
+        .effective_balance = 10000000000,
+        .slashed = false,
+        .activation_eligibility_epoch = 0,
+        .activation_epoch = 0,
+        .exit_epoch = 10,
+        .withdrawable_epoch = 10,
+    };
+
+    const validator2 = consensus.Validator{
+        .pubkey = undefined,
+        .withdrawal_credentials = undefined,
+        .effective_balance = 100000000000,
+        .slashed = false,
+        .activation_eligibility_epoch = 0,
+        .activation_epoch = 0,
+        .exit_epoch = 20,
+        .withdrawable_epoch = 20,
+    };
+
+    for (0..500000) |_| {
+        try validators.append(validator1);
+        try validators.append(validator2);
+    }
+
+    const state = consensus.BeaconState{
+        .altair = altair.BeaconState{
+            .genesis_time = 0,
+            .genesis_validators_root = .{0} ** 32,
+            .slot = 100,
+            .fork = undefined,
+            .block_roots = undefined,
+            .state_roots = undefined,
+            .historical_roots = undefined,
+            .eth1_data = undefined,
+            .eth1_data_votes = undefined,
+            .eth1_deposit_index = 0,
+            .validators = validators.items,
+            .balances = undefined,
+            .randao_mixes = undefined,
+            .slashings = undefined,
+            .previous_epoch_attestations = undefined,
+            .current_epoch_attestations = undefined,
+            .justification_bits = undefined,
+            .previous_justified_checkpoint = undefined,
+            .current_justified_checkpoint = undefined,
+            .finalized_checkpoint = &finalized_checkpoint,
+            .latest_block_header = undefined,
+            .inactivity_scores = undefined,
+            .current_sync_committee = undefined,
+            .next_sync_committee = undefined,
+        },
+    };
+
+    const total = getTotalActiveBalance(&state, std.testing.allocator);
+    try std.testing.expectEqual(
+        50000000000000000,
+        total,
+    );
+}

src/consensus/helpers/domain.zig

@@ -2,9 +2,14 @@ const std = @import("std");
 const primitives = @import("../../primitives/types.zig");
 const consensus = @import("../../consensus/types.zig");
 const configs = @import("../../configs/config.zig");
+const ssz = @import("../../ssz/ssz.zig");
 
 /// computeForkDataRoot return the 32-byte fork data root for the ``current_version`` and ``genesis_validators_root``.
 /// This is used primarily in signature domains to avoid collisions across forks/chains.
+/// @param current_version The current fork version
+/// @param genesis_validators_root The genesis validators root
+/// @param allocator The allocator to use
+/// @return The 32-byte fork data root
 /// Spec pseudocode definition:
 /// def compute_fork_data_root(current_version: Version, genesis_validators_root: Root) -> Root:
 ///     """
@@ -15,20 +20,25 @@ const configs = @import("../../configs/config.zig");
 ///        current_version=current_version,
 ///        genesis_validators_root=genesis_validators_root,
 ///   ))
-pub fn computeForkDataRoot(current_version: primitives.Version, genesis_validators_root: primitives.Root) primitives.Root {
+pub fn computeForkDataRoot(current_version: primitives.Version, genesis_validators_root: primitives.Root, allocator: std.mem.Allocator) !primitives.Root {
     const fork_data = consensus.ForkData{
         .current_version = current_version,
         .genesis_validators_root = genesis_validators_root,
     };
 
     std.log.debug("ForkData: {}\n", .{fork_data});
-    // todo: implement hash_tree_root
-    return @as(primitives.Root, .{0} ** 32);
+    var out: primitives.Root = undefined;
+    try ssz.hashTreeRoot(fork_data, &out, allocator);
+    return out;
 }
 
 /// computeForkDigest returns the 4-byte fork digest for the `currentVersion` and `genesisValidatorsRoot`.
 /// This is a digest primarily used for domain separation on the p2p layer.
 /// 4-bytes suffices for practical separation of forks/chains.
+/// @param currentVersion The current fork version
+/// @param genesisValidatorsRoot The genesis validators root
+/// @param allocator The allocator to use
+/// @return The 4-byte fork digest
 /// Spec pseudocode definition:
 /// def compute_fork_digest(current_version: Version, genesis_validators_root: Root) -> ForkDigest:
 ///    """
@@ -37,12 +47,17 @@ pub fn computeForkDataRoot(current_version: primitives.Version, genesis_validato
 ///    4-bytes suffices for practical separation of forks/chains.
 ///    """
 ///    return ForkDigest(compute_fork_data_root(current_version, genesis_validators_root)[:4])
-pub fn computeForkDigest(currentVersion: primitives.Version, genesisValidatorsRoot: primitives.Root) primitives.ForkDigest {
-    const forkDataRoot = computeForkDataRoot(currentVersion, genesisValidatorsRoot);
+pub fn computeForkDigest(currentVersion: primitives.Version, genesisValidatorsRoot: primitives.Root, allocator: std.mem.Allocator) !primitives.ForkDigest {
+    const forkDataRoot = try computeForkDataRoot(currentVersion, genesisValidatorsRoot, allocator);
     return forkDataRoot[0..4].*;
 }
 
 /// computeDomain returns the domain for the ``domain_type`` and ``fork_version`` and ``genesis_validators_root``.
+/// @param domain_type - The domain type.
+/// @param fork_version - The fork version.
+/// @param genesis_validators_root - The genesis validators root.
+/// @param allocator - The allocator to use.
+/// @return The domain.
 /// Spec pseudocode definition:
 /// def compute_domain(domain_type: DomainType, fork_version: Version=None, genesis_validators_root: Root=None) -> Domain:
 ///    """
@@ -54,14 +69,14 @@ pub fn computeForkDigest(currentVersion: primitives.Version, genesisValidatorsRo
 ///       genesis_validators_root = Root()  # all bytes zero by default
 ///    fork_data_root = compute_fork_data_root(fork_version, genesis_validators_root)
 ///    return Domain(domain_type + fork_data_root[:28])
-pub fn computeDomain(domain_type: primitives.DomainType, fork_version: ?primitives.Version, genesis_validators_root: ?primitives.Root) primitives.Domain {
+pub fn computeDomain(domain_type: primitives.DomainType, fork_version: ?primitives.Version, genesis_validators_root: ?primitives.Root, allocator: std.mem.Allocator) !primitives.Domain {
     const DOMAIN_TYPE_LENGTH: usize = 4;
     const FORK_DATA_ROOT_LENGTH: usize = 28;
 
     const effective_fork_version = fork_version orelse configs.ActiveConfig.get().GENESIS_FORK_VERSION;
     const effective_genesis_validators_root = genesis_validators_root orelse @as(primitives.Root, .{0} ** 32);
 
-    const fork_data_root = computeForkDataRoot(effective_fork_version, effective_genesis_validators_root);
+    const fork_data_root = try computeForkDataRoot(effective_fork_version, effective_genesis_validators_root, allocator);
 
     const result = blk: {
         var temp: primitives.Domain = undefined;
@@ -76,21 +91,21 @@ pub fn computeDomain(domain_type: primitives.DomainType, fork_version: ?primitiv
 test "test computeForkDigest" {
     const currentVersion = .{0} ** 4;
     const genesisValidatorsRoot = .{0} ** 32;
-    const forkDigest = computeForkDigest(currentVersion, genesisValidatorsRoot);
+    const forkDigest =try computeForkDigest(currentVersion, genesisValidatorsRoot, std.testing.allocator);
     try std.testing.expectEqual(4, forkDigest.len);
 }
 
 test "test computeForkDataRoot" {
     const currentVersion = .{0} ** 4;
     const genesisValidatorsRoot = .{0} ** 32;
-    const forkDataRoot = computeForkDataRoot(currentVersion, genesisValidatorsRoot);
+    const forkDataRoot =try computeForkDataRoot(currentVersion, genesisValidatorsRoot, std.testing.allocator);
     try std.testing.expectEqual(32, forkDataRoot.len);
 }
 
 test "test computeDomain" {
     const domainType = .{0} ** 4;
     const forkVersion = .{0} ** 4;
     const genesisValidatorsRoot = .{0} ** 32;
-    const domain = computeDomain(domainType, forkVersion, genesisValidatorsRoot);
+    const domain = try computeDomain(domainType, forkVersion, genesisValidatorsRoot, std.testing.allocator);
     try std.testing.expectEqual(32, domain.len);
 }