perf(collections): replace SortedMap with an internally sorted impl

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, {});
             }
         }
 

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,21 +4672,21 @@ 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);
     }
 
     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,14 +5433,16 @@ 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) };
 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, {});
         }
     }
 }

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 {