fix(RefCountedSet): add NeedsRehash error and fix PSL counting bug

Prevent bad input from causing repeated OutOfMemory errors by erroring
with NeedsRehash instead when there are unused dead IDs available.

Additionally, properly decrement PSL stats when reviving dead IDs.

src/inspector/page.zig

@@ -37,7 +37,7 @@ pub fn render(page: *const terminal.Page) void {
         }
         {
             _ = cimgui.c.igTableSetColumnIndex(1);
-            cimgui.c.igText("%d", page.styles.count(page.memory));
+            cimgui.c.igText("%d", page.styles.count());
         }
     }
     {

src/terminal/Screen.zig

@@ -1175,7 +1175,7 @@ pub fn setAttribute(self: *Screen, attr: sgr.Attribute) !void {
 pub fn manualStyleUpdate(self: *Screen) !void {
     var page = &self.cursor.page_pin.page.data;
 
-    // std.log.warn("active styles={}", .{page.styles.count(page.memory)});
+    // std.log.warn("active styles={}", .{page.styles.count()});
 
     // Release our previous style if it was not default.
     if (self.cursor.style_id != style.default_id) {
@@ -1201,12 +1201,17 @@ pub fn manualStyleUpdate(self: *Screen) !void {
     const id = page.styles.add(
         page.memory,
         self.cursor.style,
-    ) catch id: {
+    ) catch |err| id: {
-        // Our style map is full. Let's allocate a new
+        // Our style map is full or needs to be rehashed,
-        // page by doubling the size and then try again.
+        // so we allocate a new page, which will rehash,
+        // and double the style capacity for it if it was
+        // full.
         const node = try self.pages.adjustCapacity(
             self.cursor.page_pin.page,
-            .{ .styles = page.capacity.styles * 2 },
+            switch (err) {
+                error.OutOfMemory => .{ .styles = page.capacity.styles * 2 },
+                error.NeedsRehash => .{},
+            },
         );
 
         page = &node.data;
@@ -2388,17 +2393,17 @@ test "Screen cursorCopy style deref" {
 
     var s2 = try Screen.init(alloc, 10, 10, 0);
     defer s2.deinit();
-    const page = s2.cursor.page_pin.page.data;
+    const page = &s2.cursor.page_pin.page.data;
 
     // Bold should create our style
     try s2.setAttribute(.{ .bold = {} });
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
     try testing.expect(s2.cursor.style.flags.bold);
 
     // Copy default style, should release our style
     try s2.cursorCopy(s.cursor);
     try testing.expect(!s2.cursor.style.flags.bold);
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Screen cursorCopy style copy" {
@@ -2411,10 +2416,10 @@ test "Screen cursorCopy style copy" {
 
     var s2 = try Screen.init(alloc, 10, 10, 0);
     defer s2.deinit();
-    const page = s2.cursor.page_pin.page.data;
+    const page = &s2.cursor.page_pin.page.data;
     try s2.cursorCopy(s.cursor);
     try testing.expect(s2.cursor.style.flags.bold);
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 }
 
 test "Screen style basics" {
@@ -2423,19 +2428,19 @@ test "Screen style basics" {
 
     var s = try Screen.init(alloc, 80, 24, 1000);
     defer s.deinit();
-    const page = s.cursor.page_pin.page.data;
+    const page = &s.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 
     // Set a new style
     try s.setAttribute(.{ .bold = {} });
     try testing.expect(s.cursor.style_id != 0);
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
     try testing.expect(s.cursor.style.flags.bold);
 
     // Set another style, we should still only have one since it was unused
     try s.setAttribute(.{ .italic = {} });
     try testing.expect(s.cursor.style_id != 0);
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
     try testing.expect(s.cursor.style.flags.italic);
 }
 
@@ -2445,18 +2450,18 @@ test "Screen style reset to default" {
 
     var s = try Screen.init(alloc, 80, 24, 1000);
     defer s.deinit();
-    const page = s.cursor.page_pin.page.data;
+    const page = &s.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 
     // Set a new style
     try s.setAttribute(.{ .bold = {} });
     try testing.expect(s.cursor.style_id != 0);
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 
     // Reset to default
     try s.setAttribute(.{ .reset_bold = {} });
     try testing.expect(s.cursor.style_id == 0);
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Screen style reset with unset" {
@@ -2465,18 +2470,18 @@ test "Screen style reset with unset" {
 
     var s = try Screen.init(alloc, 80, 24, 1000);
     defer s.deinit();
-    const page = s.cursor.page_pin.page.data;
+    const page = &s.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 
     // Set a new style
     try s.setAttribute(.{ .bold = {} });
     try testing.expect(s.cursor.style_id != 0);
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 
     // Reset to default
     try s.setAttribute(.{ .unset = {} });
     try testing.expect(s.cursor.style_id == 0);
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Screen clearRows active one line" {
@@ -2522,13 +2527,13 @@ test "Screen clearRows active styled line" {
     try s.setAttribute(.{ .unset = {} });
 
     // We should have one style
-    const page = s.cursor.page_pin.page.data;
+    const page = &s.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 
     s.clearRows(.{ .active = .{} }, null, false);
 
     // We should have none because active cleared it
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 
     const str = try s.dumpStringAlloc(alloc, .{ .screen = .{} });
     defer alloc.free(str);

src/terminal/Terminal.zig

@@ -2826,8 +2826,8 @@ test "Terminal: print over wide char with bold" {
     try t.print(0x1F600); // Smiley face
     // verify we have styles in our style map
     {
-        const page = t.screen.cursor.page_pin.page.data;
+        const page = &t.screen.cursor.page_pin.page.data;
-        try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+        try testing.expectEqual(@as(usize, 1), page.styles.count());
     }
 
     // Go back and overwrite with no style
@@ -2837,8 +2837,8 @@ test "Terminal: print over wide char with bold" {
 
     // verify our style is gone
     {
-        const page = t.screen.cursor.page_pin.page.data;
+        const page = &t.screen.cursor.page_pin.page.data;
-        try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+        try testing.expectEqual(@as(usize, 0), page.styles.count());
     }
 
     try testing.expect(t.isDirty(.{ .screen = .{ .x = 0, .y = 0 } }));
@@ -2856,8 +2856,8 @@ test "Terminal: print over wide char with bg color" {
     try t.print(0x1F600); // Smiley face
     // verify we have styles in our style map
     {
-        const page = t.screen.cursor.page_pin.page.data;
+        const page = &t.screen.cursor.page_pin.page.data;
-        try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+        try testing.expectEqual(@as(usize, 1), page.styles.count());
     }
 
     // Go back and overwrite with no style
@@ -2867,8 +2867,8 @@ test "Terminal: print over wide char with bg color" {
 
     // verify our style is gone
     {
-        const page = t.screen.cursor.page_pin.page.data;
+        const page = &t.screen.cursor.page_pin.page.data;
-        try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+        try testing.expectEqual(@as(usize, 0), page.styles.count());
     }
 
     try testing.expect(t.isDirty(.{ .screen = .{ .x = 0, .y = 0 } }));
@@ -3322,7 +3322,7 @@ test "Terminal: overwrite multicodepoint grapheme clears grapheme data" {
     try testing.expectEqual(@as(usize, 2), t.screen.cursor.x);
 
     // We should have one cell with graphemes
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
     try testing.expectEqual(@as(usize, 1), page.graphemeCount());
 
     // Move back and overwrite wide
@@ -3362,7 +3362,7 @@ test "Terminal: overwrite multicodepoint grapheme tail clears grapheme data" {
     try testing.expectEqual(@as(usize, 2), t.screen.cursor.x);
 
     // We should have one cell with graphemes
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
     try testing.expectEqual(@as(usize, 1), page.graphemeCount());
 
     // Move back and overwrite wide
@@ -4534,8 +4534,8 @@ test "Terminal: insertLines handles style refs" {
     try t.setAttribute(.{ .unset = {} });
 
     // verify we have styles in our style map
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 
     t.setCursorPos(2, 2);
     t.insertLines(1);
@@ -4547,7 +4547,7 @@ test "Terminal: insertLines handles style refs" {
     }
 
     // verify we have no styles in our style map
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Terminal: insertLines outside of scroll region" {
@@ -5336,14 +5336,14 @@ test "Terminal: eraseChars handles refcounted styles" {
     try t.print('C');
 
     // verify we have styles in our style map
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 
     t.setCursorPos(1, 1);
     t.eraseChars(2);
 
     // verify we have no styles in our style map
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Terminal: eraseChars protected attributes respected with iso" {
@@ -7043,7 +7043,7 @@ test "Terminal: bold style" {
         const cell = list_cell.cell;
         try testing.expectEqual(@as(u21, 'A'), cell.content.codepoint);
         try testing.expect(cell.style_id != 0);
-        const page = t.screen.cursor.page_pin.page.data;
+        const page = &t.screen.cursor.page_pin.page.data;
         try testing.expect(page.styles.refCount(page.memory, t.screen.cursor.style_id) > 1);
     }
 }
@@ -7067,8 +7067,8 @@ test "Terminal: garbage collect overwritten" {
     }
 
     // verify we have no styles in our style map
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 0), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 0), page.styles.count());
 }
 
 test "Terminal: do not garbage collect old styles in use" {
@@ -7089,8 +7089,8 @@ test "Terminal: do not garbage collect old styles in use" {
     }
 
     // verify we have no styles in our style map
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
-    try testing.expectEqual(@as(usize, 1), page.styles.count(page.memory));
+    try testing.expectEqual(@as(usize, 1), page.styles.count());
 }
 
 test "Terminal: print with style marks the row as styled" {
@@ -7425,7 +7425,7 @@ test "Terminal: insertBlanks deleting graphemes" {
     try t.print(0x1F467);
 
     // We should have one cell with graphemes
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
     try testing.expectEqual(@as(usize, 1), page.graphemeCount());
 
     t.setCursorPos(1, 1);
@@ -7461,7 +7461,7 @@ test "Terminal: insertBlanks shift graphemes" {
     try t.print(0x1F467);
 
     // We should have one cell with graphemes
-    const page = t.screen.cursor.page_pin.page.data;
+    const page = &t.screen.cursor.page_pin.page.data;
     try testing.expectEqual(@as(usize, 1), page.graphemeCount());
 
     t.setCursorPos(1, 1);

src/terminal/page.zig

@@ -500,7 +500,7 @@ pub const Page = struct {
         return result;
     }
 
-    pub const CloneFromError = Allocator.Error || error{OutOfMemory};
+    pub const CloneFromError = Allocator.Error || style.Set.AddError;
 
     /// Clone the contents of another page into this page. The capacities
     /// can be different, but the size of the other page must fit into

src/terminal/ref_counted_set.zig

@@ -13,9 +13,9 @@ const fastmem = @import("../fastmem.zig");
 /// the exact memory requirement of a given capacity by calling `layout`
 /// and checking the total size.
 ///
-/// When the set exceeds capacity, `error.OutOfMemory` is returned from
+/// When the set exceeds capacity, an `OutOfMemory` or `NeedsRehash` error
-/// any memory-using methods. The caller is responsible for determining
+/// is returned from any memory-using methods. The caller is responsible
-/// a path forward.
+/// for determining a path forward.
 ///
 /// This set is reference counted. Each item in the set has an associated
 /// reference count. The caller is responsible for calling release for an
@@ -108,6 +108,9 @@ pub fn RefCountedSet(
         /// The backing store of items
         items: Offset(Item),
 
+        /// The number of living items currently stored in the set.
+        living: Id = 0,
+
         /// The next index to store an item at.
         /// Id 0 is reserved for unused items.
         next_id: Id = 1,
@@ -185,12 +188,23 @@ pub fn RefCountedSet(
             };
         }
 
+        /// Possible errors for `add` and `addWithId`.
+        pub const AddError = error{
+            /// There is not enough memory to add a new item.
+            /// Remove items or grow and reinitialize.
+            OutOfMemory,
+
+            /// The set needs to be rehashed, as there are many dead
+            /// items with lower IDs which are inaccessible for re-use.
+            NeedsRehash,
+        };
+
         /// Add an item to the set if not present and increment its ref count.
         ///
         /// Returns the item's ID.
         ///
         /// If the set has no more room, then an OutOfMemory error is returned.
-        pub fn add(self: *Self, base: anytype, value: T) error{OutOfMemory}!Id {
+        pub fn add(self: *Self, base: anytype, value: T) AddError!Id {
             const items = self.items.ptr(base);
 
             // Trim dead items from the end of the list.
@@ -199,14 +213,34 @@ pub fn RefCountedSet(
                 self.deleteItem(base, self.next_id);
             }
 
-            // If we still don't have an available ID, we're out of memory.
+            // If we still don't have an available ID, we can't continue.
-            if (self.next_id >= self.layout.cap) return error.OutOfMemory;
+            if (self.next_id >= self.layout.cap) {
+                // Arbitrarily chosen, threshold for rehashing.
+                // If less than 90% of currently allocated IDs
+                // correspond to living items, we should rehash.
+                // Otherwise, claim we're out of memory because
+                // we assume that we'll end up running out of
+                // memory or rehashing again very soon if we
+                // rehash with only a few IDs left.
+                const rehash_threshold = 0.9;
+                if (self.living < @as(Id, @intFromFloat(@as(f64, @floatFromInt(self.layout.cap)) * rehash_threshold))) {
+                    return AddError.NeedsRehash;
+                }
+
+                // If we don't have at least 10% dead items then
+                // we claim we're out of memory.
+                return AddError.OutOfMemory;
+            }
 
             const id = self.upsert(base, value, self.next_id);
             items[id].meta.ref += 1;
 
             if (id == self.next_id) self.next_id += 1;
 
+            if (items[id].meta.ref == 1) {
+                self.living += 1;
+            }
+
             return id;
         }
 
@@ -216,7 +250,7 @@ pub fn RefCountedSet(
         /// Returns the item's ID, or null if the provided ID was used.
         ///
         /// If the set has no more room, then an OutOfMemory error is returned.
-        pub fn addWithId(self: *Self, base: anytype, value: T, id: Id) error{OutOfMemory}!?Id {
+        pub fn addWithId(self: *Self, base: anytype, value: T, id: Id) AddError!?Id {
             const items = self.items.ptr(base);
 
             if (id < self.next_id) {
@@ -227,6 +261,8 @@ pub fn RefCountedSet(
 
                     items[added_id].meta.ref += 1;
 
+                    self.living += 1;
+
                     return if (added_id == id) null else added_id;
                 } else if (self.context.eql(value, items[id].value)) {
                     items[id].meta.ref += 1;
@@ -303,6 +339,10 @@ pub fn RefCountedSet(
 
             assert(item.meta.ref > 0);
             item.meta.ref -= 1;
+
+            if (item.meta.ref == 0) {
+                self.living -= 1;
+            }
         }
 
         /// Release a specified number of references to an item by its ID.
@@ -317,6 +357,10 @@ pub fn RefCountedSet(
 
             assert(item.meta.ref >= n);
             item.meta.ref -= n;
+
+            if (item.meta.ref == 0) {
+                self.living -= 1;
+            }
         }
 
         /// Get the ref count for an item by its ID.
@@ -330,42 +374,8 @@ pub fn RefCountedSet(
         }
 
         /// Get the current number of non-dead items in the set.
-        ///
+        pub fn count(self: *Self) usize {
-        /// NOT DESIGNED TO BE USED OUTSIDE OF TESTING, this is a very slow
+            return self.living;
-        /// operation, since it traverses the entire structure to count.
-        ///
-        /// Additionally, because this is a testing method, it does extra
-        /// work to verify the integrity of the structure when called.
-        pub fn count(self: *const Self, base: anytype) usize {
-            const table = self.table.ptr(base);
-            const items = self.items.ptr(base);
-
-            // The number of items accessible through the table.
-            var tb_ct: usize = 0;
-
-            for (table[0..self.layout.table_cap]) |id| {
-                if (id != 0) {
-                    const item = items[id];
-                    if (item.meta.ref > 0) {
-                        tb_ct += 1;
-                    }
-                }
-            }
-
-            // The number of items accessible through the backing store.
-            // The two counts should always match- it shouldn't be possible
-            // to have untracked items in the backing store.
-            var it_ct: usize = 0;
-
-            for (items[0..self.layout.cap]) |it| {
-                if (it.meta.ref > 0) {
-                    it_ct += 1;
-                }
-            }
-
-            assert(tb_ct == it_ct);
-
-            return tb_ct;
         }
 
         /// Delete an item, removing any references from
@@ -509,6 +519,7 @@ pub fn RefCountedSet(
                     chosen_id = id;
 
                     held_item.meta.bucket = p;
+                    self.psl_stats[item.meta.psl] -= 1;
                     self.psl_stats[held_item.meta.psl] += 1;
                     self.max_psl = @max(self.max_psl, held_item.meta.psl);