zig

fork of https://codeberg.org/ziglang/zig
Log | Tree | Refs | README | LICENSE

blob 2a178d9d (11086B) - Raw


      1 const std = @import("index.zig");
      2 const debug = std.debug;
      3 const assert = debug.assert;
      4 const math = std.math;
      5 const mem = std.mem;
      6 const Allocator = mem.Allocator;
      7 const builtin = @import("builtin");
      8 
      9 const want_modification_safety = builtin.mode != builtin.Mode.ReleaseFast;
     10 const debug_u32 = if (want_modification_safety) u32 else void;
     11 
     12 pub fn HashMap(comptime K: type, comptime V: type,
     13     comptime hash: fn(key: K)u32,
     14     comptime eql: fn(a: K, b: K)bool) type
     15 {
     16     return struct {
     17         entries: []Entry,
     18         size: usize,
     19         max_distance_from_start_index: usize,
     20         allocator: &Allocator,
     21         // this is used to detect bugs where a hashtable is edited while an iterator is running.
     22         modification_count: debug_u32,
     23 
     24         const Self = this;
     25 
     26         pub const Entry = struct {
     27             used: bool,
     28             distance_from_start_index: usize,
     29             key: K,
     30             value: V,
     31         };
     32 
     33         pub const Iterator = struct {
     34             hm: &const Self,
     35             // how many items have we returned
     36             count: usize,
     37             // iterator through the entry array
     38             index: usize,
     39             // used to detect concurrent modification
     40             initial_modification_count: debug_u32,
     41 
     42             pub fn next(it: &Iterator) ?&Entry {
     43                 if (want_modification_safety) {
     44                     assert(it.initial_modification_count == it.hm.modification_count); // concurrent modification
     45                 }
     46                 if (it.count >= it.hm.size) return null;
     47                 while (it.index < it.hm.entries.len) : (it.index += 1) {
     48                     const entry = &it.hm.entries[it.index];
     49                     if (entry.used) {
     50                         it.index += 1;
     51                         it.count += 1;
     52                         return entry;
     53                     }
     54                 }
     55                 unreachable; // no next item
     56             }
     57 
     58             // Reset the iterator to the initial index
     59             pub fn reset(it: &Iterator) void {
     60                 it.count = 0;
     61                 it.index = 0;
     62                 // Resetting the modification count too
     63                 it.initial_modification_count = it.hm.modification_count;
     64             }
     65         };
     66 
     67         pub fn init(allocator: &Allocator) Self {
     68             return Self {
     69                 .entries = []Entry{},
     70                 .allocator = allocator,
     71                 .size = 0,
     72                 .max_distance_from_start_index = 0,
     73                 .modification_count = if (want_modification_safety) 0 else {},
     74             };
     75         }
     76 
     77         pub fn deinit(hm: &const Self) void {
     78             hm.allocator.free(hm.entries);
     79         }
     80 
     81         pub fn clear(hm: &Self) void {
     82             for (hm.entries) |*entry| {
     83                 entry.used = false;
     84             }
     85             hm.size = 0;
     86             hm.max_distance_from_start_index = 0;
     87             hm.incrementModificationCount();
     88         }
     89 
     90         pub fn count(hm: &const Self) usize {
     91             return hm.size;
     92         }
     93 
     94         /// Returns the value that was already there.
     95         pub fn put(hm: &Self, key: K, value: &const V) !?V {
     96             if (hm.entries.len == 0) {
     97                 try hm.initCapacity(16);
     98             }
     99             hm.incrementModificationCount();
    100 
    101             // if we get too full (60%), double the capacity
    102             if (hm.size * 5 >= hm.entries.len * 3) {
    103                 const old_entries = hm.entries;
    104                 try hm.initCapacity(hm.entries.len * 2);
    105                 // dump all of the old elements into the new table
    106                 for (old_entries) |*old_entry| {
    107                     if (old_entry.used) {
    108                         _ = hm.internalPut(old_entry.key, old_entry.value);
    109                     }
    110                 }
    111                 hm.allocator.free(old_entries);
    112             }
    113 
    114             return hm.internalPut(key, value);
    115         }
    116 
    117         pub fn get(hm: &const Self, key: K) ?&Entry {
    118             if (hm.entries.len == 0) {
    119                 return null;
    120             }
    121             return hm.internalGet(key);
    122         }
    123 
    124         pub fn contains(hm: &const Self, key: K) bool {
    125             return hm.get(key) != null;
    126         }
    127 
    128         pub fn remove(hm: &Self, key: K) ?&Entry {
    129             if (hm.entries.len == 0) return null;
    130             hm.incrementModificationCount();
    131             const start_index = hm.keyToIndex(key);
    132             {var roll_over: usize = 0; while (roll_over <= hm.max_distance_from_start_index) : (roll_over += 1) {
    133                 const index = (start_index + roll_over) % hm.entries.len;
    134                 var entry = &hm.entries[index];
    135 
    136                 if (!entry.used)
    137                     return null;
    138 
    139                 if (!eql(entry.key, key)) continue;
    140 
    141                 while (roll_over < hm.entries.len) : (roll_over += 1) {
    142                     const next_index = (start_index + roll_over + 1) % hm.entries.len;
    143                     const next_entry = &hm.entries[next_index];
    144                     if (!next_entry.used or next_entry.distance_from_start_index == 0) {
    145                         entry.used = false;
    146                         hm.size -= 1;
    147                         return entry;
    148                     }
    149                     *entry = *next_entry;
    150                     entry.distance_from_start_index -= 1;
    151                     entry = next_entry;
    152                 }
    153                 unreachable; // shifting everything in the table
    154             }}
    155             return null;
    156         }
    157 
    158         pub fn iterator(hm: &const Self) Iterator {
    159             return Iterator {
    160                 .hm = hm,
    161                 .count = 0,
    162                 .index = 0,
    163                 .initial_modification_count = hm.modification_count,
    164             };
    165         }
    166 
    167         fn initCapacity(hm: &Self, capacity: usize) !void {
    168             hm.entries = try hm.allocator.alloc(Entry, capacity);
    169             hm.size = 0;
    170             hm.max_distance_from_start_index = 0;
    171             for (hm.entries) |*entry| {
    172                 entry.used = false;
    173             }
    174         }
    175 
    176         fn incrementModificationCount(hm: &Self) void {
    177             if (want_modification_safety) {
    178                 hm.modification_count +%= 1;
    179             }
    180         }
    181 
    182         /// Returns the value that was already there.
    183         fn internalPut(hm: &Self, orig_key: K, orig_value: &const V) ?V {
    184             var key = orig_key;
    185             var value = *orig_value;
    186             const start_index = hm.keyToIndex(key);
    187             var roll_over: usize = 0;
    188             var distance_from_start_index: usize = 0;
    189             while (roll_over < hm.entries.len) : ({roll_over += 1; distance_from_start_index += 1;}) {
    190                 const index = (start_index + roll_over) % hm.entries.len;
    191                 const entry = &hm.entries[index];
    192 
    193                 if (entry.used and !eql(entry.key, key)) {
    194                     if (entry.distance_from_start_index < distance_from_start_index) {
    195                         // robin hood to the rescue
    196                         const tmp = *entry;
    197                         hm.max_distance_from_start_index = math.max(hm.max_distance_from_start_index,
    198                             distance_from_start_index);
    199                         *entry = Entry {
    200                             .used = true,
    201                             .distance_from_start_index = distance_from_start_index,
    202                             .key = key,
    203                             .value = value,
    204                         };
    205                         key = tmp.key;
    206                         value = tmp.value;
    207                         distance_from_start_index = tmp.distance_from_start_index;
    208                     }
    209                     continue;
    210                 }
    211 
    212                 var result: ?V = null;
    213                 if (entry.used) {
    214                     result = entry.value;
    215                 } else {
    216                     // adding an entry. otherwise overwriting old value with
    217                     // same key
    218                     hm.size += 1;
    219                 }
    220 
    221                 hm.max_distance_from_start_index = math.max(distance_from_start_index, hm.max_distance_from_start_index);
    222                 *entry = Entry {
    223                     .used = true,
    224                     .distance_from_start_index = distance_from_start_index,
    225                     .key = key,
    226                     .value = value,
    227                 };
    228                 return result;
    229             }
    230             unreachable; // put into a full map
    231         }
    232 
    233         fn internalGet(hm: &const Self, key: K) ?&Entry {
    234             const start_index = hm.keyToIndex(key);
    235             {var roll_over: usize = 0; while (roll_over <= hm.max_distance_from_start_index) : (roll_over += 1) {
    236                 const index = (start_index + roll_over) % hm.entries.len;
    237                 const entry = &hm.entries[index];
    238 
    239                 if (!entry.used) return null;
    240                 if (eql(entry.key, key)) return entry;
    241             }}
    242             return null;
    243         }
    244 
    245         fn keyToIndex(hm: &const Self, key: K) usize {
    246             return usize(hash(key)) % hm.entries.len;
    247         }
    248     };
    249 }
    250 
    251 test "basic hash map usage" {
    252     var direct_allocator = std.heap.DirectAllocator.init();
    253     defer direct_allocator.deinit();
    254 
    255     var map = HashMap(i32, i32, hash_i32, eql_i32).init(&direct_allocator.allocator);
    256     defer map.deinit();
    257 
    258     assert((map.put(1, 11) catch unreachable) == null);
    259     assert((map.put(2, 22) catch unreachable) == null);
    260     assert((map.put(3, 33) catch unreachable) == null);
    261     assert((map.put(4, 44) catch unreachable) == null);
    262     assert((map.put(5, 55) catch unreachable) == null);
    263 
    264     assert(??(map.put(5, 66) catch unreachable) == 55);
    265     assert(??(map.put(5, 55) catch unreachable) == 66);
    266 
    267     assert(map.contains(2));
    268     assert((??map.get(2)).value == 22);
    269     _ = map.remove(2);
    270     assert(map.remove(2) == null);
    271     assert(map.get(2) == null);
    272 }
    273 
    274 test "iterator hash map" {
    275     var direct_allocator = std.heap.DirectAllocator.init();
    276     defer direct_allocator.deinit();
    277 
    278     var reset_map = HashMap(i32, i32, hash_i32, eql_i32).init(&direct_allocator.allocator);
    279     defer reset_map.deinit();
    280 
    281     assert((reset_map.put(1, 11) catch unreachable) == null);
    282     assert((reset_map.put(2, 22) catch unreachable) == null);
    283     assert((reset_map.put(3, 33) catch unreachable) == null);
    284 
    285     var keys = []i32 { 1, 2, 3 };
    286     var values = []i32 { 11, 22, 33 };
    287 
    288     var it = reset_map.iterator();
    289     var count : usize = 0;
    290     while (it.next()) |next| {
    291         assert(next.key == keys[count]);
    292         assert(next.value == values[count]);
    293         count += 1;
    294     }
    295 
    296     assert(count == 3);
    297     assert(it.next() == null);
    298     it.reset();
    299     count = 0;
    300     while (it.next()) |next| {
    301         assert(next.key == keys[count]);
    302         assert(next.value == values[count]);
    303         count += 1;
    304         if (count == 2) break;
    305     }
    306 
    307     it.reset();
    308     var entry = ?? it.next();
    309     assert(entry.key == keys[0]);
    310     assert(entry.value == values[0]);
    311 }
    312 
    313 fn hash_i32(x: i32) u32 {
    314     return @bitCast(u32, x);
    315 }
    316 
    317 fn eql_i32(a: i32, b: i32) bool {
    318     return a == b;
    319 }