diff --git a/src/consensus/fork_choice.zig b/src/consensus/fork_choice.zig index d313769643..b273c1cf8e 100644 --- a/src/consensus/fork_choice.zig +++ b/src/consensus/fork_choice.zig @@ -5,8 +5,7 @@ const builtin = @import("builtin"); const Instant = sig.time.Instant; const Hash = sig.core.Hash; const Pubkey = sig.core.Pubkey; -const SortedMap = sig.utils.collections.SortedMapUnmanaged; -const SortedMapCustom = sig.utils.collections.SortedMapUnmanagedCustom; +const SortedMap = sig.utils.collections.SortedMap; const SlotAndHash = sig.core.hash.SlotAndHash; const Slot = sig.core.Slot; const EpochStakesMap = sig.core.EpochStakesMap; @@ -42,7 +41,7 @@ const ForkInfo = struct { /// forks, unlike `heaviest_slot` deepest_slot: SlotAndHash, parent: ?SlotAndHash, - children: Children, + children: ChildSet, /// The latest ancestor of this node that has been marked invalid by being a duplicate. /// If the slot itself is a duplicate, this is set to the slot itself. latest_duplicate_ancestor: ?Slot, @@ -52,11 +51,9 @@ const ForkInfo = struct { /// and all competing forks for the same slot are invalid. is_duplicate_confirmed: bool, - const Children = SortedMapCustom(SlotAndHash, void, .{ - .orderFn = SlotAndHash.order, - }); + pub const ChildSet = SortedMap(SlotAndHash, void, .{}); - fn deinit(self: *const ForkInfo, allocator: std.mem.Allocator) void { + fn deinit(self: *ForkInfo, allocator: std.mem.Allocator) void { self.children.deinit(allocator); } @@ -162,9 +159,10 @@ pub const ForkChoice = struct { return self; } - pub fn deinit(self: *const ForkChoice, allocator: std.mem.Allocator) void { - for (self.fork_infos.values()) |fork_info| { - fork_info.deinit(allocator); + pub fn deinit(self: ForkChoice, allocator: std.mem.Allocator) void { + var it = self.fork_infos.iterator(); + while (it.next()) |fork_info| { + fork_info.value_ptr.deinit(allocator); } var fork_infos = self.fork_infos; @@ -594,9 +592,13 @@ pub const ForkChoice = struct { var remove_set = try self.subtreeDiff(allocator, &self.tree_root, new_root); defer remove_set.deinit(allocator); - for (remove_set.keys()) |node_key| { - if (!self.fork_infos.contains(node_key)) { - return error.MissingForkInfo; + { + var iter = remove_set.iterator(); + while (iter.next()) |entry| { + const node_key = entry.key_ptr.*; + if (!self.fork_infos.contains(node_key)) { + return error.MissingForkInfo; + } } } @@ -608,9 +610,11 @@ pub const ForkChoice = struct { // At this point, both the subtree to be removed and new root // are confirmed to be in the fork choice. - for (remove_set.keys()) |node_key| { + var iter = remove_set.iterator(); + while (iter.next()) |entry| { + const node_key = entry.key_ptr.*; // SAFETY: Previous contains check ensures this won't panic. - const kv = self.fork_infos.fetchSwapRemove(node_key).?; + var kv = self.fork_infos.fetchSwapRemove(node_key).?; kv.value.deinit(allocator); } @@ -657,10 +661,10 @@ pub const ForkChoice = struct { try self.fork_infos.ensureUnusedCapacity(allocator, 1); // Create the new root parent's fork info - var root_parent_children: ForkInfo.Children = .empty; + var root_parent_children: ForkInfo.ChildSet = .empty; + errdefer root_parent_children.deinit(allocator); try root_parent_children.put(allocator, self.tree_root, {}); errdefer comptime unreachable; - self.fork_infos.putAssumeCapacityNoClobber(root_parent, .{ .stake_for_slot = 0, .stake_for_subtree = root_info.stake_for_subtree, @@ -741,10 +745,9 @@ pub const ForkChoice = struct { ); defer children_hash_keys.deinit(allocator); - const children_hash_keys_keys = children_hash_keys.keys(); - for (1..children_hash_keys_keys.len + 1) |i_plus_one| { - const rev_i = children_hash_keys.count() - i_plus_one; - const child_hash_key = children_hash_keys_keys[rev_i]; + var it = children_hash_keys.iteratorRanged(null, null, .end); + while (it.prev()) |entry| { + const child_hash_key = entry.key_ptr.*; self.markForkValid(&child_hash_key, valid_slot_hash_key.slot); self.aggregateSlot(child_hash_key); @@ -781,10 +784,9 @@ pub const ForkChoice = struct { ); defer children_hash_keys.deinit(allocator); - const children_hash_keys_keys = children_hash_keys.keys(); - for (1..children_hash_keys_keys.len + 1) |i_plus_one| { - const rev_i = children_hash_keys.count() - i_plus_one; - const child_hash_key = children_hash_keys_keys[rev_i]; + var it = children_hash_keys.iteratorRanged(null, null, .end); + while (it.prev()) |entry| { + const child_hash_key = entry.key_ptr.*; self.markForkInvalid(child_hash_key, invalid_slot_hash_key.slot); self.aggregateSlot(child_hash_key); @@ -964,7 +966,10 @@ pub const ForkChoice = struct { const parent = maybe_parent orelse return true; var children = self.getChildren(&parent) orelse return false; - for (children.keys()) |child| { + var iter = children.iterator(); + while (iter.next()) |entry| { + const child = entry.key_ptr.*; + // child must exist in `self.fork_infos` const child_weight = self.stakeForSubtree(&child) orelse return error.MissingChild; @@ -1006,7 +1011,10 @@ pub const ForkChoice = struct { // Get the other chidren of the parent. i.e. siblings of the deepest_child. var children = self.getChildren(&parent) orelse return false; - for (children.keys()) |child| { + var iter = children.iterator(); + while (iter.next()) |entry| { + const child = entry.key_ptr.*; + const child_height = self.getHeight(&child) orelse return false; const child_weight = self.stakeForSubtree(&child) orelse return false; @@ -1046,9 +1054,9 @@ pub const ForkChoice = struct { fn getChildren( self: *const ForkChoice, slot_hash_key: *const SlotAndHash, - ) ?*ForkInfo.Children { - const fork_info = self.fork_infos.getPtr(slot_hash_key.*) orelse return null; - return &fork_info.children; + ) ?ForkInfo.ChildSet { + const fork_info = self.fork_infos.get(slot_hash_key.*) orelse return null; + return fork_info.children; } pub fn latestInvalidAncestor( @@ -1116,20 +1124,28 @@ pub const ForkChoice = struct { allocator: std.mem.Allocator, root1: *const SlotAndHash, root2: *const SlotAndHash, - ) (std.mem.Allocator.Error || error{MissingChild})!SortedMap(SlotAndHash, void) { + ) (std.mem.Allocator.Error || error{MissingChild})!ForkInfo.ChildSet { if (!self.containsBlock(root1)) return .empty; var pending_keys: std.ArrayListUnmanaged(SlotAndHash) = .empty; defer pending_keys.deinit(allocator); try pending_keys.append(allocator, root1.*); - var reachable_set: SortedMap(SlotAndHash, void) = .empty; + var reachable_set: ForkInfo.ChildSet = .empty; errdefer reachable_set.deinit(allocator); while (pending_keys.pop()) |current_key| { - if (current_key.equals(root2.*)) continue; - const children = self.getChildren(¤t_key) orelse return error.MissingChild; - try pending_keys.appendSlice(allocator, children.keys()); + if (current_key.equals(root2.*)) { + continue; + } + + var children = self.getChildren(¤t_key) orelse return error.MissingChild; + + var iter = children.iterator(); + while (iter.next()) |entry| { + try pending_keys.append(allocator, entry.key_ptr.*); + } + try reachable_set.put(allocator, current_key, {}); } @@ -1204,7 +1220,9 @@ pub const ForkChoice = struct { var deepest_child_slot_key: SlotAndHash = slot_hash_key; // Iterate over the children of the current fork - for (fork_info.children.keys()) |child_key| { + var it = fork_info.children.iterator(); + while (it.next()) |entry| { + const child_key = entry.key_ptr.*; const child_fork_info = self.fork_infos.get(child_key) orelse { std.debug.panic("Child must exist in fork_info map", .{}); }; @@ -1474,7 +1492,7 @@ pub const ForkChoice = struct { // Remove child link so that this slot cannot be chosen as best or deepest const parent_info = self.fork_infos.getPtr(parent) orelse return error.ParentNotFound; - std.debug.assert(parent_info.children.orderedRemove(slot_hash_key)); + std.debug.assert(parent_info.children.remove(slot_hash_key)); { // Insert aggregate operations up to the root var parent_iter = self.ancestorIterator(slot_hash_key); @@ -1497,13 +1515,17 @@ pub const ForkChoice = struct { var current_fork_info = current_kv.value; try split_tree_fork_infos.put(allocator, current_node, current_fork_info); - try to_visit.appendSlice(allocator, current_fork_info.children.keys()); + + var iter = current_fork_info.children.iterator(); + while (iter.next()) |child| { + try to_visit.append(allocator, child.key_ptr.*); + } } // Remove link from parent const parent_fork_info = self.fork_infos.getPtr(parent) orelse return error.ParentNotFound; - _ = parent_fork_info.children.swapRemoveNoSort(slot_hash_key); + _ = parent_fork_info.children.remove(slot_hash_key); // Update the root of the new tree with the proper info, now that we have finished // aggregating @@ -1607,23 +1629,22 @@ test "HeaviestSubtreeForkChoice.subtreeDiff" { ); defer diff.deinit(allocator); - const items = diff.items(); - const slot_and_hashes = items[0]; - - try std.testing.expectEqual(3, slot_and_hashes.len); + var iter = diff.iterator(); try std.testing.expectEqual( - slot_and_hashes[0], - SlotAndHash{ .slot = 3, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 3, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[1], - SlotAndHash{ .slot = 5, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 5, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[2], - SlotAndHash{ .slot = 6, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 6, .hash = Hash.ZEROES }, ); + + try std.testing.expectEqual(null, iter.next()); } // The set reachable from slot 1, excluding subtree 3, is just 1 and @@ -1636,23 +1657,22 @@ test "HeaviestSubtreeForkChoice.subtreeDiff" { ); defer diff.deinit(allocator); - const items = diff.items(); - const slot_and_hashes = items[0]; // Access the keys slice - - try std.testing.expectEqual(3, slot_and_hashes.len); + var iter = diff.iterator(); try std.testing.expectEqual( - slot_and_hashes[0], - SlotAndHash{ .slot = 1, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 1, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[1], - SlotAndHash{ .slot = 2, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 2, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[2], - SlotAndHash{ .slot = 4, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 4, .hash = Hash.ZEROES }, ); + + try std.testing.expectEqual(null, iter.next()); } // The set reachable from slot 1, excluding leaf 6, is just everything @@ -1665,35 +1685,34 @@ test "HeaviestSubtreeForkChoice.subtreeDiff" { ); defer diff.deinit(allocator); - const items = diff.items(); - const slot_and_hashes = items[0]; // Access the keys slice - - try std.testing.expectEqual(6, slot_and_hashes.len); + var iter = diff.iterator(); try std.testing.expectEqual( - slot_and_hashes[0], - SlotAndHash{ .slot = 0, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 0, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[1], - SlotAndHash{ .slot = 1, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 1, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[2], - SlotAndHash{ .slot = 2, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 2, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[3], - SlotAndHash{ .slot = 3, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 3, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[4], - SlotAndHash{ .slot = 4, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 4, .hash = Hash.ZEROES }, ); try std.testing.expectEqual( - slot_and_hashes[5], - SlotAndHash{ .slot = 5, .hash = .ZEROES }, + iter.next().?.key_ptr.*, + SlotAndHash{ .slot = 5, .hash = Hash.ZEROES }, ); + + try std.testing.expectEqual(null, iter.next()); } { @@ -2233,6 +2252,8 @@ test "HeaviestSubtreeForkChoice.setRootAndAddOutdatedVotes" { &EpochSchedule.INIT, ); + // v FAILURE + try std.testing.expectEqual( stake, fork_choice.stakeForSlot(&.{ .slot = 3, .hash = .ZEROES }).?, @@ -2560,11 +2581,11 @@ test "HeaviestSubtreeForkChoice.addNewLeafSlot_duplicate" { const child: SlotAndHash = .{ .slot = 11, .hash = .initRandom(random) }; try fork_choice.addNewLeafSlot(allocator, child, duplicate_parent); { - var children_ = fork_choice.getChildren(&duplicate_parent).?; - const children = children_.keys(); + const children = fork_choice.getChildren(&duplicate_parent).?; + const min_child = children.minEntry().?.key_ptr.*; - try std.testing.expectEqual(child.slot, children[0].slot); - try std.testing.expectEqual(child.hash, children[0].hash); + try std.testing.expectEqual(child.slot, min_child.slot); + try std.testing.expectEqual(child.hash, min_child.hash); } try std.testing.expectEqual( @@ -2587,11 +2608,11 @@ test "HeaviestSubtreeForkChoice.addNewLeafSlot_duplicate" { // Re-adding same duplicate slot should not overwrite existing one try fork_choice.addNewLeafSlot(allocator, duplicate_parent, .{ .slot = 4, .hash = .ZEROES }); { - var children_ = fork_choice.getChildren(&duplicate_parent).?; - const children = children_.keys(); + const children = fork_choice.getChildren(&duplicate_parent).?; + const min_child = children.minEntry().?.key_ptr.*; - try std.testing.expectEqual(child.slot, children[0].slot); - try std.testing.expectEqual(child.hash, children[0].hash); + try std.testing.expectEqual(child.slot, min_child.slot); + try std.testing.expectEqual(child.hash, min_child.hash); } try std.testing.expectEqual(child, fork_choice.heaviestOverallSlot()); @@ -3163,8 +3184,8 @@ test "HeaviestSubtreeForkChoice.addRootParent" { var children = fork_choice.getChildren(&.{ .slot = 2, .hash = .ZEROES }).?; try std.testing.expectEqual(1, children.count()); try std.testing.expectEqual( - SlotAndHash{ .slot = 3, .hash = .ZEROES }, - children.keys()[0], + SlotAndHash{ .slot = 3, .hash = Hash.ZEROES }, + children.minEntry().?.key_ptr.*, ); try std.testing.expectEqual( @@ -4831,10 +4852,7 @@ const linear_fork_tuples = [_]TreeNode{ }; fn compareSlotHashKey(_: void, a: SlotAndHash, b: SlotAndHash) bool { - if (a.slot == b.slot) { - return a.hash.order(&b.hash) == .lt; - } - return a.slot < b.slot; + return a.order(b) == .lt; } const TestDuplicateForks = struct { @@ -4925,14 +4943,16 @@ const TestDuplicateForks = struct { } // Verify children of slot 4 - var dup_children_4 = fork_choice.getChildren(&.{ - .slot = 4, - .hash = .ZEROES, - }).?; - - std.mem.sort(SlotAndHash, dup_children_4.mutableKeys(), {}, compareSlotHashKey); - std.debug.assert(dup_children_4.keys()[0].equals(dupe_leaves_desc_from_4.items[0])); - std.debug.assert(dup_children_4.keys()[1].equals(dupe_leaves_desc_from_4.items[1])); + { + var dup_children_4 = fork_choice.getChildren(&.{ .slot = 4, .hash = Hash.ZEROES }).?; + var dup_iter = dup_children_4.iterator(); + std.debug.assert( + dup_iter.next().?.key_ptr.equals(dupe_leaves_desc_from_4.items[0]), + ); + std.debug.assert( + dup_iter.next().?.key_ptr.equals(dupe_leaves_desc_from_4.items[1]), + ); + } var dup_children_5: std.ArrayListUnmanaged(SlotAndHash) = .empty; defer dup_children_5.deinit(gpa); @@ -4942,9 +4962,11 @@ const TestDuplicateForks = struct { .hash = .ZEROES, }).?; - for (children_5.keys()) |key| { - if (key.slot == duplicate_slot) { - try dup_children_5.append(gpa, key); + { + var iter = children_5.iterator(); + while (iter.next()) |entry| { + const key = entry.key_ptr.*; + if (key.slot == duplicate_slot) dup_children_5.append(gpa, key) catch unreachable; } } @@ -4961,9 +4983,11 @@ const TestDuplicateForks = struct { .hash = .ZEROES, }).?; - for (children_6.keys()) |key| { - if (key.slot == duplicate_slot) { - try dup_children_6.append(gpa, key); + { + var iter = children_6.iterator(); + while (iter.next()) |entry| { + const key = entry.key_ptr.*; + if (key.slot == duplicate_slot) dup_children_6.append(gpa, key) catch unreachable; } } diff --git a/src/consensus/optimistic_vote_verifier.zig b/src/consensus/optimistic_vote_verifier.zig index 403be39953..24aefc5e33 100644 --- a/src/consensus/optimistic_vote_verifier.zig +++ b/src/consensus/optimistic_vote_verifier.zig @@ -38,7 +38,7 @@ pub const OptimisticVotesTracker = struct { /// Analogous to [OptimisticConfirmationVerifier](https://github.com/anza-xyz/agave/blob/9d8bf065f7aad8257addfc5639ae5cea4e743204/core/src/optimistic_confirmation_verifier.rs#L11) pub const OptimisticConfirmationVerifier = struct { snapshot_start_slot: Slot, - unchecked_slots: sig.utils.collections.SortedSetUnmanaged(sig.core.hash.SlotAndHash), + unchecked_slots: sig.utils.collections.SortedSet(sig.core.hash.SlotAndHash, .{}), last_optimistic_slot_ts: sig.time.Instant, pub fn deinit( @@ -70,24 +70,28 @@ pub const OptimisticConfirmationVerifier = struct { ancestors: *const sig.core.Ancestors, }, ) ![]const sig.core.hash.SlotAndHash { + var after_root: sig.utils.collections.SortedSet(sig.core.hash.SlotAndHash, .{}) = .empty; + var after_root_moved: bool = false; + defer if (!after_root_moved) after_root.deinit(allocator); + var before_or_equal_root: std.ArrayListUnmanaged(sig.core.hash.SlotAndHash) = .empty; defer before_or_equal_root.deinit(allocator); + try before_or_equal_root.ensureUnusedCapacity(allocator, self.unchecked_slots.count()); - var after_root: sig.utils.collections.SortedSetUnmanaged(sig.core.hash.SlotAndHash) = .empty; - - const items = self.unchecked_slots.items(); - try before_or_equal_root.ensureTotalCapacityPrecise(allocator, items.len); - for (items) |sah| { + var iter = self.unchecked_slots.iterator(); + while (iter.next()) |entry| { + const sah = entry.key_ptr.*; if (sah.slot > root.slot) { - try after_root.put(allocator, sah); + try after_root.put(allocator, sah, {}); } else { before_or_equal_root.appendAssumeCapacity(sah); } } - const old_set = self.unchecked_slots; + var old_set = self.unchecked_slots; self.unchecked_slots = after_root; old_set.deinit(allocator); + after_root_moved = true; var optimistic_root_not_rooted: std.ArrayListUnmanaged(sig.core.hash.SlotAndHash) = .empty; errdefer optimistic_root_not_rooted.deinit(allocator); @@ -127,7 +131,7 @@ pub const OptimisticConfirmationVerifier = struct { ) catch |err| { logger.err().logf("insertOptimisticSlot: {s}", .{@errorName(err)}); }; - try self.unchecked_slots.put(allocator, slot_and_hash); + try self.unchecked_slots.put(allocator, slot_and_hash, {}); } } diff --git a/src/consensus/replay_tower.zig b/src/consensus/replay_tower.zig index 76d4cf1c23..f430776d39 100644 --- a/src/consensus/replay_tower.zig +++ b/src/consensus/replay_tower.zig @@ -14,7 +14,7 @@ const Epoch = sig.core.Epoch; const EpochStakesMap = sig.core.EpochStakesMap; const SlotAndHash = sig.core.hash.SlotAndHash; const SlotHistory = sig.runtime.sysvar.SlotHistory; -const SortedSetUnmanaged = sig.utils.collections.SortedSetUnmanaged; +const SortedSet = sig.utils.collections.SortedSet; const TowerSync = sig.runtime.program.vote.state.TowerSync; const Vote = sig.runtime.program.vote.state.Vote; const VoteStateUpdate = sig.runtime.program.vote.state.VoteStateUpdate; @@ -450,7 +450,7 @@ pub const ReplayTower = struct { allocator: std.mem.Allocator, switch_slot: Slot, ancestors: *const AutoArrayHashMapUnmanaged(Slot, Ancestors), - descendants: *const AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + descendants: *const AutoArrayHashMapUnmanaged(Slot, SlotSet), progress: *const ProgressMap, total_stake: u64, epoch_vote_accounts: *const StakeAndVoteAccountsMap, @@ -593,7 +593,7 @@ pub const ReplayTower = struct { const switch_proof = Hash.ZEROES; var locked_out_stake: u64 = 0; - var locked_out_vote_accounts: SortedSetUnmanaged(Pubkey) = .empty; + var locked_out_vote_accounts: SortedSet(Pubkey, .{}) = .empty; defer locked_out_vote_accounts.deinit(allocator); for (descendants.keys(), descendants.values()) |candidate_slot, *candidate_descendants| { @@ -609,11 +609,14 @@ pub const ReplayTower = struct { // recent frozen bank on this fork to use, so we can ignore this one. Otherwise, // even if this bank has descendants, if they have not yet been frozen / stats computed, // then use this bank as a representative for the fork. - const is_descendant_computed = for (candidate_descendants.items()) |d| { - if (progress.getForkStats(d)) |stat| { - if (stat.computed) break true; + const is_descendant_computed = blk: { + var iter = candidate_descendants.iterator(); + while (iter.next()) |entry| { + const stat = progress.getForkStats(entry.key_ptr.*) orelse continue; + if (stat.computed) break :blk true; } - } else false; + break :blk false; + }; if (is_descendant_computed) continue; // 3) Don't consider lockouts on the `last_vote` itself @@ -693,7 +696,7 @@ pub const ReplayTower = struct { ) > SWITCH_FORK_THRESHOLD) { return SwitchForkDecision{ .switch_proof = switch_proof }; } - try locked_out_vote_accounts.put(allocator, voted_pubkey); + try locked_out_vote_accounts.put(allocator, voted_pubkey, {}); } } } @@ -753,7 +756,7 @@ pub const ReplayTower = struct { }; } - try locked_out_vote_accounts.put(allocator, vote_account_pubkey); + try locked_out_vote_accounts.put(allocator, vote_account_pubkey, {}); } } } @@ -782,7 +785,7 @@ pub const ReplayTower = struct { allocator: std.mem.Allocator, switch_slot: Slot, ancestors: *const AutoArrayHashMapUnmanaged(Slot, Ancestors), - descendants: *const AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + descendants: *const AutoArrayHashMapUnmanaged(Slot, SlotSet), progress: *const ProgressMap, total_stake: u64, epoch_vote_accounts: *const StakeAndVoteAccountsMap, @@ -1398,7 +1401,7 @@ pub const ReplayTower = struct { heaviest_slot_on_same_voted_fork: ?Slot, heaviest_epoch: Epoch, ancestors: *const AutoArrayHashMapUnmanaged(u64, Ancestors), - descendants: *const AutoArrayHashMapUnmanaged(u64, SortedSetUnmanaged(u64)), + descendants: *const AutoArrayHashMapUnmanaged(u64, SlotSet), progress: *const ProgressMap, latest_validator_votes: *const LatestValidatorVotes, fork_choice: *const HeaviestSubtreeForkChoice, @@ -1693,7 +1696,7 @@ pub fn collectClusterVoteState( defer zone.deinit(); // The state we are interested in. - var vote_slots: SortedSetUnmanaged(Slot) = .empty; + var vote_slots: SlotSet = .empty; defer vote_slots.deinit(allocator); var voted_stakes = VotedStakes.empty; var total_stake: u64 = 0; @@ -1753,16 +1756,16 @@ pub fn collectClusterVoteState( try vote_state.processNextVoteSlot(bank_slot); for (vote_state.votes.constSlice()) |lockout_vote| { - try vote_slots.put(allocator, lockout_vote.slot); + try vote_slots.put(allocator, lockout_vote.slot, {}); } if (start_root != vote_state.root) { if (start_root) |root| { - try vote_slots.put(allocator, root); + try vote_slots.put(allocator, root, {}); } } if (vote_state.root) |root| { - try vote_slots.put(allocator, root); + try vote_slots.put(allocator, root, {}); } // The last vote in the vote stack is a simulated vote on bank_slot, which @@ -1792,7 +1795,7 @@ pub fn collectClusterVoteState( try populateAncestorVotedStakes( allocator, &voted_stakes, - vote_slots.items(), + &vote_slots, ancestors, ); @@ -1827,13 +1830,17 @@ pub fn collectClusterVoteState( pub fn populateAncestorVotedStakes( allocator: std.mem.Allocator, voted_stakes: *VotedStakes, - vote_slots: []const Slot, + vote_slots: *const SlotSet, ancestors: *const AutoArrayHashMapUnmanaged(Slot, Ancestors), ) !void { // If there's no ancestors, that means this slot must be from before the current root, // in which case the lockouts won't be calculated in bank_weight anyways, so ignore // this slot - for (vote_slots) |vote_slot| { + + var iter = vote_slots.iterator(); + while (iter.next()) |entry| { + const vote_slot = entry.key_ptr.*; + if (ancestors.getPtr(vote_slot)) |slot_ancestors| { _ = try voted_stakes.getOrPutValue(allocator, vote_slot, 0); @@ -4148,7 +4155,7 @@ test "switch threshold" { try ancestors.put(allocator, 112, try createAncestor(allocator, &.{ 43, 2, 1, 0 })); // Create descendants map - var descendants: AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)) = .empty; + var descendants: AutoArrayHashMapUnmanaged(Slot, SlotSet) = .empty; defer { for (descendants.values()) |*set| set.deinit(allocator); descendants.deinit(allocator); @@ -4157,13 +4164,13 @@ test "switch threshold" { // Helper to add descendants const addDescendants = struct { fn call( - desc: *AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + desc: *AutoArrayHashMapUnmanaged(Slot, SlotSet), alloc: std.mem.Allocator, slot: Slot, desc_slots: []const Slot, ) !void { - var set: SortedSetUnmanaged(Slot) = .empty; - for (desc_slots) |s| try set.put(alloc, s); + var set: SlotSet = .empty; + for (desc_slots) |s| try set.put(alloc, s, {}); try desc.put(alloc, slot, set); } }.call; @@ -4525,7 +4532,7 @@ test "switch threshold" { // Test 7b: Adding unfrozen descendant (10000) shouldn't remove slot 14 from consideration { const desc_14 = descendants.getPtr(14).?; - try desc_14.put(allocator, 10000); + try desc_14.put(allocator, 10000, {}); const decision2 = try tower.makeCheckSwitchThresholdDecision( allocator, @@ -4544,7 +4551,7 @@ test "switch threshold" { ); // Clean up descendant - _ = desc_14.orderedRemove(10000); + _ = desc_14.remove(10000); } if (progress.map.getPtr(14)) |fork_progress| { @@ -4665,7 +4672,7 @@ test "switch threshold use gossip votes" { try ancestors.put(allocator, 112, try createAncestor(allocator, &.{ 43, 2, 1, 0 })); // Create descendants map - var descendants: AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)) = .empty; + var descendants: AutoArrayHashMapUnmanaged(Slot, SlotSet) = .empty; defer { for (descendants.values()) |*set| set.deinit(allocator); descendants.deinit(allocator); @@ -4673,13 +4680,13 @@ test "switch threshold use gossip votes" { const addDescendants = struct { fn call( - desc: *AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + desc: *AutoArrayHashMapUnmanaged(Slot, SlotSet), alloc: std.mem.Allocator, slot: Slot, desc_slots: []const Slot, ) !void { - var set: SortedSetUnmanaged(Slot) = .empty; - for (desc_slots) |s| try set.put(alloc, s); + var set: SlotSet = .empty; + for (desc_slots) |s| try set.put(alloc, s, {}); try desc.put(alloc, slot, set); } }.call; @@ -5426,6 +5433,8 @@ fn fillProgressMapForkStats( } } +const SlotSet = SortedSet(Slot, .{}); + pub const MAX_TEST_TREE_LEN = 100; const SlotTracker = sig.replay.trackers.SlotTracker; const Tree = struct { root: SlotAndHash, data: std.BoundedArray(TreeNode, MAX_TEST_TREE_LEN) }; @@ -5433,7 +5442,7 @@ pub const TestFixture = struct { slot_tracker: SlotTracker, fork_choice: HeaviestSubtreeForkChoice, ancestors: AutoArrayHashMapUnmanaged(Slot, Ancestors) = .{}, - descendants: AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)) = .{}, + descendants: AutoArrayHashMapUnmanaged(Slot, SlotSet) = .{}, progress: ProgressMap = ProgressMap.INIT, epoch_stakes: EpochStakesMap, node_pubkeys: std.ArrayListUnmanaged(Pubkey), @@ -5491,7 +5500,7 @@ pub const TestFixture = struct { self.epoch_stakes.deinit(allocator); } - for (self.descendants.values()) |set| set.deinit(allocator); + for (self.descendants.values()) |*set| set.deinit(allocator); self.descendants.deinit(allocator); for (self.ancestors.values()) |*set| set.deinit(allocator); @@ -5618,7 +5627,7 @@ pub const TestFixture = struct { // Populate decendants var extended_descendants = try getDescendants(allocator, input_tree); defer { - for (extended_descendants.values()) |set| set.deinit(allocator); + for (extended_descendants.values()) |*set| set.deinit(allocator); extended_descendants.deinit(allocator); } try extendForkTree(allocator, &self.descendants, extended_descendants); @@ -5646,12 +5655,12 @@ pub const TestFixture = struct { fn getDescendants(allocator: std.mem.Allocator, tree: Tree) !std.AutoArrayHashMapUnmanaged( Slot, - SortedSetUnmanaged(Slot), + SlotSet, ) { if (!builtin.is_test) { @compileError("getDescendants should only be used in test"); } - var descendants = std.AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)){}; + var descendants = std.AutoArrayHashMapUnmanaged(Slot, SlotSet){}; var children_map = std.AutoHashMap(Slot, std.ArrayList(Slot)).init(allocator); defer { @@ -5689,19 +5698,19 @@ fn getDescendants(allocator: std.mem.Allocator, tree: Tree) !std.AutoArrayHashMa if (current.processed) { _ = stack.pop(); - var descendant_list: SortedSetUnmanaged(Slot) = .empty; + var descendant_list: SlotSet = .empty; errdefer descendant_list.deinit(allocator); if (children_map.get(current.slot)) |children| { for (children.items) |item| { - try descendant_list.put(allocator, item); + try descendant_list.put(allocator, item, {}); } for (children.items) |child| { if (descendants.get(child)) |child_descendants| { - var cd = child_descendants; - for (cd.items()) |item| { - try descendant_list.put(allocator, item); + var iter = child_descendants.iterator(); + while (iter.next()) |entry| { + try descendant_list.put(allocator, entry.key_ptr.*, {}); } } } @@ -5792,8 +5801,8 @@ fn getAncestors(allocator: std.mem.Allocator, tree: Tree) !std.AutoArrayHashMapU pub fn extendForkTree( allocator: std.mem.Allocator, - original: *std.AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), - extension: std.AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + original: *std.AutoArrayHashMapUnmanaged(Slot, SlotSet), + extension: std.AutoArrayHashMapUnmanaged(Slot, SlotSet), ) !void { if (!builtin.is_test) { @compileError("extendForkTree should only be used in test"); @@ -5808,8 +5817,10 @@ pub fn extendForkTree( continue; }; - for (extension_children.items()) |extension_child| { - try original_children.put(allocator, extension_child); + var iter = extension_children.iterator(); + while (iter.next()) |entry| { + const extension_child = entry.key_ptr.*; + try original_children.put(allocator, extension_child, {}); } } } diff --git a/src/core/hash.zig b/src/core/hash.zig index 8862ec0ffd..f7dcb6b13a 100644 --- a/src/core/hash.zig +++ b/src/core/hash.zig @@ -11,20 +11,20 @@ pub const SlotAndHash = struct { slot: Slot, hash: Hash, + pub const empty: SlotAndHash = .{ + .slot = std.math.maxInt(Slot), + .hash = .ZEROES, + }; + pub fn tuple(self: SlotAndHash) struct { Slot, Hash } { return .{ self.slot, self.hash }; } pub fn order(a: SlotAndHash, b: SlotAndHash) std.math.Order { - if (a.slot == b.slot and a.hash.order(&b.hash) == .eq) { - return .eq; - } else if (a.slot < b.slot or a.slot == b.slot and (a.hash.order(&b.hash) == .lt)) { - return .lt; - } else if (a.slot > b.slot or a.slot == b.slot and (a.hash.order(&b.hash) == .gt)) { - return .gt; - } else { - unreachable; + if (a.slot == b.slot) { + return a.hash.order(&b.hash); } + return std.math.order(a.slot, b.slot); } pub fn equals(a: SlotAndHash, b: SlotAndHash) bool { diff --git a/src/ledger/Reader.zig b/src/ledger/Reader.zig index 8e6b2438d4..ea7e85b7f1 100644 --- a/src/ledger/Reader.zig +++ b/src/ledger/Reader.zig @@ -6,6 +6,8 @@ const tracy = @import("tracy"); // std const Allocator = std.mem.Allocator; const ArrayList = std.ArrayList; +const ArrayListUnmanaged = std.ArrayListUnmanaged; + const AutoHashMap = std.AutoHashMap; // sig common @@ -42,6 +44,8 @@ const schema = ledger_mod.schema.schema; const key_serializer = ledger_mod.database.key_serializer; const shredder = ledger_mod.shredder; +const DataIndexes = SortedSet(u32, .{}); + const DEFAULT_TICKS_PER_SECOND = sig.core.time.DEFAULT_TICKS_PER_SECOND; const Logger = sig.trace.Logger("reader"); @@ -95,9 +99,9 @@ pub fn slotRangeConnected( schema.slot_meta, starting_slot, ) orelse return false; - defer start_slot_meta.deinit(); + defer start_slot_meta.deinit(allocator); // need a reference so the start_slot_meta.deinit works correctly - var child_slots: *ArrayList(Slot) = &start_slot_meta.child_slots; + var child_slots: *ArrayListUnmanaged(Slot) = &start_slot_meta.child_slots; // TODO: revisit this with more extensive testing. how does agave work fine with // supposed bugs? it may be worth opening a PR in agave with the presumed fix @@ -109,13 +113,13 @@ pub fn slotRangeConnected( const slot = child_slots.items[i]; if (try self.ledger.db.get(allocator, schema.slot_meta, slot)) |_slot_meta| { var slot_meta = _slot_meta; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); if (slot_meta.isFull()) { std.debug.assert(last_slot == slot - 1); // this append is the same as agave, but is it redundant? // does the list already have these slots? - try child_slots.appendSlice(slot_meta.child_slots.items); + try child_slots.appendSlice(allocator, slot_meta.child_slots.items); } else { return false; // this is missing from agave, which seems like a bug } @@ -461,7 +465,7 @@ pub fn getCompleteBlockWithEntries( .logf("getCompleteBlockWithEntries failed for slot {} (missing SlotMeta)", .{slot}); return error.SlotUnavailable; }; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); if (!slot_meta.isFull()) { self.logger.debug() .logf("getCompleteBlockWithEntries failed for slot {} (slot not full)", .{slot}); @@ -1138,7 +1142,7 @@ fn getCompletedRanges( return .{ CompletedRanges.init(allocator), null }; } var slot_meta: SlotMeta = maybe_slot_meta.?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); // Find all the ranges for the completed data blocks const completed_ranges = try getCompletedDataRanges( @@ -1156,7 +1160,7 @@ fn getCompletedRanges( fn getCompletedDataRanges( allocator: Allocator, start_index: u32, - completed_data_indexes: *SortedSet(u32), + completed_data_indexes: *DataIndexes, consumed: u32, ) Allocator.Error!CompletedRanges { // `consumed` is the next missing shred index, but shred `i` existing in @@ -1164,10 +1168,14 @@ fn getCompletedDataRanges( std.debug.assert(!completed_data_indexes.contains(consumed)); var ranges = CompletedRanges.init(allocator); var begin: u32 = start_index; - for (completed_data_indexes.range(start_index, consumed)) |index| { + + var iter = completed_data_indexes.iteratorRanged(start_index, consumed, .start); + while (iter.next()) |entry| { + const index = entry.key_ptr.*; try ranges.append(.{ begin, index + 1 }); begin = index + 1; } + return ranges; } @@ -1325,10 +1333,10 @@ pub fn getSlotsSince( for (slots) |slot| { if (try self.ledger.db.get(allocator, schema.slot_meta, slot)) |meta| { var child_slots = meta.child_slots; - errdefer child_slots.deinit(); + errdefer child_slots.deinit(allocator); var cdi = meta.completed_data_indexes; - cdi.deinit(); - try map.put(allocator, slot, child_slots.moveToUnmanaged()); + cdi.deinit(allocator); + try map.put(allocator, slot, child_slots); } } return map; @@ -1629,19 +1637,22 @@ pub const AncestorIterator = struct { } pub fn next(self: *AncestorIterator) !?Slot { - if (self.next_slot) |slot| { - if (slot == 0) { - self.next_slot = null; - } else if (try self.db.get(self.allocator, schema.slot_meta, slot)) |slot_meta| { - defer slot_meta.deinit(); - self.next_slot = slot_meta.parent_slot; - } else { - self.next_slot = null; - } + const slot = self.next_slot orelse return null; + if (slot == 0) { + self.next_slot = null; return slot; } - return null; + + var slot_meta: SlotMeta = try self.db.get(self.allocator, schema.slot_meta, slot) orelse { + self.next_slot = null; + return slot; + }; + defer slot_meta.deinit(self.allocator); + + self.next_slot = slot_meta.parent_slot; + + return slot; } }; @@ -1890,7 +1901,7 @@ test slotMetaIterator { var slot_metas = ArrayList(SlotMeta).init(allocator); defer { for (slot_metas.items) |*slot_meta| { - slot_meta.deinit(); + slot_meta.deinit(allocator); } slot_metas.deinit(); } @@ -1901,13 +1912,13 @@ test slotMetaIterator { const roots: [3]Slot = .{ 1, 2, 3 }; var parent_slot: ?Slot = null; for (roots, 0..) |slot, i| { - var slot_meta = SlotMeta.init(allocator, slot, parent_slot); + var slot_meta = SlotMeta.init(slot, parent_slot); // ensure isFull() is true slot_meta.last_index = 1; slot_meta.consecutive_received_from_0 = slot_meta.last_index.? + 1; // update next slots if (i + 1 < roots.len) { - try slot_meta.child_slots.append(roots[i + 1]); + try slot_meta.child_slots.append(allocator, roots[i + 1]); } try write_batch.put(schema.slot_meta, slot_meta.slot, slot_meta); // connect the chain @@ -1923,7 +1934,7 @@ test slotMetaIterator { var index: u64 = 0; while (try iter.next()) |entry| { var slot_meta = entry[1]; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqual(slot_metas.items[index].slot, slot_meta.slot); try std.testing.expectEqual(slot_metas.items[index].last_index, slot_meta.last_index); @@ -1972,14 +1983,14 @@ test slotRangeConnected { // 1 -> 2 -> 3 var parent_slot: ?Slot = null; for (roots, 0..) |slot, i| { - var slot_meta = SlotMeta.init(allocator, slot, parent_slot); - defer slot_meta.deinit(); + var slot_meta = SlotMeta.init(slot, parent_slot); + defer slot_meta.deinit(allocator); // ensure isFull() is true slot_meta.last_index = 1; slot_meta.consecutive_received_from_0 = slot_meta.last_index.? + 1; // update next slots if (i + 1 < roots.len) { - try slot_meta.child_slots.append(roots[i + 1]); + try slot_meta.child_slots.append(allocator, roots[i + 1]); } try write_batch.put(schema.slot_meta, slot_meta.slot, slot_meta); // connect the chain @@ -1995,8 +2006,8 @@ test slotRangeConnected { try std.testing.expectEqual(true, is_connected); // insert a non-full last_slot - var slot_meta = SlotMeta.init(allocator, 4, parent_slot); - defer slot_meta.deinit(); + var slot_meta = SlotMeta.init(4, parent_slot); + defer slot_meta.deinit(allocator); // ensure isFull() is FALSE slot_meta.last_index = 1; try write_batch2.put(schema.slot_meta, slot_meta.slot, slot_meta); @@ -2019,7 +2030,7 @@ test highestSlot { { // insert a shred const shred_slot = 10; - var slot_meta = SlotMeta.init(allocator, shred_slot, null); + var slot_meta = SlotMeta.init(shred_slot, null); slot_meta.last_index = 21; slot_meta.received = 1; @@ -2039,7 +2050,7 @@ test highestSlot { { // insert another shred at a higher slot - var slot_meta2 = SlotMeta.init(allocator, 100, null); + var slot_meta2 = SlotMeta.init(100, null); slot_meta2.last_index = 21; slot_meta2.received = 1; @@ -2076,7 +2087,7 @@ test lowestSlot { shred.data.common.index = shred_index; // insert a shred - var slot_meta = SlotMeta.init(allocator, shred_slot, null); + var slot_meta = SlotMeta.init(shred_slot, null); slot_meta.last_index = 21; slot_meta.received = 1; @@ -2162,7 +2173,7 @@ test findMissingDataIndexes { shred.data.common.variant = variant; try ledger_mod.shred.overwriteShredForTest(allocator, &shred, &(.{2} ** 100)); - var slot_meta = SlotMeta.init(allocator, shred_slot, null); + var slot_meta = SlotMeta.init(shred_slot, null); slot_meta.last_index = 4; var write_batch = try state.db.initWriteBatch(); @@ -2661,7 +2672,7 @@ test getConfirmedSignaturesForAddress { try write_batch.put(schema.rooted_slots, slot, true); // add a slot meta for genesis check - var slot_meta = SlotMeta.init(allocator, 0, null); + var slot_meta = SlotMeta.init(0, null); slot_meta.received = 1; try write_batch.put(schema.slot_meta, 0, slot_meta); diff --git a/src/ledger/ResultWriter.zig b/src/ledger/ResultWriter.zig index a1f38cb860..fe9da46c27 100644 --- a/src/ledger/ResultWriter.zig +++ b/src/ledger/ResultWriter.zig @@ -361,8 +361,8 @@ pub fn setAndChainConnectedOnRootAndNextSlots( root: Slot, ) !void { var root_slot_meta: SlotMeta = try self.ledger.db.get(allocator, schema.slot_meta, root) orelse - SlotMeta.init(allocator, root, null); - defer root_slot_meta.deinit(); + SlotMeta.init(root, null); + defer root_slot_meta.deinit(allocator); // If the slot was already connected, there is nothing to do as this slot's // children are also assumed to be appropriately connected @@ -391,7 +391,7 @@ pub fn setAndChainConnectedOnRootAndNextSlots( self.logger.err().logf("Slot {} is a child but has no SlotMeta in ledger", .{slot}); return error.CorruptedLedger; }; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); if (slot_meta.setParentConnected()) { try child_slots.appendSlice(slot_meta.child_slots.items); @@ -608,9 +608,9 @@ test "scanAndFixRoots" { const roots: [2]Slot = .{ 1, 3 }; try setRoots(&result_writer, &roots); - const slot_meta_1 = SlotMeta.init(allocator, 1, null); - const slot_meta_2 = SlotMeta.init(allocator, 2, 1); - const slot_meta_3 = SlotMeta.init(allocator, 3, 2); + const slot_meta_1 = SlotMeta.init(1, null); + const slot_meta_2 = SlotMeta.init(2, 1); + const slot_meta_3 = SlotMeta.init(3, 2); var write_batch = try state.db.initWriteBatch(); defer write_batch.deinit(); @@ -639,7 +639,7 @@ test "setAndChainConnectedOnRootAndNextSlots" { // 1 is a root const roots: [1]Slot = .{1}; try setRoots(&result_writer, &roots); - const slot_meta_1 = SlotMeta.init(allocator, 1, null); + const slot_meta_1 = SlotMeta.init(1, null); var write_batch = try state.db.initWriteBatch(); defer write_batch.deinit(); try write_batch.put(schema.slot_meta, slot_meta_1.slot, slot_meta_1); @@ -662,15 +662,15 @@ test "setAndChainConnectedOnRootAndNextSlots" { try setRoots(&result_writer, &other_roots); for (other_roots, 0..) |slot, i| { - var slot_meta = SlotMeta.init(allocator, slot, parent_slot); - defer slot_meta.deinit(); + var slot_meta = SlotMeta.init(slot, parent_slot); + defer slot_meta.deinit(allocator); // ensure isFull() is true slot_meta.last_index = 1; slot_meta.consecutive_received_from_0 = slot_meta.last_index.? + 1; // update next slots if (i + 1 < other_roots.len) { - try slot_meta.child_slots.append(other_roots[i + 1]); + try slot_meta.child_slots.append(allocator, other_roots[i + 1]); } try write_batch2.put(schema.slot_meta, slot_meta.slot, slot_meta); @@ -684,7 +684,7 @@ test "setAndChainConnectedOnRootAndNextSlots" { for (other_roots) |slot| { var db_slot_meta = (try state.db.get(allocator, schema.slot_meta, slot)) orelse return error.MissingSlotMeta; - defer db_slot_meta.deinit(); + defer db_slot_meta.deinit(allocator); try std.testing.expectEqual(true, db_slot_meta.isConnected()); } } @@ -704,24 +704,24 @@ test "setAndChainConnectedOnRootAndNextSlots: disconnected" { const roots: [3]Slot = .{ 1, 2, 3 }; try setRoots(&result_writer, &roots); - var slot_meta_1 = SlotMeta.init(allocator, 1, null); - defer slot_meta_1.deinit(); + var slot_meta_1 = SlotMeta.init(1, null); + defer slot_meta_1.deinit(allocator); slot_meta_1.last_index = 1; slot_meta_1.consecutive_received_from_0 = 1 + 1; - try slot_meta_1.child_slots.append(2); + try slot_meta_1.child_slots.append(allocator, 2); try write_batch.put(schema.slot_meta, slot_meta_1.slot, slot_meta_1); // 2 is not full - var slot_meta_2 = SlotMeta.init(allocator, 2, 1); - defer slot_meta_2.deinit(); + var slot_meta_2 = SlotMeta.init(2, 1); + defer slot_meta_2.deinit(allocator); slot_meta_2.last_index = 1; slot_meta_2.consecutive_received_from_0 = 0; // ! NOT FULL - try slot_meta_2.child_slots.append(3); + try slot_meta_2.child_slots.append(allocator, 3); try write_batch.put(schema.slot_meta, slot_meta_2.slot, slot_meta_2); // 3 is full - var slot_meta_3 = SlotMeta.init(allocator, 3, 2); - defer slot_meta_3.deinit(); + var slot_meta_3 = SlotMeta.init(3, 2); + defer slot_meta_3.deinit(allocator); slot_meta_3.last_index = 1; slot_meta_3.consecutive_received_from_0 = 1 + 1; try write_batch.put(schema.slot_meta, slot_meta_3.slot, slot_meta_3); @@ -733,18 +733,18 @@ test "setAndChainConnectedOnRootAndNextSlots: disconnected" { // should be connected var db_slot_meta_1 = (try state.db.get(allocator, schema.slot_meta, 1)) orelse return error.MissingSlotMeta; - defer db_slot_meta_1.deinit(); + defer db_slot_meta_1.deinit(allocator); try std.testing.expectEqual(true, db_slot_meta_1.isConnected()); var db_slot_meta_2: SlotMeta = (try state.db.get(allocator, schema.slot_meta, 2)) orelse return error.MissingSlotMeta; - defer db_slot_meta_2.deinit(); + defer db_slot_meta_2.deinit(allocator); try std.testing.expectEqual(true, db_slot_meta_2.isParentConnected()); try std.testing.expectEqual(false, db_slot_meta_2.isConnected()); var db_slot_meta_3: SlotMeta = (try state.db.get(allocator, schema.slot_meta, 3)) orelse return error.MissingSlotMeta; - defer db_slot_meta_3.deinit(); + defer db_slot_meta_3.deinit(allocator); try std.testing.expectEqual(false, db_slot_meta_3.isParentConnected()); try std.testing.expectEqual(false, db_slot_meta_3.isConnected()); } diff --git a/src/ledger/benchmarks.zig b/src/ledger/benchmarks.zig index d7c0a9b0d2..d97a3a1f70 100644 --- a/src/ledger/benchmarks.zig +++ b/src/ledger/benchmarks.zig @@ -379,14 +379,14 @@ pub const BenchmarkLedgerSlow = struct { const slot_per_epoch = 432_000; var parent_slot: ?Slot = null; for (1..(slot_per_epoch + 1)) |slot| { - var slot_meta = SlotMeta.init(allocator, slot, parent_slot); - defer slot_meta.deinit(); + var slot_meta = SlotMeta.init(slot, parent_slot); + defer slot_meta.deinit(allocator); // ensure isFull() is true slot_meta.last_index = 1; slot_meta.consecutive_received_from_0 = slot_meta.last_index.? + 1; // update next slots if (slot < (slot_per_epoch + 1)) { - try slot_meta.child_slots.append(slot + 1); + try slot_meta.child_slots.append(allocator, slot + 1); } try write_batch.put(schema.slot_meta, slot_meta.slot, slot_meta); // connect the chain diff --git a/src/ledger/cleanup_service.zig b/src/ledger/cleanup_service.zig index 963da7caca..7c1086a8c7 100644 --- a/src/ledger/cleanup_service.zig +++ b/src/ledger/cleanup_service.zig @@ -429,12 +429,12 @@ test "findSlotsToClean" { defer state.deinit(); // set highest and lowest slot by inserting slot_meta - var lowest_slot_meta = ledger.meta.SlotMeta.init(allocator, 10, null); - defer lowest_slot_meta.deinit(); + var lowest_slot_meta = ledger.meta.SlotMeta.init(10, null); + defer lowest_slot_meta.deinit(allocator); lowest_slot_meta.received = 10; - var highest_slot_meta = ledger.meta.SlotMeta.init(allocator, 20, null); - defer highest_slot_meta.deinit(); + var highest_slot_meta = ledger.meta.SlotMeta.init(20, null); + defer highest_slot_meta.deinit(allocator); highest_slot_meta.received = 20; { diff --git a/src/ledger/database/hashmap.zig b/src/ledger/database/hashmap.zig index 9b7be779eb..f7d6d44e8d 100644 --- a/src/ledger/database/hashmap.zig +++ b/src/ledger/database/hashmap.zig @@ -46,27 +46,34 @@ pub fn SharedHashMapDB(comptime column_families: []const ColumnFamily) type { defer allocator.free(actual_path); const disk_allocator = try allocator.create(DiskMemoryAllocator); + errdefer allocator.destroy(disk_allocator); disk_allocator.* = DiskMemoryAllocator{ .dir = try std.fs.cwd().makeOpenPath(actual_path, .{}), .logger = logger.withScope(@typeName(DiskMemoryAllocator)), .mmap_ratio = 8, }; + const batch_allocator = try allocator.create(BatchAllocator); + errdefer allocator.destroy(batch_allocator); batch_allocator.* = BatchAllocator{ .backing_allocator = disk_allocator.allocator(), .batch_size = 1 << 30, }; var maps = try allocator.alloc(SharedHashMap, column_families.len); - const lock = try allocator.create(RwLock); - lock.* = .{}; errdefer { for (maps) |*m| m.deinit(); allocator.free(maps); } + inline for (0..column_families.len) |i| { - maps[i] = try SharedHashMap.init(batch_allocator.allocator()); + maps[i] = SharedHashMap.init(batch_allocator.allocator()); } + + const lock = try allocator.create(RwLock); + errdefer allocator.destroy(lock); + lock.* = .{}; + return .{ .fast_allocator = allocator, .storage_allocator = batch_allocator.allocator(), @@ -222,13 +229,16 @@ pub fn SharedHashMapDB(comptime column_families: []const ColumnFamily) type { }, .delete_range => |delete_range_ix| { const cf_index, const start, const end = delete_range_ix; - const keys, _ = self.maps[cf_index].map.range(start, end); - const to_delete = try batch.fast_allocator.alloc([]const u8, keys.len); - defer batch.fast_allocator.free(to_delete); - for (keys, 0..) |key, i| { - to_delete[i] = key; + + var iter = self.maps[cf_index].map.iteratorRanged(start, end, .start); + var to_delete: std.ArrayListUnmanaged([]const u8) = .empty; + defer to_delete.deinit(batch.fast_allocator); + + while (iter.next()) |entry| { + try to_delete.append(batch.fast_allocator, entry.key_ptr.*); } - for (to_delete) |delete_key| { + + for (to_delete.items) |delete_key| { self.maps[cf_index].delete(delete_key); } }, @@ -341,40 +351,57 @@ pub fn SharedHashMapDB(comptime column_families: []const ColumnFamily) type { shared_map.lock.lockShared(); defer shared_map.lock.unlockShared(); - const keys, const vals = if (start) |start_| b: { - const search_bytes = try key_serializer.serializeAlloc(self.fast_allocator, start_); - defer self.fast_allocator.free(search_bytes); - break :b switch (direction) { - .forward => map.rangeCustom(.{ .inclusive = search_bytes }, null), - .reverse => map.rangeCustom(null, .{ .inclusive = search_bytes }), - }; - } else map.items(); - std.debug.assert(keys.len == vals.len); - - const copied_keys = try self.storage_allocator.alloc([]const u8, keys.len); - errdefer self.storage_allocator.free(copied_keys); - const copied_vals = try self.storage_allocator.alloc(RcSlice(u8), vals.len); - errdefer self.storage_allocator.free(copied_vals); - for (0..keys.len) |i| { - errdefer for (0..i) |n| { - self.storage_allocator.free(copied_keys[n]); - copied_vals[n].deinit(self.storage_allocator); - }; - copied_keys[i] = try self.storage_allocator.dupe(u8, keys[i]); - errdefer self.storage_allocator.free(copied_keys[i]); - copied_vals[i] = vals[i].acquire(); + var iter = if (start == null) + map.iterator() + else switch (direction) { + .forward => map.iteratorRanged( + try key_serializer.serializeAlloc(self.fast_allocator, start.?), + null, + .start, + ), + .reverse => map.iteratorRanged( + null, + try key_serializer.serializeAlloc(self.fast_allocator, start.?), + .start, + ), + }; + defer { + if (iter.start) |start_bytes| self.fast_allocator.free(start_bytes); + if (iter.end) |end_bytes| self.fast_allocator.free(end_bytes); + } + + const len = blk: { + var iter_copied = iter; + break :blk iter_copied.countForwardsInclusive(); + }; + + var copied_keys: std.ArrayListUnmanaged([]const u8) = try .initCapacity(self.storage_allocator, len); + errdefer { + for (copied_keys.items) |key| self.storage_allocator.free(key); + copied_keys.deinit(self.storage_allocator); + } + + var copied_vals: std.ArrayListUnmanaged(RcSlice(u8)) = try .initCapacity(self.storage_allocator, len); + errdefer { + for (copied_vals.items) |val| val.deinit(self.storage_allocator); + copied_vals.deinit(self.storage_allocator); + } + + while (iter.nextInclusive()) |entry| { + copied_keys.appendAssumeCapacity(try self.storage_allocator.dupe(u8, entry.key_ptr.*)); + copied_vals.appendAssumeCapacity(entry.value_ptr.acquire()); } return .{ .allocator = self.fast_allocator, .storage_allocator = self.storage_allocator, - .keys = copied_keys, - .vals = copied_vals, + .keys = try copied_keys.toOwnedSlice(self.storage_allocator), + .vals = try copied_vals.toOwnedSlice(self.storage_allocator), .cursor = switch (direction) { .forward => 0, - .reverse => keys.len, + .reverse => len, }, - .size = keys.len, + .size = len, }; } @@ -463,26 +490,26 @@ fn serializeValue(allocator: Allocator, value: anytype) !RcSlice(u8) { const SharedHashMap = struct { /// must be the same as SharedHashmapDB.storage_allocator allocator: Allocator, - map: SortedMap([]const u8, RcSlice(u8)), + map: SortedMap([]const u8, RcSlice(u8), .{}), lock: RwLock, const Self = @This(); - fn init(allocator: Allocator) Allocator.Error!Self { + fn init(allocator: Allocator) Self { return .{ .allocator = allocator, - .map = SortedMap([]const u8, RcSlice(u8)).init(allocator), + .map = .empty, .lock = .{}, }; } pub fn deinit(self: *Self) void { - const keys, const values = self.map.items(); - for (keys, values) |key, value| { - self.allocator.free(key); - value.deinit(self.allocator); + var iter = self.map.iterator(); + while (iter.next()) |entry| { + self.allocator.free(entry.key_ptr.*); + entry.value_ptr.deinit(self.allocator); } - self.map.deinit(); + self.map.deinit(self.allocator); } pub fn count(self: *Self) usize { @@ -495,7 +522,7 @@ const SharedHashMap = struct { pub fn put(self: *Self, key: []const u8, value: RcSlice(u8)) Allocator.Error!void { self.lock.lock(); defer self.lock.unlock(); - const entry = try self.map.getOrPut(key); + const entry = try self.map.getOrPut(self.allocator, key); if (entry.found_existing) { self.allocator.free(key); entry.value_ptr.deinit(self.allocator); @@ -512,7 +539,7 @@ const SharedHashMap = struct { const key, const value = lock: { self.lock.lock(); defer self.lock.unlock(); - const entry = self.map.fetchSwapRemove(key_) orelse return; + const entry = self.map.fetchRemove(key_) orelse return; break :lock .{ entry.key, entry.value }; }; self.allocator.free(key); diff --git a/src/ledger/meta.zig b/src/ledger/meta.zig index 6a0abe3c85..a100683afc 100644 --- a/src/ledger/meta.zig +++ b/src/ledger/meta.zig @@ -33,16 +33,18 @@ pub const SlotMeta = struct { parent_slot: ?Slot, /// The list of slots, each of which contains a block that derives /// from this one. - child_slots: std.ArrayList(Slot), + child_slots: std.ArrayListUnmanaged(Slot), /// Connected status flags of this slot connected_flags: ConnectedFlags, /// Shreds indices which are marked data complete. That is, those that have the /// [`ShredFlags::DATA_COMPLETE_SHRED`][`crate::shred::ShredFlags::DATA_COMPLETE_SHRED`] set. - completed_data_indexes: SortedSet(u32), + completed_data_indexes: DataIndexes, + + const DataIndexes = SortedSet(u32, .{}); const Self = @This(); - pub fn init(allocator: std.mem.Allocator, slot: Slot, parent_slot: ?Slot) Self { + pub fn init(slot: Slot, parent_slot: ?Slot) Self { const connected_flags = if (slot == 0) // Slot 0 is the start, mark it as having its' parent connected // such that slot 0 becoming full will be updated as connected @@ -57,21 +59,22 @@ pub const SlotMeta = struct { .received = 0, .first_shred_timestamp_milli = 0, .last_index = null, - .child_slots = std.ArrayList(Slot).init(allocator), - .completed_data_indexes = SortedSet(u32).init(allocator), + .child_slots = .empty, + .completed_data_indexes = .empty, }; } - pub fn deinit(self: Self) void { - self.child_slots.deinit(); - self.completed_data_indexes.deinit(); + pub fn deinit(self: *Self, allocator: std.mem.Allocator) void { + self.child_slots.deinit(allocator); + self.completed_data_indexes.deinit(allocator); } pub fn clone(self: Self, allocator: Allocator) Allocator.Error!Self { - var child_slots = try std.ArrayList(Slot).initCapacity( + var child_slots = try std.ArrayListUnmanaged(Slot).initCapacity( allocator, self.child_slots.items.len, ); + errdefer child_slots.deinit(allocator); child_slots.appendSliceAssumeCapacity(self.child_slots.items); return .{ .slot = self.slot, @@ -82,7 +85,7 @@ pub const SlotMeta = struct { .first_shred_timestamp_milli = self.first_shred_timestamp_milli, .last_index = self.last_index, .child_slots = child_slots, - .completed_data_indexes = try self.completed_data_indexes.clone(), + .completed_data_indexes = try self.completed_data_indexes.clone(allocator), }; } @@ -218,8 +221,15 @@ pub const ErasureMeta = struct { } { const c_start, const c_end = self.codeShredsIndices(); const d_start, const d_end = self.dataShredsIndices(); - const num_code = index.code_index.range(c_start, c_end).len; - const num_data = index.data_index.range(d_start, d_end).len; + + const num_code = blk: { + var iter = index.code_index.iteratorRanged(c_start, c_end, .start); + break :blk iter.countForwards(); + }; + const num_data = blk: { + var iter = index.data_index.iteratorRanged(d_start, d_end, .start); + break :blk iter.countForwards(); + }; const data_missing = self.config.num_data -| num_data; const num_needed = data_missing -| num_code; @@ -267,21 +277,21 @@ pub const Index = struct { data_index: ShredIndex, code_index: ShredIndex, - pub fn init(allocator: std.mem.Allocator, slot: Slot) Index { + pub fn init(slot: Slot) Index { return .{ .slot = slot, - .data_index = ShredIndex.init(allocator), - .code_index = ShredIndex.init(allocator), + .data_index = .empty, + .code_index = .empty, }; } - pub fn deinit(self: *Index) void { - self.data_index.deinit(); - self.code_index.deinit(); + pub fn deinit(self: *Index, allocator: std.mem.Allocator) void { + self.data_index.deinit(allocator); + self.code_index.deinit(allocator); } }; -pub const ShredIndex = SortedSet(u64); +pub const ShredIndex = SortedSet(u64, .{}); pub const TransactionStatusMeta = sig.ledger.transaction_status.TransactionStatusMeta; diff --git a/src/ledger/shred.zig b/src/ledger/shred.zig index ed089fa7c6..135e3d410a 100644 --- a/src/ledger/shred.zig +++ b/src/ledger/shred.zig @@ -567,6 +567,11 @@ pub const ErasureSetId = struct { /// aka "fec_set_index" erasure_set_index: u64, + pub const empty: ErasureSetId = .{ + .slot = std.math.maxInt(Slot), + .erasure_set_index = std.math.maxInt(u64), + }; + pub fn order(a: ErasureSetId, b: ErasureSetId) std.math.Order { if (a.slot == b.slot and a.erasure_set_index == b.erasure_set_index) { return .eq; diff --git a/src/ledger/shred_inserter/ShredInserter.zig b/src/ledger/shred_inserter/ShredInserter.zig index cd915d73e3..dd4d84e0bb 100644 --- a/src/ledger/shred_inserter/ShredInserter.zig +++ b/src/ledger/shred_inserter/ShredInserter.zig @@ -44,6 +44,8 @@ const handleChaining = lib.slot_chaining.handleChaining; const recover = lib.recovery.recover; const newlinesToSpaces = sig.utils.fmt.newlinesToSpaces; +const DataIndexes = SortedSet(u32, .{}); + const DEFAULT_TICKS_PER_SECOND = sig.core.time.DEFAULT_TICKS_PER_SECOND; const Logger = sig.trace.Logger("shred_inserter"); @@ -209,6 +211,7 @@ pub fn insertShreds( errdefer newly_completed_data_sets.deinit(); for (shreds, is_repaired) |shred, is_repair| { const shred_source: ShredSource = if (is_repair) .repaired else .turbine; + switch (shred) { .data => |data_shred| { if (options.shred_tracker) |tracker| { @@ -365,8 +368,13 @@ pub fn insertShreds( // var merkle_chaining_timer = Timer.start(); - const em0_keys, const em0_values = state.erasure_metas.items(); - for (em0_keys, em0_values) |erasure_set, working_em| if (working_em == .dirty) { + var iter = state.erasure_metas.iterator(); + while (iter.next()) |entry| { + const erasure_set = entry.key_ptr.*; + const working_em = entry.value_ptr.*; + + if (working_em != .dirty) continue; + const slot = erasure_set.slot; const erasure_meta: ErasureMeta = working_em.dirty; if (try self.hasDuplicateShredsInSlot(slot)) { @@ -386,7 +394,7 @@ pub fn insertShreds( erasure_meta, state.merkleRootMetas(), ); - }; + } ////////////////////////////////////////////////////// // check backward chaining for each merkle root @@ -464,7 +472,12 @@ fn checkInsertCodeShred( // .record_shred(shred.slot(), shred.erasure_set_index(), shred_source, None); _ = shred_source; - const was_inserted = !std.meta.isError(insertCodeShred(index_meta, shred, write_batch)); + const was_inserted = !std.meta.isError(insertCodeShred( + state.allocator, + index_meta, + shred, + write_batch, + )); if (was_inserted) { index_meta_working_set_entry.did_insert_occur = true; @@ -728,6 +741,7 @@ fn checkInsertDataShred( /// agave: insert_coding_shred fn insertCodeShred( + allocator: std.mem.Allocator, index_meta: *meta.Index, shred: CodeShred, write_batch: *WriteBatch, @@ -738,7 +752,7 @@ fn insertCodeShred( assertOk(shred.sanitize()); try write_batch.put(schema.code_shred, .{ slot, shred_index }, shred.payload); - try index_meta.code_index.put(shred_index); + try index_meta.code_index.put(allocator, shred_index, {}); } /// Check if the shred already exists in ledger @@ -856,7 +870,7 @@ fn insertDataShred( } else slot_meta.consecutive_received_from_0; try write_batch.put(schema.data_shred, .{ slot, index }, shred.payload); - try data_index.put(index); + try data_index.put(allocator, index, {}); var newly_completed_data_sets = ArrayList(CompletedDataSetInfo).init(allocator); const shred_indices = try updateSlotMeta( @@ -899,7 +913,7 @@ fn sendSlotFullTiming(self: *const ShredInserter, slot: Slot) void { fn tryShredRecovery( self: *const ShredInserter, allocator: Allocator, - erasure_metas: *SortedMap(ErasureSetId, WorkingEntry(ErasureMeta)), + erasure_metas: *SortedMap(ErasureSetId, WorkingEntry(ErasureMeta), .{}), index_working_set: *AutoHashMap(u64, IndexMetaWorkingSetEntry), shred_store: ShredWorkingStore, reed_solomon_cache: *ReedSolomonCache, @@ -909,13 +923,17 @@ fn tryShredRecovery( // 2. For new data shreds, check if an erasure set exists. If not, don't try recovery // 3. Before trying recovery, check if enough number of shreds have been received // 3a. Enough number of shreds = (#data + #code shreds) > erasure.num_data - const keys, const values = erasure_metas.items(); // let index = &mut index_meta_entry.index; - for (keys, values) |erasure_set, *working_erasure_meta| { - const erasure_meta = working_erasure_meta.asRef(); + + var iter = erasure_metas.iterator(); + while (iter.next()) |entry| { + const erasure_meta: *const ErasureMeta = entry.value_ptr.asRef(); + const erasure_set = entry.key_ptr.*; + var index_meta_entry = index_working_set.get(erasure_set.slot) orelse { return error.Unwrap; // TODO: consider all the unwraps }; + switch (erasure_meta.status(&index_meta_entry.index)) { .can_recover => return try recoverShreds( .from(self.logger), @@ -948,6 +966,7 @@ fn tryShredRecovery( }, } } + return std.ArrayList(Shred).init(allocator); } @@ -1099,24 +1118,29 @@ fn updateCompletedDataIndexes( new_shred_index: u32, received_data_shreds: *meta.ShredIndex, /// Shreds indices which are marked data complete. - completed_data_indexes: *SortedSet(u32), + completed_data_indexes: *DataIndexes, ) Allocator.Error!ArrayList([2]u32) { var shred_indices = ArrayList(u32).init(allocator); defer shred_indices.deinit(); - const subslice = completed_data_indexes.range(null, new_shred_index); - const start_shred_index = if (subslice.len == 0) 0 else subslice[subslice.len - 1]; + + const start_shred_index = blk: { + var iter = completed_data_indexes.iteratorRanged(0, new_shred_index, .end); + break :blk if (iter.prev()) |entry| entry.key_ptr.* else 0; + }; + // Consecutive entries i, k, j in this vector represent potential ranges [i, k), // [k, j) that could be completed data ranges try shred_indices.append(start_shred_index); // `new_shred_index` is data complete, so need to insert here into the // `completed_data_indexes` if (is_last_in_data) { - try completed_data_indexes.put(new_shred_index); + try completed_data_indexes.put(allocator, new_shred_index, {}); try shred_indices.append(new_shred_index + 1); } - const new_subslice = completed_data_indexes.range(new_shred_index + 1, null); - if (new_subslice.len != 0) { - try shred_indices.append(new_subslice[0]); + + { + var iter = completed_data_indexes.iteratorRanged(new_shred_index + 1, null, .start); + if (iter.next()) |entry| try shred_indices.append(entry.key_ptr.*); } var ret = ArrayList([2]u32).init(allocator); @@ -1125,7 +1149,9 @@ fn updateCompletedDataIndexes( const begin = shred_indices.items[i]; const end = shred_indices.items[i + 1]; const num_shreds: usize = @intCast(end - begin); - if (received_data_shreds.range(begin, end).len == num_shreds) { + + var iter = received_data_shreds.iteratorRanged(begin, end, .start); + if (iter.countForwards() == num_shreds) { try ret.append(.{ begin, end - 1 }); } i += 1; @@ -1345,7 +1371,7 @@ test "chaining basic" { result.deinit(); { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 1)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{}, slot_meta.child_slots.items); try std.testing.expect(!slot_meta.isConnected()); try std.testing.expectEqual(0, slot_meta.parent_slot); @@ -1357,7 +1383,7 @@ test "chaining basic" { result.deinit(); { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 1)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{2}, slot_meta.child_slots.items); try std.testing.expect(!slot_meta.isConnected()); // since 0 is not yet inserted try std.testing.expectEqual(0, slot_meta.parent_slot); @@ -1365,7 +1391,7 @@ test "chaining basic" { } { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 2)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{}, slot_meta.child_slots.items); try std.testing.expect(!slot_meta.isConnected()); // since 0 is not yet inserted try std.testing.expectEqual(1, slot_meta.parent_slot); @@ -1377,7 +1403,7 @@ test "chaining basic" { result.deinit(); { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 0)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{1}, slot_meta.child_slots.items); try std.testing.expect(slot_meta.isConnected()); try std.testing.expectEqual(0, slot_meta.parent_slot); @@ -1385,7 +1411,7 @@ test "chaining basic" { } { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 1)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{2}, slot_meta.child_slots.items); try std.testing.expect(slot_meta.isConnected()); try std.testing.expectEqual(0, slot_meta.parent_slot); @@ -1393,7 +1419,7 @@ test "chaining basic" { } { var slot_meta = (try state.ledger.db.get(allocator, schema.slot_meta, 2)).?; - defer slot_meta.deinit(); + defer slot_meta.deinit(allocator); try std.testing.expectEqualSlices(u64, &.{}, slot_meta.child_slots.items); try std.testing.expect(slot_meta.isConnected()); try std.testing.expectEqual(1, slot_meta.parent_slot); @@ -1586,6 +1612,7 @@ test "recovery" { for (data_shreds) |data_shred| { const key = .{ data_shred.data.common.slot, data_shred.data.common.index }; const actual_shred = try state.ledger.db.getBytes(schema.data_shred, key); + defer actual_shred.?.deinit(); try std.testing.expectEqualSlices(u8, data_shred.payload(), actual_shred.?.data); } diff --git a/src/ledger/shred_inserter/slot_chaining.zig b/src/ledger/shred_inserter/slot_chaining.zig index 9302a60385..68a45e3af9 100644 --- a/src/ledger/shred_inserter/slot_chaining.zig +++ b/src/ledger/shred_inserter/slot_chaining.zig @@ -16,7 +16,6 @@ const SlotMeta = ledger.meta.SlotMeta; const SlotMetaWorkingSetEntry = shred_inserter.working_state.SlotMetaWorkingSetEntry; const WriteBatch = LedgerDB.WriteBatch; -const deinitMapRecursive = shred_inserter.working_state.deinitMapRecursive; const isNewlyCompletedSlot = shred_inserter.working_state.isNewlyCompletedSlot; /// agave: handle_chaining @@ -51,13 +50,19 @@ pub fn handleChaining( for (keys[delete_i..count]) |k| { if (working_set.fetchRemove(k)) |entry| { var slot_meta_working_set_entry = entry.value; - slot_meta_working_set_entry.deinit(); + slot_meta_working_set_entry.deinit(allocator); } } // handle chaining var new_chained_slots = AutoHashMap(u64, SlotMeta).init(allocator); - defer deinitMapRecursive(&new_chained_slots); + defer { + var deinit_iter = new_chained_slots.iterator(); + while (deinit_iter.next()) |entry| { + entry.value_ptr.deinit(allocator); + } + new_chained_slots.deinit(); + } for (keys[0..keep_i]) |slot| { try handleChainingForSlot( allocator, @@ -111,7 +116,7 @@ fn handleChainingForSlot( // This is a newly inserted slot/orphan so run the chaining logic to link it to a // newly discovered parent - try chainNewSlotToPrevSlot(prev_slot_meta, slot, slot_meta); + try chainNewSlotToPrevSlot(allocator, prev_slot_meta, slot, slot_meta); // If the parent of `slot` is a newly inserted orphan, insert it into the orphans // column family @@ -170,7 +175,7 @@ fn findSlotMetaElseCreate( entry.value_ptr.* = if (try db.get(allocator, schema.slot_meta, slot)) |m| m else - SlotMeta.init(allocator, slot, null); + SlotMeta.init(slot, null); return entry.value_ptr; } @@ -220,11 +225,12 @@ fn traverseChildrenMut( /// agave: chain_new_slot_to_prev_slot fn chainNewSlotToPrevSlot( + allocator: Allocator, prev_slot_meta: *SlotMeta, current_slot: Slot, current_slot_meta: *SlotMeta, ) !void { - try prev_slot_meta.child_slots.append(current_slot); + try prev_slot_meta.child_slots.append(allocator, current_slot); if (prev_slot_meta.isConnected()) { _ = current_slot_meta.setParentConnected(); } diff --git a/src/ledger/shred_inserter/working_state.zig b/src/ledger/shred_inserter/working_state.zig index 28c203e578..8b7310af59 100644 --- a/src/ledger/shred_inserter/working_state.zig +++ b/src/ledger/shred_inserter/working_state.zig @@ -81,7 +81,7 @@ pub const PendingInsertShredsState = struct { db: *LedgerDB, write_batch: WriteBatch, just_inserted_shreds: AutoHashMap(ShredId, Shred), - erasure_metas: SortedMap(ErasureSetId, WorkingEntry(ErasureMeta)), + erasure_metas: SortedMap(ErasureSetId, WorkingEntry(ErasureMeta), .{}), merkle_root_metas: AutoHashMap(ErasureSetId, WorkingEntry(MerkleRootMeta)), slot_meta_working_set: AutoHashMap(u64, SlotMetaWorkingSetEntry), index_working_set: AutoHashMap(u64, IndexMetaWorkingSetEntry), @@ -103,12 +103,12 @@ pub const PendingInsertShredsState = struct { .db = db, .logger = logger.withScope(@typeName(Self)), .write_batch = try db.initWriteBatch(), - .just_inserted_shreds = AutoHashMap(ShredId, Shred).init(allocator), // TODO capacity = shreds.len - .erasure_metas = SortedMap(ErasureSetId, WorkingEntry(ErasureMeta)).init(allocator), - .merkle_root_metas = AutoHashMap(ErasureSetId, WorkingEntry(MerkleRootMeta)).init(allocator), - .slot_meta_working_set = AutoHashMap(u64, SlotMetaWorkingSetEntry).init(allocator), - .index_working_set = AutoHashMap(u64, IndexMetaWorkingSetEntry).init(allocator), - .duplicate_shreds = ArrayList(PossibleDuplicateShred).init(allocator), + .just_inserted_shreds = .init(allocator), // TODO capacity = shreds.len + .erasure_metas = .empty, + .merkle_root_metas = .init(allocator), + .slot_meta_working_set = .init(allocator), + .index_working_set = .init(allocator), + .duplicate_shreds = .init(allocator), .metrics = metrics, }; } @@ -116,10 +116,19 @@ pub const PendingInsertShredsState = struct { /// duplicate_shreds is not deinitialized. ownership is transfered to caller pub fn deinit(self: *Self) void { self.just_inserted_shreds.deinit(); - self.erasure_metas.deinit(); + self.erasure_metas.deinit(self.allocator); self.merkle_root_metas.deinit(); - deinitMapRecursive(&self.slot_meta_working_set); - deinitMapRecursive(&self.index_working_set); + + { + var iter = self.slot_meta_working_set.iterator(); + while (iter.next()) |entry| entry.value_ptr.deinit(self.allocator); + self.slot_meta_working_set.deinit(); + } + { + var iter = self.index_working_set.iterator(); + while (iter.next()) |entry| entry.value_ptr.deinit(self.allocator); + self.index_working_set.deinit(); + } self.write_batch.deinit(); } @@ -131,7 +140,7 @@ pub const PendingInsertShredsState = struct { if (try self.db.get(self.allocator, schema.index, slot)) |item| { entry.value_ptr.* = .{ .index = item }; } else { - entry.value_ptr.* = IndexMetaWorkingSetEntry.init(self.allocator, slot); + entry.value_ptr.* = IndexMetaWorkingSetEntry.init(slot); } } if (self.metrics) |m| m.index_meta_time_us.add(timer.read().asMicros()); @@ -161,7 +170,7 @@ pub const PendingInsertShredsState = struct { }; } else { entry.value_ptr.* = .{ - .new_slot_meta = SlotMeta.init(self.allocator, slot, parent_slot), + .new_slot_meta = SlotMeta.init(slot, parent_slot), }; } } @@ -327,7 +336,7 @@ pub const MerkleRootMetaWorkingStore = struct { pub const ErasureMetaWorkingStore = struct { allocator: Allocator, db: *LedgerDB, - working_entries: *SortedMap(ErasureSetId, WorkingEntry(ErasureMeta)), + working_entries: *SortedMap(ErasureSetId, WorkingEntry(ErasureMeta), .{}), const Self = @This(); @@ -343,7 +352,7 @@ pub const ErasureMetaWorkingStore = struct { erasure_set_id: ErasureSetId, code_shred: CodeShred, ) !*const ErasureMeta { - const erasure_meta_entry = try self.working_entries.getOrPut(erasure_set_id); + const erasure_meta_entry = try self.working_entries.getOrPut(self.allocator, erasure_set_id); if (!erasure_meta_entry.found_existing) { if (try self.db.get(self.allocator, schema.erasure_meta, erasure_set_id)) |meta_| { erasure_meta_entry.value_ptr.* = .{ .clean = meta_ }; @@ -361,7 +370,7 @@ pub const ErasureMetaWorkingStore = struct { pub fn load(self: Self, erasure_set_id: ErasureSetId) !void { if (!self.working_entries.contains(erasure_set_id)) { if (try self.db.get(self.allocator, schema.erasure_meta, erasure_set_id)) |meta_| { - try self.working_entries.put(erasure_set_id, .{ .clean = meta_ }); + try self.working_entries.put(self.allocator, erasure_set_id, .{ .clean = meta_ }); } } } @@ -376,15 +385,19 @@ pub const ErasureMetaWorkingStore = struct { // Check the previous entry from the in memory map to see if it is the consecutive // set to `erasure set` - const id_range, const meta_range = self.working_entries.range( - .{ .slot = slot, .erasure_set_index = 0 }, - erasure_set, - ); - if (id_range.len != 0) { - const i = id_range.len - 1; - const last_meta = meta_range[i].asRef(); - if (@as(u32, @intCast(erasure_set_index)) == last_meta.nextErasureSetIndex()) { - return .{ id_range[i], last_meta.* }; + + { + var iter = self.working_entries.iteratorRanged( + .{ .slot = slot, .erasure_set_index = 0 }, + erasure_set, + .end, + ); + + if (iter.prev()) |entry| { + const last_meta = entry.value_ptr.asRef(); + if (@as(u32, @intCast(erasure_set_index)) == last_meta.nextErasureSetIndex()) { + return .{ entry.key_ptr.*, last_meta.* }; + } } } @@ -451,12 +464,12 @@ pub const IndexMetaWorkingSetEntry = struct { // struct was created did_insert_occur: bool = false, - pub fn init(allocator: std.mem.Allocator, slot: Slot) IndexMetaWorkingSetEntry { - return .{ .index = meta.Index.init(allocator, slot) }; + pub fn init(slot: Slot) IndexMetaWorkingSetEntry { + return .{ .index = meta.Index.init(slot) }; } - pub fn deinit(self: *IndexMetaWorkingSetEntry) void { - self.index.deinit(); + pub fn deinit(self: *IndexMetaWorkingSetEntry, allocator: std.mem.Allocator) void { + self.index.deinit(allocator); } }; @@ -474,9 +487,9 @@ pub const SlotMetaWorkingSetEntry = struct { /// this struct was created. did_insert_occur: bool = false, - pub fn deinit(self: *@This()) void { - self.new_slot_meta.deinit(); - if (self.old_slot_meta) |*old| old.deinit(); + pub fn deinit(self: *@This(), allocator: std.mem.Allocator) void { + self.new_slot_meta.deinit(allocator); + if (self.old_slot_meta) |*old| old.deinit(allocator); } }; @@ -565,14 +578,6 @@ pub const ShredWorkingStore = struct { } }; -pub fn deinitMapRecursive(map: anytype) void { - var iter = map.iterator(); - while (iter.next()) |entry| { - entry.value_ptr.deinit(); - } - map.deinit(); -} - /// agave: is_newly_completed_slot pub fn isNewlyCompletedSlot(slot_meta: *const SlotMeta, backup_slot_meta: *const ?SlotMeta) bool { return slot_meta.isFull() and ( // diff --git a/src/ledger/tests.zig b/src/ledger/tests.zig index 21c5e2ce82..85c3066782 100644 --- a/src/ledger/tests.zig +++ b/src/ledger/tests.zig @@ -377,7 +377,7 @@ pub fn insertDataForBlockTest( var result_writer = state.resultWriter(); try result_writer.setRoots(&.{ slot - 1, slot, slot + 1 }); - const parent_meta = SlotMeta.init(allocator, 0, null); + const parent_meta = SlotMeta.init(0, null); try db.put(schema.slot_meta, slot - 1, parent_meta); var expected_transactions = std.ArrayList(VersionedTransactionWithStatusMeta).init(allocator); diff --git a/src/replay/consensus/cluster_sync.zig b/src/replay/consensus/cluster_sync.zig index 70b8b7cbc6..f30027f445 100644 --- a/src/replay/consensus/cluster_sync.zig +++ b/src/replay/consensus/cluster_sync.zig @@ -18,6 +18,8 @@ const GossipVerifiedVoteHash = sig.consensus.vote_listener.GossipVerifiedVoteHas const ThresholdConfirmedSlot = sig.consensus.vote_listener.ThresholdConfirmedSlot; const LatestValidatorVotes = sig.consensus.latest_validator_votes.LatestValidatorVotes; +const SlotSet = collections.SortedSet(Slot, .{}); + const ledger_tests = sig.ledger.tests; pub const ProcessClusterSyncTimings = struct { @@ -155,19 +157,34 @@ pub const SlotData = struct { latest_validator_votes: LatestValidatorVotes, /// Analogous to [DuplicateSlotsTracker](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/repair/cluster_slot_state_verifier.rs#L18) - pub const DuplicateSlots = collections.SortedSetUnmanaged(Slot); + pub const DuplicateSlots = collections.SortedSet( + Slot, + .{}, + ); /// Analogous to [DuplicateSlotsToRepair](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/repair/cluster_slot_state_verifier.rs#L19) pub const DuplicateSlotsToRepair = std.AutoArrayHashMapUnmanaged(Slot, Hash); /// Analogous to [PurgeRepairSlotCounter](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/repair/cluster_slot_state_verifier.rs#L20) - pub const PurgeRepairSlotCounters = collections.SortedMapUnmanaged(Slot, usize); + pub const PurgeRepairSlotCounters = collections.SortedMap( + Slot, + usize, + .{}, + ); /// Analogous to [EpochSlotsFrozenSlots](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/repair/cluster_slot_state_verifier.rs#L22) - pub const EpochSlotsFrozenSlots = collections.SortedMapUnmanaged(Slot, Hash); + pub const EpochSlotsFrozenSlots = collections.SortedMap( + Slot, + Hash, + .{}, + ); /// Analogous to [DuplicateConfirmedSlots](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/repair/cluster_slot_state_verifier.rs#L24) - pub const DuplicateConfirmedSlots = collections.SortedMapUnmanaged(Slot, Hash); + pub const DuplicateConfirmedSlots = collections.SortedMap( + Slot, + Hash, + .{}, + ); pub const empty: SlotData = .{ .duplicate_confirmed_slots = .empty, @@ -179,7 +196,7 @@ pub const SlotData = struct { .latest_validator_votes = .empty, }; - pub fn deinit(self: SlotData, allocator: std.mem.Allocator) void { + pub fn deinit(self: *SlotData, allocator: std.mem.Allocator) void { self.duplicate_confirmed_slots.deinit(allocator); self.epoch_slots_frozen_slots.deinit(allocator); @@ -194,20 +211,21 @@ pub const SlotData = struct { /// Analogous to [UnfrozenGossipVerifiedVoteHashes](https://github.com/anza-xyz/agave/blob/0315eb6adc87229654159448344972cbe484d0c7/core/src/unfrozen_gossip_verified_vote_hashes.rs#L8) pub const UnfrozenGossipVerifiedVoteHashes = struct { - votes_per_slot: sig.utils.collections.SortedMapUnmanaged(Slot, HashToVotesMap), + votes_per_slot: sig.utils.collections.SortedMap( + Slot, + HashToVotesMap, + .{}, + ), const HashToVotesMap = std.AutoArrayHashMapUnmanaged(Hash, VoteList); const VoteList = std.ArrayListUnmanaged(Pubkey); pub const empty: UnfrozenGossipVerifiedVoteHashes = .{ .votes_per_slot = .empty }; - pub fn deinit(self: UnfrozenGossipVerifiedVoteHashes, allocator: std.mem.Allocator) void { - var votes_per_slot = self.votes_per_slot; - for (votes_per_slot.values()) |*htvm| { - for (htvm.values()) |*vl| vl.deinit(allocator); - htvm.deinit(allocator); - } - votes_per_slot.deinit(allocator); + pub fn deinit(self: *UnfrozenGossipVerifiedVoteHashes, allocator: std.mem.Allocator) void { + var iter = self.votes_per_slot.iterator(); + while (iter.next()) |entry| entry.value_ptr.deinit(allocator); + self.votes_per_slot.deinit(allocator); } /// Update `latest_validator_votes_for_frozen_slots` if gossip has seen a newer vote for a frozen slot. @@ -253,7 +271,7 @@ pub const UnfrozenGossipVerifiedVoteHashes = struct { // frozen later const vps_gop = try self.votes_per_slot.getOrPut(allocator, vote_slot); errdefer if (!vps_gop.found_existing) { - std.debug.assert(self.votes_per_slot.orderedRemove(vps_gop.key_ptr.*)); + std.debug.assert(self.votes_per_slot.remove(vps_gop.key_ptr.*)); }; const hash_to_votes: *HashToVotesMap = vps_gop.value_ptr; @@ -1073,7 +1091,7 @@ pub const check_slot_agrees_with_cluster = struct { return; } else { // Otherwise, add it to the set of processed slots, and proceed. - try duplicate_slots_tracker.put(allocator, slot); + try duplicate_slots_tracker.put(allocator, slot, {}); } // TODO: consider putting a prometheus metric here, similar to how agave @@ -1340,13 +1358,13 @@ const state_change = struct { } _ = duplicate_slots_to_repair.swapRemove(slot); - _ = purge_repair_slot_counter.orderedRemove(slot); + _ = purge_repair_slot_counter.remove(slot); } }; const Descendants = std.AutoArrayHashMapUnmanaged( Slot, - sig.utils.collections.SortedMapUnmanaged(Slot, void), + SlotSet, ); fn descendantsDeinit(allocator: std.mem.Allocator, descendants: Descendants) void { for (descendants.values()) |*child_set| child_set.deinit(allocator); @@ -1533,11 +1551,31 @@ const TestData = struct { errdefer descendants.deinit(allocator); errdefer for (descendants.values()) |*child_set| child_set.deinit(allocator); try descendants.ensureUnusedCapacity(allocator, 4); - descendants.putAssumeCapacity(0, try .init(allocator, &.{ 1, 2, 3 }, &.{})); - descendants.putAssumeCapacity(1, try .init(allocator, &.{ 3, 2 }, &.{})); - descendants.putAssumeCapacity(2, try .init(allocator, &.{3}, &.{})); + descendants.putAssumeCapacity(0, blk: { + var set: SlotSet = .empty; + errdefer set.deinit(allocator); + try set.put(allocator, 1, {}); + try set.put(allocator, 2, {}); + try set.put(allocator, 3, {}); + + break :blk set; + }); + descendants.putAssumeCapacity(1, blk: { + var set: SlotSet = .empty; + errdefer set.deinit(allocator); + try set.put(allocator, 3, {}); + try set.put(allocator, 2, {}); + + break :blk set; + }); + descendants.putAssumeCapacity(2, blk: { + var set: SlotSet = .empty; + errdefer set.deinit(allocator); + try set.put(allocator, 3, {}); + + break :blk set; + }); descendants.putAssumeCapacity(3, .empty); - for (descendants.values()) |*slot_set| slot_set.sort(); return .{ .slot_tracker = slot_tracker, @@ -1574,19 +1612,27 @@ test "apply state changes" { .slot = duplicate_slot, .hash = duplicate_slot_hash, }).?); - for ([_][]const Slot{ - descendants.getPtr(duplicate_slot).?.keys(), - &.{duplicate_slot}, - }) |child_slot_set| { - for (child_slot_set) |child_slot| { - try std.testing.expectEqual( - duplicate_slot, - heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ - .slot = child_slot, - .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, - }).?, - ); - } + + var iter = descendants.getPtr(duplicate_slot).?.iterator(); + while (iter.next()) |entry| { + const child_slot = entry.key_ptr.*; + try std.testing.expectEqual( + duplicate_slot, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, + }).?, + ); + } + { + const child_slot = duplicate_slot; + try std.testing.expectEqual( + duplicate_slot, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, + }).?, + ); } var duplicate_slots_to_repair: SlotData.DuplicateSlotsToRepair = .empty; @@ -1755,20 +1801,29 @@ test "apply state changes duplicate confirmed matches frozen" { try confirmed_non_dupe_frozen_hash.finalize(duplicate_slot, ledger.resultWriter()); } - for ([_][]const Slot{ - descendants.getPtr(duplicate_slot).?.keys(), - &.{duplicate_slot}, - }) |child_slot_set| { - for (child_slot_set) |child_slot| { - try std.testing.expectEqual( - null, - heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ - .slot = child_slot, - .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, - }), - ); - } + var iter = descendants.iterator(); + while (iter.next()) |entry| { + const child_slot = entry.key_ptr.*; + try std.testing.expectEqual( + null, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, + }), + ); } + + { + const child_slot = duplicate_slot; + try std.testing.expectEqual( + null, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.slots.get(child_slot).?.state.hash.readCopy().?, + }), + ); + } + try std.testing.expectEqual(true, heaviest_subtree_fork_choice.isCandidate(&.{ .slot = duplicate_slot, .hash = our_duplicate_slot_hash, @@ -1854,19 +1909,29 @@ test "apply state changes slot frozen and duplicate confirmed matches frozen" { try confirmed_non_dupe_frozen_hash.finalize(duplicate_slot, ledger.resultWriter()); } - for ([_][]const Slot{ - descendants.getPtr(duplicate_slot).?.keys(), - &.{duplicate_slot}, - }) |child_slot_set| { - for (child_slot_set) |child_slot| { - try std.testing.expectEqual( - null, - heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ - .slot = child_slot, - .hash = slot_tracker.get(child_slot).?.state.hash.readCopy().?, - }), - ); - } + var iter = descendants.iterator(); + while (iter.next()) |entry| { + const child_slot = entry.key_ptr.*; + + try std.testing.expectEqual( + null, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.get(child_slot).?.state.hash.readCopy().?, + }), + ); + } + + { + const child_slot = duplicate_slot; + + try std.testing.expectEqual( + null, + heaviest_subtree_fork_choice.latestInvalidAncestor(&.{ + .slot = child_slot, + .hash = slot_tracker.get(child_slot).?.state.hash.readCopy().?, + }), + ); } try std.testing.expectEqual(true, heaviest_subtree_fork_choice.isCandidate(&.{ diff --git a/src/replay/consensus/core.zig b/src/replay/consensus/core.zig index 7d90513e16..c2af72aba3 100644 --- a/src/replay/consensus/core.zig +++ b/src/replay/consensus/core.zig @@ -14,7 +14,7 @@ pub const Logger = sig.trace.Logger("consensus"); const Channel = sig.sync.Channel; const RwMux = sig.sync.RwMux; -const SortedSetUnmanaged = sig.utils.collections.SortedSetUnmanaged; +const SortedSet = sig.utils.collections.SortedSet; const Ancestors = sig.core.Ancestors; const Epoch = sig.core.Epoch; @@ -32,6 +32,8 @@ const SocketAddr = sig.net.SocketAddr; const Tower = sig.consensus.tower.Tower; const VoteStateVersions = sig.runtime.program.vote.state.VoteStateVersions; +const SlotSet = SortedSet(Slot, .{}); + /// UDP sockets used to send vote transactions. pub const VoteSockets = struct { ipv4: sig.net.UdpSocket, @@ -119,7 +121,7 @@ pub const TowerConsensus = struct { /// functions; ie, it isn't used for any persistent data arena_state: std.heap.ArenaAllocator.State, - pub fn deinit(self: TowerConsensus, allocator: Allocator) void { + pub fn deinit(self: *TowerConsensus, allocator: Allocator) void { self.replay_tower.deinit(allocator); self.fork_choice.deinit(allocator); @@ -423,8 +425,6 @@ pub const TowerConsensus = struct { params.duplicate_confirmed_slots.clearRetainingCapacity(); params.gossip_verified_vote_hashes.clearRetainingCapacity(); - const SlotSet = SortedSetUnmanaged(Slot); - const asc_desc_zone = tracy.Zone.init( @src(), .{ .name = "TowerConsensus.process: ancestors/descendants" }, @@ -447,7 +447,7 @@ pub const TowerConsensus = struct { try ancestor_gop.value_ptr.addSlot(arena, ancestor_slot); const descendants_gop = try descendants.getOrPutValue(arena, ancestor_slot, .empty); - try descendants_gop.value_ptr.put(arena, slot); + try descendants_gop.value_ptr.put(arena, slot, {}); } } break :cluster_sync_and_ancestors_descendants .{ ancestors, descendants }; @@ -509,7 +509,7 @@ pub const TowerConsensus = struct { ledger: *Ledger, gossip_table: ?*sig.sync.RwMux(sig.gossip.GossipTable), ancestors: *const std.AutoArrayHashMapUnmanaged(Slot, Ancestors), - descendants: *const std.AutoArrayHashMapUnmanaged(Slot, SortedSetUnmanaged(Slot)), + descendants: *const std.AutoArrayHashMapUnmanaged(Slot, SlotSet), slot_tracker: *SlotTracker, epoch_tracker: *const EpochTracker, progress_map: *ProgressMap, @@ -7067,8 +7067,10 @@ test "successful fork switch (switch_proof)" { } var ancestors_map = std.AutoArrayHashMapUnmanaged(Slot, sig.core.Ancestors).empty; - var descendants_map = - std.AutoArrayHashMapUnmanaged(Slot, sig.utils.collections.SortedSetUnmanaged(Slot)).empty; + var descendants_map: std.AutoArrayHashMapUnmanaged( + Slot, + sig.utils.collections.SortedSet(Slot, .{}), + ) = .empty; defer { var it = ancestors_map.iterator(); while (it.next()) |entry| entry.value_ptr.deinit(allocator); @@ -7089,7 +7091,7 @@ test "successful fork switch (switch_proof)" { for (slot_ancestors.keys()) |a| { try gop.value_ptr.addSlot(allocator, a); const dg = try descendants_map.getOrPutValue(allocator, a, .empty); - try dg.value_ptr.put(allocator, slot); + try dg.value_ptr.put(allocator, slot, {}); } } @@ -7244,7 +7246,7 @@ test "successful fork switch (switch_proof)" { var ancestors_map2 = std.AutoArrayHashMapUnmanaged(Slot, sig.core.Ancestors).empty; var descendants_map2 = - std.AutoArrayHashMapUnmanaged(Slot, sig.utils.collections.SortedSetUnmanaged(Slot)).empty; + std.AutoArrayHashMapUnmanaged(Slot, SlotSet).empty; defer { var it = ancestors_map2.iterator(); while (it.next()) |entry| entry.value_ptr.deinit(allocator); @@ -7265,7 +7267,7 @@ test "successful fork switch (switch_proof)" { for (slot_ancestors.keys()) |a| { try gop.value_ptr.addSlot(allocator, a); const dg = try descendants_map2.getOrPutValue(allocator, a, .empty); - try dg.value_ptr.put(allocator, slot); + try dg.value_ptr.put(allocator, slot, {}); } } diff --git a/src/replay/service.zig b/src/replay/service.zig index 2c8add27e6..f96cfdbe6d 100644 --- a/src/replay/service.zig +++ b/src/replay/service.zig @@ -686,9 +686,9 @@ test trackNewSlots { .{ 6, 4, &.{} }, }) |item| { const slot, const parent, const children = item; - var meta = sig.ledger.meta.SlotMeta.init(allocator, slot, parent); - defer meta.deinit(); - try meta.child_slots.appendSlice(children); + var meta: sig.ledger.meta.SlotMeta = .init(slot, parent); + defer meta.deinit(allocator); + try meta.child_slots.appendSlice(allocator, children); try ledger.db.put(sig.ledger.schema.schema.slot_meta, slot, meta); } diff --git a/src/runtime/program/vote/state.zig b/src/runtime/program/vote/state.zig index ffee70bcda..1afc11b709 100644 --- a/src/runtime/program/vote/state.zig +++ b/src/runtime/program/vote/state.zig @@ -12,7 +12,7 @@ const Slot = sig.core.Slot; const Epoch = sig.core.Epoch; const Pubkey = sig.core.Pubkey; const Hash = sig.core.hash.Hash; -const SortedMap = sig.utils.collections.SortedMapUnmanaged; +const SortedMap = sig.utils.collections.SortedMap; const AccountSharedData = sig.runtime.AccountSharedData; const SlotHashes = sig.runtime.sysvar.SlotHashes; @@ -351,7 +351,9 @@ pub fn deserializeTowerSync( /// [agave] https://github.com/anza-xyz/solana-sdk/blob/52d80637e13bca19ed65920fbda154993c37dbbe/vote-interface/src/authorized_voters.rs#L11 pub const AuthorizedVoters = struct { - voters: SortedMap(Epoch, Pubkey), + const Map = SortedMap(Epoch, Pubkey, .{}); + + voters: Map, pub const EMPTY: AuthorizedVoters = .{ .voters = .empty }; pub const @"!bincode-config": sig.bincode.FieldConfig(AuthorizedVoters) = .{ @@ -360,7 +362,7 @@ pub const AuthorizedVoters = struct { }; pub fn init(allocator: Allocator, epoch: Epoch, pubkey: Pubkey) !AuthorizedVoters { - var authorized_voters: SortedMap(Epoch, Pubkey) = .empty; + var authorized_voters: Map = .empty; try authorized_voters.put(allocator, epoch, pubkey); return .{ .voters = authorized_voters }; } @@ -427,9 +429,8 @@ pub const AuthorizedVoters = struct { } for (expired_keys.items) |key| { - _ = self.voters.swapRemoveNoSort(key); + _ = self.voters.remove(key); } - self.voters.sort(); // Have to uphold this invariant b/c this is // 1) The check for whether the vote state is initialized @@ -444,21 +445,13 @@ pub const AuthorizedVoters = struct { } pub fn first(self: *AuthorizedVoters) ?struct { Epoch, Pubkey } { - var voter_iter = self.voters.iterator(); - if (voter_iter.next()) |entry| { - return .{ entry.key_ptr.*, entry.value_ptr.* }; - } else { - return null; - } + const first_voter = self.voters.minEntry() orelse return null; + return .{ first_voter.key_ptr.*, first_voter.value_ptr.* }; } pub fn last(self: *const AuthorizedVoters) ?struct { Epoch, Pubkey } { - const last_epoch = self.voters.max orelse return null; - if (self.voters.get(last_epoch)) |last_pubkey| { - return .{ last_epoch, last_pubkey }; - } else { - return null; - } + const last_voter = self.voters.maxEntry() orelse return null; + return .{ last_voter.key_ptr.*, last_voter.value_ptr.* }; } pub fn len(self: *const AuthorizedVoters) usize { @@ -481,12 +474,9 @@ pub const AuthorizedVoters = struct { if (self.voters.get(epoch)) |pubkey| { return .{ pubkey, true }; } else { - _, const values = self.voters.range(0, epoch); - if (values.len == 0) { - return null; - } - const last_voter = values[values.len - 1]; - return .{ last_voter, false }; + var it = self.voters.iteratorRanged(null, epoch, .end); + const last_voter = it.prev() orelse return null; + return .{ last_voter.value_ptr.*, false }; } } @@ -514,10 +504,11 @@ pub const AuthorizedVoters = struct { pub fn serialize(writer: anytype, data: anytype, _: sig.bincode.Params) !void { var authorized_voters: AuthorizedVoters = data; try writer.writeInt(usize, authorized_voters.len(), .little); - const items = authorized_voters.voters.items(); - for (items[0], items[1]) |k, v| { - try writer.writeInt(u64, k, .little); - try writer.writeAll(&v.data); + + var iter = authorized_voters.voters.iterator(); + while (iter.next()) |entry| { + try writer.writeInt(u64, entry.key_ptr.*, .little); + try writer.writeAll(&entry.value_ptr.data); } } @@ -525,7 +516,10 @@ pub const AuthorizedVoters = struct { if (self.count() != other.count()) return false; var self_voters = self.voters; var other_voters = other.voters; - for (self_voters.keys()) |key| { + + var it = self_voters.iterator(); + while (it.next()) |entry| { + const key = entry.key_ptr.*; const self_value = self_voters.get(key).?; const other_value = other_voters.get(key) orelse return false; if (!self_value.equals(&other_value)) return false; @@ -640,7 +634,7 @@ pub const VoteStateVersions = union(enum) { pub fn convertToCurrent(self: VoteStateVersions, allocator: Allocator) !VoteState { switch (self) { .v0_23_5 => |state| { - const authorized_voters = try AuthorizedVoters.init( + var authorized_voters = try AuthorizedVoters.init( allocator, state.voter_epoch, state.voter, @@ -669,7 +663,7 @@ pub const VoteStateVersions = union(enum) { }; }, .v1_14_11 => |state| { - const authorized_voters = try state.voters.clone(allocator); + var authorized_voters = try state.voters.clone(allocator); errdefer authorized_voters.deinit(allocator); const votes = try VoteStateVersions.landedVotesFromLockouts( @@ -937,7 +931,7 @@ pub const VoteState = struct { var votes = try self.votes.clone(allocator); errdefer votes.deinit(allocator); - const voters = try self.voters.clone(allocator); + var voters = try self.voters.clone(allocator); errdefer voters.deinit(allocator); return .{ diff --git a/src/time/stake_weighted_timestamp.zig b/src/time/stake_weighted_timestamp.zig index de075a433d..6d999f645f 100644 --- a/src/time/stake_weighted_timestamp.zig +++ b/src/time/stake_weighted_timestamp.zig @@ -7,7 +7,7 @@ const Pubkey = sig.core.Pubkey; const Slot = sig.core.Slot; const StakeAndVoteAccountsMap = sig.core.stakes.StakeAndVoteAccountsMap; -const SortedMap = sig.utils.collections.SortedMapUnmanaged; +const SortedMap = sig.utils.collections.SortedMap; /// Analogous to [MaxAllowableDrift](https://github.com/anza-xyz/agave/blob/e0bd9224fe60d8caa35bcca8daf6c8103ce424ec/runtime/src/stake_weighted_timestamp.rs#L21) pub const MaxAllowableDrift = struct { @@ -38,18 +38,26 @@ pub fn calculateStakeWeightedTimestamp( max_allowable_drift: MaxAllowableDrift, fix_estimate_into_u64: bool, ) Allocator.Error!?i64 { - var stakes_per_timestamp = SortedMap(i64, u128).empty; + var stakes_per_timestamp: SortedMap( + i64, + u128, + .{}, + ) = .empty; defer stakes_per_timestamp.deinit(allocator); var total_stake: u128 = 0; for (recent_timestamps) |timestamp_entry| { const vote_pubkey, const timestamp_slot, const timestamp = timestamp_entry; + const offset_s = (ns_per_slot *| (slot -| timestamp_slot)) / 1_000_000_000; const estimate_s = timestamp +| @as(i64, @intCast(offset_s)); const stake = if (vote_accounts.getPtr(vote_pubkey)) |vote_account_entry| vote_account_entry.stake else 0; + + if (estimate_s == std.math.maxInt(i64)) unreachable; + const entry = try stakes_per_timestamp.getOrPut(allocator, estimate_s); if (entry.found_existing) entry.value_ptr.* +|= stake @@ -63,11 +71,12 @@ pub fn calculateStakeWeightedTimestamp( var stake_accumulator: u128 = 0; var estimate_s: i64 = 0; - for (stakes_per_timestamp.keys()) |timestamp_s| { - const stake = stakes_per_timestamp.get(timestamp_s).?; + var iter = stakes_per_timestamp.iterator(); + while (iter.next()) |entry| { + const stake = entry.value_ptr.*; stake_accumulator +|= stake; if (stake_accumulator > total_stake / 2) { - estimate_s = timestamp_s; + estimate_s = entry.key_ptr.*; break; } } @@ -193,7 +202,7 @@ test "uses median: low-staked outliers" { { const recent_timestamps = [_]struct { Pubkey, Slot, i64 }{ .{ pubkey_0, slot, recent_timestamp }, - .{ pubkey_1, slot, std.math.maxInt(i64) }, + .{ pubkey_1, slot, std.math.maxInt(i64) - 1 }, .{ pubkey_2, slot, recent_timestamp }, .{ pubkey_3, slot, recent_timestamp }, .{ pubkey_4, slot, recent_timestamp }, @@ -216,7 +225,7 @@ test "uses median: low-staked outliers" { { const recent_timestamps = [_]struct { Pubkey, Slot, i64 }{ .{ pubkey_0, slot, 0 }, - .{ pubkey_1, slot, std.math.maxInt(i64) }, + .{ pubkey_1, slot, std.math.maxInt(i64) - 1 }, .{ pubkey_2, slot, recent_timestamp }, .{ pubkey_3, slot, recent_timestamp }, .{ pubkey_4, slot, recent_timestamp }, @@ -273,7 +282,7 @@ test "uses median: high-staked outliers" { const recent_timestamps = [_]struct { Pubkey, Slot, i64 }{ .{ pubkey_0, slot, 0 }, - .{ pubkey_1, slot, std.math.maxInt(i64) }, + .{ pubkey_1, slot, std.math.maxInt(i64) - 1 }, .{ pubkey_2, slot, recent_timestamp }, }; diff --git a/src/utils/collections.zig b/src/utils/collections.zig index c21bcefa5c..4b7db054b1 100644 --- a/src/utils/collections.zig +++ b/src/utils/collections.zig @@ -1,5 +1,6 @@ const std = @import("std"); const sig = @import("../sig.zig"); +const sorted_map = @import("sorted_map.zig"); const Allocator = std.mem.Allocator; const ArrayList = std.ArrayList; @@ -290,629 +291,10 @@ test "SplitUnionList: addOne, get, and swapRemove" { _ = list.swapRemove(.{ .tag = .one, .index = 0 }); } -/// DEPRECATED: use the unmanaged variant instead -pub fn SortedSet(comptime T: type) type { - return SortedSetCustom(T, .{}); -} - -/// DEPRECATED: use the unmanaged variant instead -pub fn SortedSetCustom(comptime T: type, comptime config: SortedMapConfig(T)) type { - return struct { - map: SortedMapCustom(T, void, config), - const SortedSetSelf = @This(); - - pub fn init(allocator: Allocator) SortedSetSelf { - return .{ .map = .init(allocator) }; - } - - pub fn deinit(self: SortedSetSelf) void { - self.map.deinit(); - } - - pub fn clone(self: SortedSetSelf) !SortedSetSelf { - return .{ .map = try self.map.clone() }; - } - - pub fn eql(self: *SortedSetSelf, other: *SortedSetSelf) bool { - return self.map.eql(&other.map); - } - - pub fn put(self: *SortedSetSelf, item: T) !void { - try self.map.put(item, {}); - } - - pub fn orderedRemove(self: *SortedSetSelf, item: T) bool { - return self.map.orderedRemove(item); - } - - pub fn contains(self: SortedSetSelf, item: T) bool { - return self.map.contains(item); - } - - pub fn count(self: SortedSetSelf) usize { - return self.map.count(); - } - - pub fn items(self: *SortedSetSelf) []const T { - return self.map.keys(); - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn range(self: *SortedSetSelf, start: ?T, end: ?T) []const T { - return self.map.range(start, end)[0]; - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn rangeCustom(self: *SortedSetSelf, start: ?Bound(T), end: ?Bound(T)) []const T { - return self.map.rangeCustom(start, end)[0]; - } - }; -} - -/// A set that guarantees the contained items will be sorted whenever -/// accessed through public methods like `items` and `range`. -/// -/// Compatible with numbers, slices of numbers, and types that have an "order" method -pub fn SortedSetUnmanaged(comptime T: type) type { - return SortedSetUnmanagedCustom(T, .{}); -} - -/// A set that guarantees the contained items will be sorted whenever -/// accessed through public methods like `items` and `range`. -pub fn SortedSetUnmanagedCustom(comptime T: type, comptime config: SortedMapConfig(T)) type { - return struct { - map: SortedMapUnmanagedCustom(T, void, config), - const SortedSetSelf = @This(); - - pub const empty: SortedSetSelf = .{ - .map = .empty, - }; - - pub fn deinit(self: SortedSetSelf, allocator: std.mem.Allocator) void { - self.map.deinit(allocator); - } - - pub fn clearRetainingCapacity(self: *SortedSetSelf) void { - return self.map.inner.clearRetainingCapacity(); - } - - pub fn clone( - self: SortedSetSelf, - allocator: std.mem.Allocator, - ) std.mem.Allocator.Error!SortedSetSelf { - return .{ .map = try self.map.clone(allocator) }; - } - - pub fn eql(self: *SortedSetSelf, other: *SortedSetSelf) bool { - return self.map.eql(&other.map); - } - - pub fn put( - self: *SortedSetSelf, - allocator: std.mem.Allocator, - item: T, - ) std.mem.Allocator.Error!void { - try self.map.put(allocator, item, {}); - } +pub const SortedSet = sorted_map.SortedSet; +pub const SortedMap = sorted_map.SortedMap; - pub fn orderedRemove(self: *SortedSetSelf, item: T) bool { - return self.map.orderedRemove(item); - } - - pub fn pop(self: *SortedSetSelf) ?T { - self.map.sort(); - const kv = self.map.inner.pop() orelse return null; - return kv.key; - } - - pub fn contains(self: SortedSetSelf, item: T) bool { - return self.map.contains(item); - } - - pub fn count(self: SortedSetSelf) usize { - return self.map.count(); - } - - pub fn items(self: *SortedSetSelf) []const T { - return self.map.keys(); - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn range(self: *SortedSetSelf, start: ?T, end: ?T) []const T { - return self.map.range(start, end)[0]; - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn rangeCustom(self: *SortedSetSelf, start: ?Bound(T), end: ?Bound(T)) []const T { - return self.map.rangeCustom(start, end)[0]; - } - }; -} - -/// DEPRECATED: use the unmanaged variant instead -pub fn SortedMap(comptime K: type, comptime V: type) type { - return SortedMapCustom(K, V, .{}); -} - -/// DEPRECATED: use the unmanaged variant instead -pub fn SortedMapCustom( - comptime K: type, - comptime V: type, - comptime config: SortedMapConfig(K), -) type { - return struct { - allocator: std.mem.Allocator, - unmanaged: Unmanaged, - const SortedMapSelf = @This(); - - const Unmanaged = SortedMapUnmanagedCustom(K, V, config); - - pub const @"!bincode-config": sig.bincode.FieldConfig(SortedMapSelf) = .{ - .deserializer = bincodeDeserialize, - .serializer = bincodeSerialize, - .free = bincodeFree, - }; - - pub fn init(allocator: Allocator) SortedMapSelf { - return .{ - .allocator = allocator, - .unmanaged = .empty, - }; - } - - pub fn deinit(self: SortedMapSelf) void { - var self_mut = self; - self_mut.unmanaged.deinit(self.allocator); - } - - pub fn clone(self: SortedMapSelf) std.mem.Allocator.Error!SortedMapSelf { - return .{ - .allocator = self.allocator, - .unmanaged = try self.unmanaged.clone(self.allocator), - }; - } - - pub fn eql(self: *SortedMapSelf, other: *SortedMapSelf) bool { - return self.unmanaged.eql(&other.unmanaged); - } - - pub fn get(self: SortedMapSelf, key: K) ?V { - return self.unmanaged.get(key); - } - - pub fn getEntry(self: SortedMapSelf, key: K) ?Unmanaged.Entry { - return self.unmanaged.getEntry(key); - } - - pub fn fetchSwapRemove(self: *SortedMapSelf, key: K) ?Unmanaged.Inner.KV { - return self.unmanaged.fetchSwapRemove(key); - } - - pub fn swapRemoveNoSort(self: *SortedMapSelf, key: K) bool { - return self.unmanaged.swapRemoveNoSort(key); - } - - pub fn getOrPut( - self: *SortedMapSelf, - key: K, - ) std.mem.Allocator.Error!Unmanaged.Inner.GetOrPutResult { - return self.unmanaged.getOrPut(self.allocator, key); - } - - pub fn put(self: *SortedMapSelf, key: K, value: V) std.mem.Allocator.Error!void { - try self.unmanaged.put(self.allocator, key, value); - } - - pub fn orderedRemove(self: *SortedMapSelf, key: K) bool { - return self.unmanaged.orderedRemove(key); - } - - pub fn contains(self: SortedMapSelf, key: K) bool { - return self.unmanaged.contains(key); - } - - pub fn count(self: SortedMapSelf) usize { - return self.unmanaged.count(); - } - - pub fn keys(self: *SortedMapSelf) []const K { - return self.unmanaged.keys(); - } - - pub fn mutableKeys(self: *SortedMapSelf) []K { - return self.unmanaged.mutableKeys(); - } - - pub fn items(self: *SortedMapSelf) struct { []const K, []const V } { - return self.unmanaged.items(); - } - - pub fn iterator(self: *SortedMapSelf) Unmanaged.Inner.Iterator { - return self.unmanaged.iterator(); - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn range(self: *SortedMapSelf, start: ?K, end: ?K) struct { []const K, []const V } { - return self.unmanaged.range(start, end); - } - - /// subslice of items ranging from start to end - pub fn rangeCustom( - self: *SortedMapSelf, - start_bound: ?Bound(K), - end_bound: ?Bound(K), - ) struct { []const K, []const V } { - return self.unmanaged.rangeCustom(start_bound, end_bound); - } - - pub fn sort(self: *SortedMapSelf) void { - self.unmanaged.sort(); - } - - fn bincodeDeserialize( - limit_allocator: *sig.bincode.LimitAllocator, - reader: anytype, - params: sig.bincode.Params, - ) !SortedMapSelf { - const unmanaged = - try sig.bincode.readWithLimit(limit_allocator, Unmanaged, reader, params); - return .{ - .allocator = limit_allocator.backing_allocator, // patch persistent. - .unmanaged = unmanaged, - }; - } - - fn bincodeSerialize( - writer: anytype, - data: anytype, - params: sig.bincode.Params, - ) !void { - try sig.bincode.write(writer, data.unmanaged, params); - } - - fn bincodeFree(_: std.mem.Allocator, data: anytype) void { - data.deinit(); - } - }; -} - -/// A map that guarantees the contained items will be sorted by key -/// whenever accessed through public methods like `keys` and `range`. -/// -/// Compatible with numbers, slices of numbers, and types that have an "order" method -pub fn SortedMapUnmanaged(comptime K: type, comptime V: type) type { - return SortedMapUnmanagedCustom(K, V, .{}); -} - -/// A map that guarantees the contained items will be sorted by key -/// whenever accessed through public methods like `keys` and `range`. -/// -/// TODO consider reimplementing with something faster (e.g. binary tree) -pub fn SortedMapUnmanagedCustom( - comptime K: type, - comptime V: type, - comptime config: SortedMapConfig(K), -) type { - const order = config.orderFn; - - return struct { - inner: Inner, - max: ?K, - is_sorted: bool, - - const SortedMapSelf = @This(); - - const Inner = std.ArrayHashMapUnmanaged(K, V, config.Context, config.store_hash); - - pub const Entry = Inner.Entry; - - pub const empty: SortedMapSelf = .{ - .inner = .empty, - .max = null, - .is_sorted = true, - }; - - pub fn deinit(self: SortedMapSelf, allocator: std.mem.Allocator) void { - var self_mut = self; - self_mut.inner.deinit(allocator); - } - - pub fn init( - allocator: std.mem.Allocator, - keys_init: []const K, - values_init: []const V, - ) std.mem.Allocator.Error!SortedMapSelf { - var result: SortedMapSelf = .empty; - errdefer result.deinit(allocator); - try result.inner.reinit(allocator, keys_init, values_init); - result.sort(); - return result; - } - - pub fn clone( - self: SortedMapSelf, - allocator: std.mem.Allocator, - ) std.mem.Allocator.Error!SortedMapSelf { - return .{ - .inner = try self.inner.clone(allocator), - .max = self.max, - .is_sorted = self.is_sorted, - }; - } - - pub fn eql(self: *SortedMapSelf, other: *SortedMapSelf) bool { - if (self.count() != other.count()) return false; - self.sort(); - other.sort(); - for ( - self.inner.keys(), - self.inner.values(), - other.inner.keys(), - other.inner.values(), - ) |sk, sv, ok, ov| { - if (sk != ok or sv != ov) return false; - } - return true; - } - - pub fn get(self: SortedMapSelf, key: K) ?V { - return self.inner.get(key); - } - - pub fn getPtr(self: SortedMapSelf, key: K) ?*V { - return self.inner.getPtr(key); - } - - pub fn getEntry(self: SortedMapSelf, key: K) ?Inner.Entry { - return self.inner.getEntry(key); - } - - pub fn fetchSwapRemove(self: *SortedMapSelf, key: K) ?Inner.KV { - const item = self.inner.fetchSwapRemove(key); - if (item != null and !self.resetMaxOnRemove(key)) { - self.is_sorted = false; - } - return item; - } - - pub fn swapRemoveNoSort(self: *SortedMapSelf, key: K) bool { - const was_removed = self.inner.swapRemove(key); - if (was_removed and !self.resetMaxOnRemove(key)) { - self.is_sorted = false; - } - return was_removed; - } - - pub fn getOrPut( - self: *SortedMapSelf, - allocator: std.mem.Allocator, - key: K, - ) std.mem.Allocator.Error!Inner.GetOrPutResult { - const result = try self.inner.getOrPut(allocator, key); - if (self.max == null or order(key, self.max.?) == .gt) { - self.max = key; - } else { - self.is_sorted = false; - } - return result; - } - - pub fn put( - self: *SortedMapSelf, - allocator: std.mem.Allocator, - key: K, - value: V, - ) std.mem.Allocator.Error!void { - try self.ensureUnusedCapacity(allocator, 1); - self.putAssumeCapacity(key, value); - } - - pub fn putAssumeCapacity(self: *SortedMapSelf, key: K, value: V) void { - self.inner.putAssumeCapacity(key, value); - if (self.max == null or order(key, self.max.?) == .gt) { - self.max = key; - } else { - self.is_sorted = false; - } - } - - /// Inserts a new `Entry` into the hash map, returning the previous one, if any. - pub fn fetchPut( - self: *SortedMapSelf, - allocator: Allocator, - key: K, - value: V, - ) std.mem.Allocator.Error!?Inner.KV { - const gop = try self.getOrPut(allocator, key); - const result: ?Inner.KV = if (!gop.found_existing) null else .{ - .key = gop.key_ptr.*, - .value = gop.value_ptr.*, - }; - gop.key_ptr.* = key; - gop.value_ptr.* = value; - return result; - } - - pub fn orderedRemove(self: *SortedMapSelf, key: K) bool { - const was_removed = self.inner.orderedRemove(key); - if (was_removed) _ = self.resetMaxOnRemove(key); - return was_removed; - } - - /// - returns whether the key was the prior max. - /// - don't call this unless an item was definitely removed. - fn resetMaxOnRemove(self: *SortedMapSelf, removed_key: K) bool { - std.debug.assert(self.max != null); - if (self.count() == 0) { - self.max = null; - return true; - } else switch (order(removed_key, self.max.?)) { - .eq => { - const sorted_keys = self.keys(); - self.max = sorted_keys[sorted_keys.len - 1]; - return true; - }, - .gt => unreachable, - .lt => return false, - } - } - - pub fn contains(self: SortedMapSelf, key: K) bool { - return self.inner.contains(key); - } - - pub fn count(self: SortedMapSelf) usize { - return self.inner.count(); - } - - pub fn keys(self: *SortedMapSelf) []const K { - self.sort(); - return self.inner.keys(); - } - - pub fn mutableKeys(self: *SortedMapSelf) []K { - self.sort(); - return self.inner.keys(); - } - - pub fn values(self: *SortedMapSelf) []V { - self.sort(); - return self.inner.values(); - } - - pub fn items(self: *SortedMapSelf) struct { []const K, []const V } { - self.sort(); - return .{ self.inner.keys(), self.inner.values() }; - } - - pub fn iterator(self: *SortedMapSelf) Inner.Iterator { - self.sort(); - return self.inner.iterator(); - } - - /// subslice of items ranging from start (inclusive) to end (exclusive) - pub fn range(self: *SortedMapSelf, start: ?K, end: ?K) struct { []const K, []const V } { - return self.rangeCustom( - if (start) |b| .{ .inclusive = b } else null, - if (end) |b| .{ .exclusive = b } else null, - ); - } - - /// subslice of items ranging from start to end - pub fn rangeCustom( - self: *SortedMapSelf, - start_bound: ?Bound(K), - end_bound: ?Bound(K), - ) struct { []const K, []const V } { - // TODO: can the code in this fn be simplified while retaining identical logic? - const len = self.count(); - if (len == 0) return .{ &.{}, &.{} }; - - // extract relevant info from bounds - const start, const incl_start = if (start_bound) |b| - .{ b.val(), b == .inclusive } - else - .{ null, false }; - const end, const excl_end = if (end_bound) |b| - .{ b.val(), b == .exclusive } - else - .{ null, false }; - - // edge case: check if bounds could permit any items - if (start) |s| if (end) |e| { - if (incl_start and !excl_end) { - if (order(e, s) == .lt) return .{ &.{}, &.{} }; - } else if (order(e, s) != .gt) return .{ &.{}, &.{} }; - }; - - self.sort(); - var keys_ = self.inner.keys(); - var values_ = self.inner.values(); - if (start) |start_| { - // .any instead of .first because uniqueness is guaranteed - const start_index = switch (binarySearch(K, keys_, start_, .any, order)) { - .found => |index| if (incl_start) index else @min(len - 1, index + 1), - .after => |index| index + 1, - .less => 0, - .greater => return .{ &.{}, &.{} }, - .empty => unreachable, // count checked above - }; - keys_ = keys_[start_index..]; - values_ = values_[start_index..]; - } - if (end) |end_| { - // .any instead of .last because uniqueness is guaranteed - const end_index = switch (binarySearch(K, keys_, end_, .any, order)) { - .found => |index| if (excl_end) index else index + 1, - .after => |index| index + 1, - .less => return .{ &.{}, &.{} }, - .greater => keys_.len, - .empty => unreachable, // count checked above - }; - keys_ = keys_[0..end_index]; - values_ = values_[0..end_index]; - } - return .{ keys_, values_ }; - } - - pub fn sort(self: *SortedMapSelf) void { - if (self.is_sorted) return; - self.inner.sort(struct { - items: std.MultiArrayList(Inner.Data).Slice, - pub fn lessThan(ctx: @This(), a_index: usize, b_index: usize) bool { - return order(ctx.items.get(a_index).key, ctx.items.get(b_index).key) == .lt; - } - }{ .items = self.inner.entries.slice() }); - self.is_sorted = true; - } - - pub fn capacity(self: *const SortedMapSelf) usize { - return self.inner.capacity(); - } - - pub fn unusedCapacity(self: *const SortedMapSelf) usize { - return self.inner.capacity() - self.count(); - } - - pub fn ensureUnusedCapacity( - self: *SortedMapSelf, - gpa: std.mem.Allocator, - additional_capacity: usize, - ) Allocator.Error!void { - try self.inner.ensureUnusedCapacity(gpa, additional_capacity); - } - }; -} - -pub fn Bound(comptime T: type) type { - return union(enum) { - inclusive: T, - exclusive: T, - - pub fn val(self: @This()) T { - return switch (self) { - inline .inclusive, .exclusive => |x| x, - }; - } - }; -} - -pub fn SortedMapConfig(comptime K: type) type { - const default_Context, const default_store_hash = if (K == []const u8 or K == []u8) - .{ std.array_hash_map.StringContext, true } - else - .{ std.array_hash_map.AutoContext(K), !std.array_hash_map.autoEqlIsCheap(K) }; - - return struct { - orderFn: fn (a: K, b: K) std.math.Order = defaultOrderFn(K), - /// passthrough to std.ArrayHashMap - Context: type = default_Context, - /// passthrough to std.ArrayHashMap - store_hash: bool = default_store_hash, - }; -} - -fn defaultOrderFn(comptime K: type) fn (lhs: K, rhs: K) std.math.Order { +pub fn defaultOrderFn(comptime K: type) fn (lhs: K, rhs: K) std.math.Order { return struct { fn orderFn(lhs: K, rhs: K) std.math.Order { switch (@typeInfo(K)) { @@ -1166,91 +548,122 @@ pub fn deinitMapAndValues(allocator: Allocator, const_map: anytype) void { const expect = std.testing.expect; const expectEqual = std.testing.expectEqual; -const expectEqualSlices = std.testing.expectEqualSlices; test SortedSet { - var set = SortedSet(u64).init(std.testing.allocator); - defer set.deinit(); + const allocator = std.testing.allocator; + + var set: SortedSet(u64, .{}) = .empty; + defer set.deinit(allocator); // add/contains try expect(!set.contains(3)); - try set.put(3); + try set.put(allocator, 3, {}); try expect(set.contains(3)); - try set.put(0); - try set.put(2); - try set.put(1); - try set.put(4); - try set.put(5); + try set.put(allocator, 0, {}); + try set.put(allocator, 2, {}); + try set.put(allocator, 1, {}); + try set.put(allocator, 4, {}); + try set.put(allocator, 5, {}); // remove - try expect(set.orderedRemove(5)); + try expect(set.remove(5)); try expect(!set.contains(5)); - try expect(!set.orderedRemove(5)); - try set.put(5); + try expect(!set.remove(5)); + try set.put(allocator, 5, {}); try expect(set.contains(5)); - // ordering - for (set.items(), 0..) |key, i| { - try expect(key == i); + var iter = set.iterator(); + var i: u64 = 0; + while (iter.next()) |entry| : (i += 1) { + try expectEqual(i, entry.key_ptr.*); + } + + var j: u64 = i; + while (iter.prev()) |entry| { + j -= 1; + try expectEqual(j, entry.key_ptr.*); } } -test "SortedSet bincode round trip does not break sorting" { - var set = SortedSet(u8).init(std.testing.allocator); - defer set.deinit(); +test "SortedMap bincode round trip does not break sorting" { + const allocator = std.testing.allocator; - try set.put(5); - try set.put(3); + const Set = SortedSet(u8, .{}); - const ser = try sig.bincode.writeAlloc(std.testing.allocator, set, .{}); - defer std.testing.allocator.free(ser); + var set: Set = .empty; + defer set.deinit(allocator); - var des = try sig.bincode.readFromSlice(std.testing.allocator, SortedSet(u8), ser, .{}); - defer des.deinit(); + try set.put(allocator, 5, {}); + try set.put(allocator, 3, {}); - const items = des.items(); - try std.testing.expectEqual(3, items[0]); - try std.testing.expectEqual(5, items[1]); + const ser = try sig.bincode.writeAlloc(allocator, set, .{}); + defer allocator.free(ser); + + var des = try sig.bincode.readFromSlice(allocator, Set, ser, .{}); + defer des.deinit(allocator); + + var iter = set.iterator(); + try std.testing.expectEqual(3, iter.next().?.key_ptr.*); + try std.testing.expectEqual(5, iter.next().?.key_ptr.*); } -test "SortedSet range" { - var set = SortedSet(u8).init(std.testing.allocator); - defer set.deinit(); - - try set.put(5); - try set.put(3); - try set.put(1); - try set.put(3); - - try expectEqualSlices(u8, &.{ 1, 3, 5 }, set.range(null, null)); - try expectEqualSlices(u8, &.{}, set.range(0, 0)); - try expectEqualSlices(u8, &.{}, set.range(10, 10)); - try expectEqualSlices(u8, &.{}, set.range(10, 11)); - try expectEqualSlices(u8, &.{}, set.range(12, 11)); - try expectEqualSlices(u8, &.{1}, set.range(null, 3)); - try expectEqualSlices(u8, &.{ 1, 3 }, set.range(null, 4)); - try expectEqualSlices(u8, &.{ 1, 3 }, set.range(null, 5)); - try expectEqualSlices(u8, &.{ 1, 3, 5 }, set.range(null, 6)); - try expectEqualSlices(u8, &.{ 1, 3, 5 }, set.range(0, null)); - try expectEqualSlices(u8, &.{ 1, 3, 5 }, set.range(1, null)); - try expectEqualSlices(u8, &.{ 3, 5 }, set.range(2, null)); - try expectEqualSlices(u8, &.{ 3, 5 }, set.range(3, null)); - try expectEqualSlices(u8, &.{5}, set.range(4, null)); - try expectEqualSlices(u8, &.{5}, set.range(5, null)); - try expectEqualSlices(u8, &.{ 1, 3, 5 }, set.range(1, 6)); - try expectEqualSlices(u8, &.{ 1, 3 }, set.range(1, 5)); - try expectEqualSlices(u8, &.{ 1, 3 }, set.range(1, 4)); - try expectEqualSlices(u8, &.{1}, set.range(1, 3)); - try expectEqualSlices(u8, &.{1}, set.range(1, 2)); - try expectEqualSlices(u8, &.{}, set.range(1, 1)); - try expectEqualSlices(u8, &.{ 3, 5 }, set.range(2, 6)); - try expectEqualSlices(u8, &.{ 3, 5 }, set.range(3, 6)); - try expectEqualSlices(u8, &.{5}, set.range(4, 6)); - try expectEqualSlices(u8, &.{5}, set.range(5, 6)); - try expectEqualSlices(u8, &.{3}, set.range(3, 4)); - try expectEqualSlices(u8, &.{}, set.range(3, 3)); - try expectEqualSlices(u8, &.{}, set.range(2, 3)); - try expectEqualSlices(u8, &.{}, set.range(2, 2)); +fn testRange( + expected: []const u8, + iterator: SortedSet(u8, .{}).Iterator, +) !void { + const allocator = std.testing.allocator; + + var found_data: std.ArrayListUnmanaged(u8) = .empty; + defer found_data.deinit(allocator); + + var iter = iterator; + while (iter.next()) |entry| { + try found_data.append(allocator, entry.key_ptr.*); + } + + try std.testing.expectEqualSlices(u8, expected, found_data.items); +} + +test "SortedMap.iteratorRanged" { + const allocator = std.testing.allocator; + + var set: SortedSet(u8, .{}) = .empty; + defer set.deinit(allocator); + + try set.put(allocator, 5, {}); + try set.put(allocator, 3, {}); + try set.put(allocator, 1, {}); + try set.put(allocator, 3, {}); + + try testRange(&.{ 1, 3, 5 }, set.iteratorRanged(null, null, .start)); + try testRange(&.{}, set.iteratorRanged(0, 0, .start)); + try testRange(&.{}, set.iteratorRanged(10, 10, .start)); + try testRange(&.{}, set.iteratorRanged(10, 11, .start)); + try testRange(&.{}, set.iteratorRanged(12, 11, .start)); + try testRange(&.{1}, set.iteratorRanged(null, 3, .start)); + try testRange(&.{ 1, 3 }, set.iteratorRanged(null, 4, .start)); + try testRange(&.{ 1, 3 }, set.iteratorRanged(null, 5, .start)); + try testRange(&.{ 1, 3, 5 }, set.iteratorRanged(null, 6, .start)); + try testRange(&.{ 1, 3, 5 }, set.iteratorRanged(0, null, .start)); + try testRange(&.{ 1, 3, 5 }, set.iteratorRanged(1, null, .start)); + try testRange(&.{ 3, 5 }, set.iteratorRanged(2, null, .start)); + try testRange(&.{ 3, 5 }, set.iteratorRanged(3, null, .start)); + try testRange(&.{5}, set.iteratorRanged(4, null, .start)); + try testRange(&.{5}, set.iteratorRanged(5, null, .start)); + try testRange(&.{ 1, 3, 5 }, set.iteratorRanged(1, 6, .start)); + try testRange(&.{ 1, 3 }, set.iteratorRanged(1, 5, .start)); + try testRange(&.{ 1, 3 }, set.iteratorRanged(1, 4, .start)); + try testRange(&.{1}, set.iteratorRanged(1, 3, .start)); + try testRange(&.{1}, set.iteratorRanged(1, 2, .start)); + try testRange(&.{}, set.iteratorRanged(1, 1, .start)); + try testRange(&.{ 3, 5 }, set.iteratorRanged(2, 6, .start)); + try testRange(&.{ 3, 5 }, set.iteratorRanged(3, 6, .start)); + try testRange(&.{5}, set.iteratorRanged(4, 6, .start)); + try testRange(&.{5}, set.iteratorRanged(5, 6, .start)); + try testRange(&.{3}, set.iteratorRanged(3, 4, .start)); + try testRange(&.{}, set.iteratorRanged(3, 3, .start)); + try testRange(&.{}, set.iteratorRanged(2, 3, .start)); + try testRange(&.{}, set.iteratorRanged(2, 2, .start)); } test binarySearch { @@ -1291,23 +704,6 @@ test "order slices" { try expectEqual(orderSlices(u8, std.math.order, &e, &b), .lt); } -test "sorted set slice range" { - var set = SortedSet([]const u8).init(std.testing.allocator); - defer set.deinit(); - try set.put(&.{ 0, 0, 10 }); - try set.put(&.{ 0, 0, 20 }); - try set.put(&.{ 0, 0, 30 }); - try set.put(&.{ 0, 0, 40 }); - - const range = set.rangeCustom(null, .{ .inclusive = &.{ 0, 0, 40 } }); - - try std.testing.expectEqual(4, range.len); - try std.testing.expectEqualSlices(u8, &.{ 0, 0, 10 }, range[0]); - try std.testing.expectEqualSlices(u8, &.{ 0, 0, 20 }, range[1]); - try std.testing.expectEqualSlices(u8, &.{ 0, 0, 30 }, range[2]); - try std.testing.expectEqualSlices(u8, &.{ 0, 0, 40 }, range[3]); -} - test "binarySearch slice of slices" { const slices = [4][]const u8{ &.{ 0, 0, 10 }, @@ -1451,38 +847,162 @@ test "Window realigns" { } test "SortedMap" { - var map = SortedMap(u64, u64).init(std.testing.allocator); - defer map.deinit(); + const allocator = std.testing.allocator; - try map.put(3, 30); - try map.put(1, 10); - try map.put(2, 20); - try map.put(4, 40); - try map.put(5, 50); + var map: SortedMap( + u64, + u64, + .{}, + ) = .empty; + defer map.deinit(allocator); - // Get the keys and values - const items = map.items(); - const keys = items[0]; - const values = items[1]; + try map.put(allocator, 3, 30); + try map.put(allocator, 1, 10); + try map.put(allocator, 2, 20); + try map.put(allocator, 4, 40); + try map.put(allocator, 5, 50); // Check that the keys and values are sorted. - for (keys, 0..) |key, i| { + var iter = map.iterator(); + var i: u64 = 0; + while (iter.next()) |entry| : (i += 1) { // Keys should be 1, 2, 3, 4, 5 - try expectEqual(key, i + 1); + try expectEqual(entry.key_ptr.*, i + 1); // Values should be 10, 20, 30, 40, 50 - try expectEqual(values[i], (i + 1) * 10); + try expectEqual(entry.value_ptr.*, (i + 1) * 10); } - // Check that the map is sorted - try expect(map.unmanaged.is_sorted); // Remove a non terminal item with no sort. - try expect(map.swapRemoveNoSort(3)); - try expect(!map.swapRemoveNoSort(3)); - try expect(map.swapRemoveNoSort(1)); + try expect(map.remove(3)); + try expect(!map.remove(3)); + try expect(map.remove(1)); +} + +test "SortedMap.put primitives" { + const allocator = std.testing.allocator; - try expect(!map.unmanaged.is_sorted); - map.sort(); - try expect(map.unmanaged.is_sorted); + var map: SortedMap(i32, u128, .{}) = .empty; + defer map.deinit(allocator); + try std.testing.expectEqual(0, map.count()); + + try map.put(allocator, 5, 500); + try std.testing.expectEqual(1, map.count()); + try std.testing.expectEqual(@as(?u128, 500), map.get(5)); + + try map.put(allocator, 5, 600); + try std.testing.expectEqual(@as(u32, 1), map.count()); + try std.testing.expectEqual(@as(?u128, 600), map.get(5)); + + try map.put(allocator, 1, 100); + try map.put(allocator, 3, 300); + try map.put(allocator, 2, 200); + try map.put(allocator, 4, 400); + try map.put(allocator, 5, 500); + + var iter = map.iterator(); + var i: u16 = 1; + while (iter.next()) |entry| : (i += 1) { + try std.testing.expectEqual(i, entry.key_ptr.*); + try std.testing.expectEqual(i * 100, entry.value_ptr.*); + } + + var large_map: SortedMap(i32, i32, .{}) = .empty; + defer large_map.deinit(allocator); + + for (0..100) |j| { + try large_map.put(allocator, @intCast(j), @intCast(j * 10)); + } + try std.testing.expectEqual(@as(u32, 100), large_map.count()); + + var large_iter = large_map.iterator(); + var count: u32 = 0; + while (large_iter.next()) |_| count += 1; + try std.testing.expectEqual(@as(u32, 100), count); +} + +test "SortedMap.getOrPut" { + const allocator = std.testing.allocator; + + var map: SortedMap(i32, i32, .{}) = .empty; + defer map.deinit(allocator); + + const entry1 = try map.getOrPut(allocator, 100); + try std.testing.expect(!entry1.found_existing); + try std.testing.expectEqual(100, entry1.key_ptr.*); + try std.testing.expectEqual(1, map.count()); + entry1.value_ptr.* = 500; + + const entry2 = try map.getOrPut(allocator, 100); + try std.testing.expect(entry2.found_existing); + try std.testing.expectEqual(100, entry2.key_ptr.*); + try std.testing.expectEqual(500, entry2.value_ptr.*); + try std.testing.expectEqual(1, map.count()); + + try std.testing.expectEqual(entry1.value_ptr, entry2.value_ptr); + + (try map.getOrPut(allocator, 200)).value_ptr.* = 2000; + (try map.getOrPut(allocator, 300)).value_ptr.* = 3000; + + try std.testing.expectEqual(3, map.count()); + try std.testing.expectEqual(2000, map.get(200).?); + try std.testing.expectEqual(3000, map.get(300).?); +} + +test "SortedMap.iteratorRanged bidirectional" { + const allocator = std.testing.allocator; + + var map: SortedMap(i32, i32, .{}) = .empty; + defer map.deinit(allocator); + + try map.put(allocator, 0, 0); + try map.put(allocator, 1, 10); + try map.put(allocator, 2, 20); + try map.put(allocator, 3, 30); + + // forward iteration + { + var iter = map.iteratorRanged(null, null, .start); + var count: u32 = 0; + var keys: [4]i32 = undefined; + + while (iter.next()) |entry| { + keys[count] = entry.key_ptr.*; + count += 1; + } + + try std.testing.expectEqual(4, count); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 2, 3 }, &keys); + } + + // backward iteration + { + var iter = map.iteratorRanged(null, null, .end); + var count: u32 = 0; + var keys: [4]i32 = undefined; + + while (iter.prev()) |entry| { + keys[count] = entry.key_ptr.*; + count += 1; + } + + try std.testing.expectEqual(4, count); + try std.testing.expectEqualSlices(i32, &.{ 3, 2, 1, 0 }, &keys); + } + + // backward ranged + { + var iter = map.iteratorRanged(1, 3, .end); + var count: u32 = 0; + var keys: [3]i32 = undefined; + + while (iter.prev()) |entry| { + keys[count] = entry.key_ptr.*; + count += 1; + } + + try std.testing.expectEqual(3, count); + try std.testing.expectEqualSlices(i32, &.{ 3, 2, 1 }, &keys); + } } test "checkAllAllocationFailures in cloneMapAndValues" { diff --git a/src/utils/sorted_map.zig b/src/utils/sorted_map.zig new file mode 100644 index 0000000000..75052c295a --- /dev/null +++ b/src/utils/sorted_map.zig @@ -0,0 +1,411 @@ +const std = @import("std"); +const sig = @import("../sig.zig"); + +const defaultOrderFn = sig.utils.collections.defaultOrderFn; + +pub fn SortedSet( + comptime Key: type, + comptime options: struct { + orderFn: fn (Key, Key) std.math.Order = defaultOrderFn(Key), + }, +) type { + return SortedMap(Key, void, .{ + .orderFn = options.orderFn, + }); +} + +pub fn SortedMap( + comptime Key: type, + comptime Value: type, + comptime options: struct { + orderFn: fn (Key, Key) std.math.Order = defaultOrderFn(Key), + }, +) type { + const orderFn = options.orderFn; + + return struct { + len: u32, + tail: u32, + free_list: u32, + + levels: u32, + prng: std.Random.DefaultPrng, + nodes: std.ArrayListUnmanaged(Node), + + pub const empty: Self = .{ + .len = 0, + .tail = 0, + .free_list = 0, + + .levels = 0, + .prng = .init(0), + .nodes = .{}, + }; + + const max_objects = 1 << 24; // 16 million objects + const max_levels = blk: { + @setEvalBranchQuota(10_000); + var levels: u32 = 0; + while (true) : (levels += 1) { + const limit = std.math.pow(f64, std.math.e, @floatFromInt(levels)); + if (limit >= max_objects) break; + } + break :blk levels; + }; + + const Self = @This(); + const Node = struct { + key: Key, + value: Value, + prev: u32, + next: [max_levels + 1]u32, + }; + + pub fn deinit(const_self: Self, allocator: std.mem.Allocator) void { + var self = const_self; + self.nodes.deinit(allocator); + } + + pub fn count(self: *const Self) usize { + return self.len; + } + + fn search(self: *const Self, key: Key, update: *[max_levels + 1]u32) ?u32 { + if (self.count() == 0) return null; + + var prev: u32 = 0; + var i = self.levels; + while (i != std.math.maxInt(u32)) : (i -%= 1) { + while (true) { + const next = self.nodes.items[prev].next[i]; + if (next == 0) break; + if (orderFn(self.nodes.items[next].key, key) != .lt) break; + prev = next; + } + update[i] = prev; + } + + const next = self.nodes.items[prev].next[0]; + if (next > 0 and orderFn(self.nodes.items[next].key, key) == .eq) { + return next; + } + + return null; + } + + pub fn contains(self: *const Self, key: Key) bool { + return self.getEntry(key) != null; + } + + pub fn get(self: *const Self, key: Key) ?Value { + const value_ptr = self.getPtr(key) orelse return null; + return value_ptr.*; + } + + pub fn getPtr(self: *const Self, key: Key) ?*Value { + const entry = self.getEntry(key) orelse return null; + return entry.value_ptr; + } + + pub const Entry = struct { + key_ptr: *Key, + value_ptr: *Value, + }; + + pub fn getEntry(self: *const Self, key: Key) ?Entry { + var update: [max_levels + 1]u32 = undefined; + const idx = self.search(key, &update) orelse return null; + + const node = &self.nodes.items[idx]; + return .{ .key_ptr = &node.key, .value_ptr = &node.value }; + } + + pub const GetOrPutResult = struct { + key_ptr: *const Key, + value_ptr: *Value, + found_existing: bool, + }; + + pub fn getOrPut(self: *Self, allocator: std.mem.Allocator, key: Key) !GetOrPutResult { + // lazy init + if (self.nodes.items.len == 0) { + @branchHint(.unlikely); + (try self.nodes.addOne(allocator)).* = .{ + .key = undefined, + .value = undefined, + .prev = 0, + .next = @splat(0), + }; + } + + var update: [max_levels + 1]u32 = @splat(0); + if (self.search(key, &update)) |idx| { + const node = &self.nodes.items[idx]; + return .{ .key_ptr = &node.key, .value_ptr = &node.value, .found_existing = true }; + } + + // allocate new node idx. + var idx = self.free_list; + if (idx > 0) { + self.free_list = self.nodes.items[idx].next[0]; + } else { + idx = @intCast(self.nodes.items.len); + _ = try self.nodes.addOne(allocator); + } + + const node = &self.nodes.items[idx]; + node.* = .{ .key = key, .value = undefined, .prev = update[0], .next = @splat(0) }; + + const lvl = self.prng.random().uintAtMost(u32, max_levels); + if (lvl > self.levels) { + for (self.levels + 1..lvl + 1) |i| update[i] = 0; + self.levels = lvl; + } + + for (0..lvl + 1) |i| { + const update_node = &self.nodes.items[update[i]]; + node.next[i] = update_node.next[i]; + update_node.next[i] = idx; + } + + if (node.next[0] > 0) { + std.debug.assert(self.nodes.items[node.next[0]].prev == update[0]); + self.nodes.items[node.next[0]].prev = idx; + } else { + std.debug.assert(self.tail == update[0]); + self.tail = idx; + } + + self.len += 1; + std.debug.assert(self.len <= max_objects); + return .{ .key_ptr = &node.key, .value_ptr = &node.value, .found_existing = false }; + } + + pub const KV = struct { + key: Key, + value: Value, + }; + + pub fn fetchPut(self: *Self, allocator: std.mem.Allocator, key: Key, value: Value) !?KV { + const gop = try self.getOrPut(allocator, key); + const result: ?KV = if (gop.found_existing) + .{ .key = gop.key_ptr.*, .value = gop.value_ptr.* } + else + null; + + gop.value_ptr.* = value; + return result; + } + + pub fn put(self: *Self, allocator: std.mem.Allocator, key: Key, value: Value) !void { + const gop = try self.getOrPut(allocator, key); + gop.value_ptr.* = value; + } + + pub fn fetchRemove(self: *Self, key: Key) ?KV { + var kv: KV = undefined; + return if (self.removeByKey(key, &kv)) kv else null; + } + + pub fn remove(self: *Self, key: Key) bool { + return self.removeByKey(key, null); + } + + fn removeByKey(self: *Self, key: Key, maybe_out_kv: ?*KV) bool { + var update: [max_levels + 1]u32 = @splat(0); + const idx = self.search(key, &update) orelse return false; + + const node = &self.nodes.items[idx]; + std.debug.assert(node.prev == update[0]); + + if (maybe_out_kv) |out_kv| { + out_kv.key = node.key; + out_kv.value = node.value; + } + + self.len -= 1; + if (node.next[0] > 0) { + std.debug.assert(idx != self.tail); + std.debug.assert(self.nodes.items[node.next[0]].prev == idx); + self.nodes.items[node.next[0]].prev = node.prev; + } else { + std.debug.assert(idx == self.tail); + self.tail = node.prev; + } + + for (0..self.levels + 1) |i| { + const update_node = &self.nodes.items[update[i]]; + if (update_node.next[i] != idx) break; + update_node.next[i] = node.next[i]; + } + + while (self.levels > 0 and self.nodes.items[0].next[self.levels] == 0) + self.levels -= 1; + + node.next[0] = self.free_list; + self.free_list = idx; + return true; + } + + pub fn clone(self: *const Self, allocator: std.mem.Allocator) !Self { + return Self{ + .len = self.len, + .tail = self.tail, + .free_list = self.free_list, + .levels = self.levels, + .prng = self.prng, + .nodes = try self.nodes.clone(allocator), + }; + } + + pub fn minEntry(self: *const Self) ?Entry { + if (self.count() == 0) return null; + const min_idx = self.nodes.items[0].next[0]; + if (min_idx == 0) return null; + const node = &self.nodes.items[min_idx]; + return .{ .key_ptr = &node.key, .value_ptr = &node.value }; + } + + pub fn maxEntry(self: *const Self) ?Entry { + if (self.count() == 0) return null; + const node = &self.nodes.items[self.tail]; + return .{ .key_ptr = &node.key, .value_ptr = &node.value }; + } + + pub fn eql(self: *const Self, other: *const Self) bool { + var self_it = self.iterator(); + var other_it = other.iterator(); + while (true) { + const a = self_it.next() orelse (return other_it.next() == null); + const b = other_it.next() orelse return false; + if (orderFn(a.key_ptr.*, b.key_ptr.*) != .eq) return false; + } + } + + pub fn iterator(self: *const Self) Iterator { + return self.iteratorRanged(null, null, .start); + } + + pub fn iteratorRanged( + self: *const Self, + maybe_start: ?Key, + maybe_end: ?Key, + begin: enum { start, end }, + ) Iterator { + const idx = blk: { + if (self.count() == 0) break :blk 0; + const key = switch (begin) { + .start => maybe_start orelse break :blk 0, // min + .end => maybe_end orelse break :blk self.tail, // max, + }; + + var update: [max_levels + 1]u32 = undefined; + break :blk self.search(key, &update) orelse update[0]; + }; + + return .{ + .map = self, + .idx = idx, + .start = maybe_start, + .end = maybe_end, + }; + } + + pub const Iterator = struct { + map: *const Self, + idx: u32, + start: ?Key, + end: ?Key, + + pub fn countForwardsInclusive(self: *const Iterator) u32 { + return self.countForwardWith(.inclusive); + } + + pub fn countForwards(self: *const Iterator) u32 { + return self.countForwardWith(.exclusive); + } + + fn countForwardWith(self: *const Iterator, include: enum { inclusive, exclusive }) u32 { + // TODO: make this not O(N) + var n: u32 = 0; + var it = self.*; + while (it.advance( + .forward, + if (include == .inclusive) .inclusive else .exclusive, + ) != null) n += 1; + return n; + } + + pub fn next(self: *Iterator) ?Entry { + return self.advance(.forward, .exclusive); + } + + pub fn nextInclusive(self: *Iterator) ?Entry { + return self.advance(.forward, .inclusive); + } + + pub fn prev(self: *Iterator) ?Entry { + return self.advance(.backward, .inclusive); + } + + fn advance( + self: *Iterator, + direction: enum { forward, backward }, + include: enum { inclusive, exclusive }, + ) ?Entry { + if (self.map.count() == 0) return null; + while (true) { + // start of range + if (self.idx == 0) { + self.idx = switch (direction) { + .backward => return null, + .forward => switch (self.map.nodes.items[0].next[0]) { + 0 => std.math.maxInt(u32), + else => |idx| idx, + }, + }; + } + // end of range + if (self.idx == std.math.maxInt(u32)) { + switch (direction) { + .forward => return null, + .backward => { + self.idx = self.map.tail; + std.debug.assert(self.idx != std.math.maxInt(u32)); + continue; + }, + } + } + // valid node + const node_idx = self.idx; + const node = &self.map.nodes.items[self.idx]; + self.idx = switch (direction) { + .backward => node.prev, + .forward => switch (node.next[0]) { + 0 => std.math.maxInt(u32), + else => |idx| idx, + }, + }; + // skip nodes outside range + if (self.start) |start_key| { + switch (direction) { + .forward => if (orderFn(start_key, node.key) == .gt) continue, + .backward => if (orderFn(start_key, node.key) == .gt) { + self.idx = node_idx; // stay on this one + return null; + }, + } + } + if (self.end) |end_key| { + switch (include) { + .inclusive => if (orderFn(node.key, end_key) == .gt) continue, + .exclusive => if (orderFn(node.key, end_key) != .lt) continue, + } + } + // found node in range + return .{ .key_ptr = &node.key, .value_ptr = &node.value }; + } + } + }; + }; +}