From 00cf0e4a66b730b734305cb05206847829c24e18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Motiejus=20Jak=C5=A1tys?= <motiejus@jakstys.lt>
Date: Thu, 17 Feb 2022 11:04:32 +0200
Subject: [PATCH] hashmap: use getOrPut

---
 src/shell.zig | 35 +++++++++++++++--------------------
 1 file changed, 15 insertions(+), 20 deletions(-)

diff --git a/src/shell.zig b/src/shell.zig
index 7ed0da5..0adf7ce 100644
--- a/src/shell.zig
+++ b/src/shell.zig
@@ -5,6 +5,7 @@ const StringArrayHashMap = std.StringArrayHashMap;
 const StringHashMap = std.StringHashMap;
 const BoundedArray = std.BoundedArray;
 const testing = std.testing;
+const StringContext = std.hash_map.StringContext;
 
 // MaxShells is the maximum number of "popular" shells.
 const MaxShells = 63;
@@ -34,7 +35,7 @@ const ShellReader = struct {
     }
 
     // get returns a shell at the given index.
-    pub fn get(self: *ShellReader, idx: u10) []const u8 {
+    pub fn get(self: *const ShellReader, idx: u10) []const u8 {
         const shellIndex = self.sectionIndex[idx];
         const start = shellIndex.offset << 2;
         const end = start + shellIndex.len + 1;
@@ -88,11 +89,11 @@ const ShellWriter = struct {
             return self;
         }
 
-        pub fn sectionIndex(self: *ShellSections) []const u8 {
+        pub fn sectionIndex(self: *const ShellSections) []const u8 {
             return @bitCast([]const u8, self.index.constSlice());
         }
 
-        pub fn sectionBlob(self: *ShellSections) []const u8 {
+        pub fn sectionBlob(self: *const ShellSections) []const u8 {
             return self.blob.constSlice();
         }
 
@@ -101,7 +102,7 @@ const ShellWriter = struct {
             self.* = undefined;
         }
 
-        pub fn getIndex(self: *ShellSections, shell: []const u8) ?u10 {
+        pub fn getIndex(self: *const ShellSections, shell: []const u8) ?u10 {
             return self.indices.get(shell);
         }
     };
@@ -123,14 +124,16 @@ const ShellWriter = struct {
     }
 
     pub fn put(self: *ShellWriter, shell: []const u8) !void {
-        // TODO getOrPutAdapted may be more elegant, not sure which
-        // context to pass.
-        if (self.counts.getPtr(shell)) |ptr| {
-            ptr.* += 1;
+        const res = try self.counts.getOrPutAdapted(shell, self.counts.ctx);
+
+        if (res.found_existing) {
+            res.value_ptr.* += 1;
         } else {
-            var ourShell = try self.allocator.alloc(u8, shell.len);
+            // TODO(motiejus): can we avoid `ourShell` variable here?
+            const ourShell = try self.allocator.alloc(u8, shell.len);
             std.mem.copy(u8, ourShell, shell);
-            try self.counts.put(ourShell, 1);
+            res.key_ptr.* = ourShell;
+            res.value_ptr.* = 1;
         }
     }
 
@@ -160,7 +163,7 @@ const ShellWriter = struct {
 
         var i: u32 = 0;
         while (i < total) {
-            const elem: []const u8 = deque.removeMax().shell;
+            const elem = deque.removeMax().shell;
             topShells.set(i, elem);
             i += 1;
         }
@@ -206,21 +209,13 @@ test "basic shellpopcon" {
     try testing.expectEqual(sections.getIndex(bash).?, 2);
     try testing.expectEqual(sections.getIndex(nobody), null);
 
-    var shellReader = ShellReader.init(
+    const shellReader = ShellReader.init(
         sections.sectionIndex(),
         sections.sectionBlob(),
     );
     try testing.expectEqualStrings(shellReader.get(0), long);
     try testing.expectEqualStrings(shellReader.get(1), zsh);
     try testing.expectEqualStrings(shellReader.get(2), bash);
-
-    for ([_][]const u8{ long, zsh, bash }) |shell| {
-        const idx = sections.getIndex(shell).?;
-        const start = sections.index.get(idx).offset << 2;
-        const end = start + sections.index.get(idx).len + 1;
-        const got = sections.blob.constSlice()[start..end];
-        try testing.expectEqualStrings(got, shell);
-    }
 }
 
 test "padding" {