Commit b25e58b0ac

fn ⌃ ⌥ <70830482+FnControlOption@users.noreply.github.com>
2021-08-31 06:32:34
std.hash_map: add getKey methods (#9607)
1 parent ede47d4
Changed files (4)
lib/std/array_hash_map.zig
@@ -293,6 +293,22 @@ pub fn ArrayHashMap(
             return self.unmanaged.getPtrAdapted(key, ctx);
         }
 
+        /// Find the actual key associated with an adapted key
+        pub fn getKey(self: Self, key: K) ?K {
+            return self.unmanaged.getKeyContext(key, self.ctx);
+        }
+        pub fn getKeyAdapted(self: Self, key: anytype, ctx: anytype) ?K {
+            return self.unmanaged.getKeyAdapted(key, ctx);
+        }
+
+        /// Find a pointer to the actual key associated with an adapted key
+        pub fn getKeyPtr(self: Self, key: K) ?*K {
+            return self.unmanaged.getKeyPtrContext(key, self.ctx);
+        }
+        pub fn getKeyPtrAdapted(self: Self, key: anytype, ctx: anytype) ?*K {
+            return self.unmanaged.getKeyPtrAdapted(key, ctx);
+        }
+
         /// Check whether a key is stored in the map
         pub fn contains(self: Self, key: K) bool {
             return self.unmanaged.containsContext(key, self.ctx);
@@ -967,6 +983,34 @@ pub fn ArrayHashMapUnmanaged(
             return if (@sizeOf(*V) == 0) @as(*V, undefined) else &self.values()[index];
         }
 
+        /// Find the actual key associated with an adapted key
+        pub fn getKey(self: Self, key: K) ?K {
+            if (@sizeOf(Context) != 0)
+                @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call getKeyContext instead.");
+            return self.getKeyContext(key, undefined);
+        }
+        pub fn getKeyContext(self: Self, key: K, ctx: Context) ?K {
+            return self.getKeyAdapted(key, ctx);
+        }
+        pub fn getKeyAdapted(self: Self, key: anytype, ctx: anytype) ?K {
+            const index = self.getIndexAdapted(key, ctx) orelse return null;
+            return self.keys()[index];
+        }
+
+        /// Find a pointer to the actual key associated with an adapted key
+        pub fn getKeyPtr(self: Self, key: K) ?*K {
+            if (@sizeOf(Context) != 0)
+                @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call getKeyPtrContext instead.");
+            return self.getKeyPtrContext(key, undefined);
+        }
+        pub fn getKeyPtrContext(self: Self, key: K, ctx: Context) ?*K {
+            return self.getKeyPtrAdapted(key, ctx);
+        }
+        pub fn getKeyPtrAdapted(self: Self, key: anytype, ctx: anytype) ?*K {
+            const index = self.getIndexAdapted(key, ctx) orelse return null;
+            return &self.keys()[index];
+        }
+
         /// Check whether a key is stored in the map
         pub fn contains(self: Self, key: K) bool {
             if (@sizeOf(Context) != 0)
lib/std/hash_map.zig
@@ -561,6 +561,21 @@ pub fn HashMap(
             return self.unmanaged.getPtrAdapted(key, ctx);
         }
 
+        /// Finds the actual key associated with an adapted key in the map
+        pub fn getKey(self: Self, key: K) ?K {
+            return self.unmanaged.getKeyContext(key, self.ctx);
+        }
+        pub fn getKeyAdapted(self: Self, key: anytype, ctx: anytype) ?K {
+            return self.unmanaged.getKeyAdapted(key, ctx);
+        }
+
+        pub fn getKeyPtr(self: Self, key: K) ?*K {
+            return self.unmanaged.getKeyPtrContext(key, self.ctx);
+        }
+        pub fn getKeyPtrAdapted(self: Self, key: anytype, ctx: anytype) ?*K {
+            return self.unmanaged.getKeyPtrAdapted(key, ctx);
+        }
+
         /// Finds the key and value associated with a key in the map
         pub fn getEntry(self: Self, key: K) ?Entry {
             return self.unmanaged.getEntryContext(key, self.ctx);
@@ -1124,6 +1139,38 @@ pub fn HashMapUnmanaged(
             result.value_ptr.* = value;
         }
 
+        /// Get an optional pointer to the actual key associated with adapted key, if present.
+        pub fn getKeyPtr(self: Self, key: K) ?*K {
+            if (@sizeOf(Context) != 0)
+                @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call getKeyPtrContext instead.");
+            return self.getKeyPtrContext(key, undefined);
+        }
+        pub fn getKeyPtrContext(self: Self, key: K, ctx: Context) ?*K {
+            return self.getKeyPtrAdapted(key, ctx);
+        }
+        pub fn getKeyPtrAdapted(self: Self, key: anytype, ctx: anytype) ?*K {
+            if (self.getIndex(key, ctx)) |idx| {
+                return &self.keys()[idx];
+            }
+            return null;
+        }
+
+        /// Get a copy of the actual key associated with adapted key, if present.
+        pub fn getKey(self: Self, key: K) ?K {
+            if (@sizeOf(Context) != 0)
+                @compileError("Cannot infer context " ++ @typeName(Context) ++ ", call getKeyContext instead.");
+            return self.getKeyContext(key, undefined);
+        }
+        pub fn getKeyContext(self: Self, key: K, ctx: Context) ?K {
+            return self.getKeyAdapted(key, ctx);
+        }
+        pub fn getKeyAdapted(self: Self, key: anytype, ctx: anytype) ?K {
+            if (self.getIndex(key, ctx)) |idx| {
+                return self.keys()[idx];
+            }
+            return null;
+        }
+
         /// Get an optional pointer to the value associated with key, if present.
         pub fn getPtr(self: Self, key: K) ?*V {
             if (@sizeOf(Context) != 0)
@@ -1948,6 +1995,7 @@ test "std.hash_map getOrPutAdapted" {
         try testing.expect(result.found_existing);
         try testing.expectEqual(real_keys[i], result.key_ptr.*);
         try testing.expectEqual(@as(u64, i) * 2, result.value_ptr.*);
+        try testing.expectEqual(real_keys[i], map.getKeyAdapted(key_str, AdaptedContext{}).?);
     }
 }
 
src/link/MachO/TextBlock.zig
@@ -656,7 +656,7 @@ fn initRelocFromObject(rel: macho.relocation_info, context: RelocContext) !Reloc
             parsed_rel.where = .local;
             parsed_rel.where_index = where_index;
         } else {
-            const n_strx = context.macho_file.strtab_dir.getAdapted(@as([]const u8, sym_name), MachO.StringSliceAdapter{
+            const n_strx = context.macho_file.strtab_dir.getKeyAdapted(@as([]const u8, sym_name), MachO.StringSliceAdapter{
                 .strtab = &context.macho_file.strtab,
             }) orelse unreachable;
             const resolv = context.macho_file.symbol_resolver.get(n_strx) orelse unreachable;
@@ -717,7 +717,7 @@ pub fn parseRelocs(self: *TextBlock, relocs: []macho.relocation_info, context: R
                 const where_index = context.object.symbol_mapping.get(rel.r_symbolnum) orelse unreachable;
                 subtractor = where_index;
             } else {
-                const n_strx = context.macho_file.strtab_dir.getAdapted(@as([]const u8, sym_name), MachO.StringSliceAdapter{
+                const n_strx = context.macho_file.strtab_dir.getKeyAdapted(@as([]const u8, sym_name), MachO.StringSliceAdapter{
                     .strtab = &context.macho_file.strtab,
                 }) orelse unreachable;
                 const resolv = context.macho_file.symbol_resolver.get(n_strx) orelse unreachable;
src/link/MachO.zig
@@ -148,7 +148,7 @@ globals_free_list: std.ArrayListUnmanaged(u32) = .{},
 stub_helper_stubs_start_off: ?u64 = null,
 
 strtab: std.ArrayListUnmanaged(u8) = .{},
-strtab_dir: std.HashMapUnmanaged(u32, u32, StringIndexContext, std.hash_map.default_max_load_percentage) = .{},
+strtab_dir: std.HashMapUnmanaged(u32, void, StringIndexContext, std.hash_map.default_max_load_percentage) = .{},
 
 got_entries: std.ArrayListUnmanaged(GotIndirectionKey) = .{},
 got_entries_map: std.AutoHashMapUnmanaged(GotIndirectionKey, u32) = .{},
@@ -938,7 +938,7 @@ fn linkWithZld(self: *MachO, comp: *Compilation) !void {
 
         {
             // Add dyld_stub_binder as the final GOT entry.
-            const n_strx = self.strtab_dir.getAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
+            const n_strx = self.strtab_dir.getKeyAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
                 .strtab = &self.strtab,
             }) orelse unreachable;
             const resolv = self.symbol_resolver.get(n_strx) orelse unreachable;
@@ -1966,7 +1966,7 @@ fn writeStubHelperCommon(self: *MachO) !void {
                 code[9] = 0xff;
                 code[10] = 0x25;
                 {
-                    const n_strx = self.strtab_dir.getAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
+                    const n_strx = self.strtab_dir.getKeyAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
                         .strtab = &self.strtab,
                     }) orelse unreachable;
                     const resolv = self.symbol_resolver.get(n_strx) orelse unreachable;
@@ -2017,7 +2017,7 @@ fn writeStubHelperCommon(self: *MachO) !void {
                 code[10] = 0xbf;
                 code[11] = 0xa9;
                 binder_blk_outer: {
-                    const n_strx = self.strtab_dir.getAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
+                    const n_strx = self.strtab_dir.getKeyAdapted(@as([]const u8, "dyld_stub_binder"), StringSliceAdapter{
                         .strtab = &self.strtab,
                     }) orelse unreachable;
                     const resolv = self.symbol_resolver.get(n_strx) orelse unreachable;
@@ -2435,7 +2435,7 @@ fn resolveSymbols(self: *MachO) !void {
     }
 
     // Fourth pass, handle synthetic symbols and flag any undefined references.
-    if (self.strtab_dir.getAdapted(@as([]const u8, "___dso_handle"), StringSliceAdapter{
+    if (self.strtab_dir.getKeyAdapted(@as([]const u8, "___dso_handle"), StringSliceAdapter{
         .strtab = &self.strtab,
     })) |n_strx| blk: {
         const resolv = self.symbol_resolver.getPtr(n_strx) orelse break :blk;
@@ -2985,7 +2985,7 @@ fn setEntryPoint(self: *MachO) !void {
     // TODO we should respect the -entry flag passed in by the user to set a custom
     // entrypoint. For now, assume default of `_main`.
     const seg = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
-    const n_strx = self.strtab_dir.getAdapted(@as([]const u8, "_main"), StringSliceAdapter{
+    const n_strx = self.strtab_dir.getKeyAdapted(@as([]const u8, "_main"), StringSliceAdapter{
         .strtab = &self.strtab,
     }) orelse {
         log.err("'_main' export not found", .{});
@@ -4616,7 +4616,7 @@ pub fn addExternFn(self: *MachO, name: []const u8) !u32 {
     const sym_name = try std.fmt.allocPrint(self.base.allocator, "_{s}", .{name});
     defer self.base.allocator.free(sym_name);
 
-    if (self.strtab_dir.getAdapted(@as([]const u8, sym_name), StringSliceAdapter{
+    if (self.strtab_dir.getKeyAdapted(@as([]const u8, sym_name), StringSliceAdapter{
         .strtab = &self.strtab,
     })) |n_strx| {
         const resolv = self.symbol_resolver.get(n_strx) orelse unreachable;
@@ -5858,7 +5858,13 @@ pub fn padToIdeal(actual_size: anytype) @TypeOf(actual_size) {
 }
 
 pub fn makeString(self: *MachO, string: []const u8) !u32 {
-    if (self.strtab_dir.getAdapted(@as([]const u8, string), StringSliceAdapter{ .strtab = &self.strtab })) |off| {
+    const gop = try self.strtab_dir.getOrPutContextAdapted(self.base.allocator, @as([]const u8, string), StringSliceAdapter{
+        .strtab = &self.strtab,
+    }, StringIndexContext{
+        .strtab = &self.strtab,
+    });
+    if (gop.found_existing) {
+        const off = gop.key_ptr.*;
         log.debug("reusing string '{s}' at offset 0x{x}", .{ string, off });
         return off;
     }
@@ -5871,9 +5877,7 @@ pub fn makeString(self: *MachO, string: []const u8) !u32 {
     self.strtab.appendSliceAssumeCapacity(string);
     self.strtab.appendAssumeCapacity(0);
 
-    try self.strtab_dir.putContext(self.base.allocator, new_off, new_off, StringIndexContext{
-        .strtab = &self.strtab,
-    });
+    gop.key_ptr.* = new_off;
 
     return new_off;
 }