master
  1const std = @import("std.zig");
  2const mem = std.mem;
  3
  4/// Static string map optimized for small sets of disparate string keys.
  5/// Works by separating the keys by length at initialization and only checking
  6/// strings of equal length at runtime.
  7pub fn StaticStringMap(comptime V: type) type {
  8    return StaticStringMapWithEql(V, defaultEql);
  9}
 10
 11/// Like `std.mem.eql`, but takes advantage of the fact that the lengths
 12/// of `a` and `b` are known to be equal.
 13pub fn defaultEql(a: []const u8, b: []const u8) bool {
 14    if (a.ptr == b.ptr) return true;
 15    for (a, b) |a_elem, b_elem| {
 16        if (a_elem != b_elem) return false;
 17    }
 18    return true;
 19}
 20
 21/// Like `std.ascii.eqlIgnoreCase` but takes advantage of the fact that
 22/// the lengths of `a` and `b` are known to be equal.
 23pub fn eqlAsciiIgnoreCase(a: []const u8, b: []const u8) bool {
 24    if (a.ptr == b.ptr) return true;
 25    for (a, b) |a_c, b_c| {
 26        if (std.ascii.toLower(a_c) != std.ascii.toLower(b_c)) return false;
 27    }
 28    return true;
 29}
 30
 31/// StaticStringMap, but accepts an equality function (`eql`).
 32/// The `eql` function is only called to determine the equality
 33/// of equal length strings. Any strings that are not equal length
 34/// are never compared using the `eql` function.
 35pub fn StaticStringMapWithEql(
 36    comptime V: type,
 37    comptime eql: fn (a: []const u8, b: []const u8) bool,
 38) type {
 39    return struct {
 40        kvs: *const KVs = &empty_kvs,
 41        len_indexes: [*]const u32 = &empty_len_indexes,
 42        len_indexes_len: u32 = 0,
 43        min_len: u32 = std.math.maxInt(u32),
 44        max_len: u32 = 0,
 45
 46        pub const KV = struct {
 47            key: []const u8,
 48            value: V,
 49        };
 50
 51        const Self = @This();
 52        const KVs = struct {
 53            keys: [*]const []const u8,
 54            values: [*]const V,
 55            len: u32,
 56        };
 57        const empty_kvs = KVs{
 58            .keys = &empty_keys,
 59            .values = &empty_vals,
 60            .len = 0,
 61        };
 62        const empty_len_indexes = [0]u32{};
 63        const empty_keys = [0][]const u8{};
 64        const empty_vals = [0]V{};
 65
 66        /// Returns a map backed by static, comptime allocated memory.
 67        ///
 68        /// `kvs_list` must be either a list of `struct { []const u8, V }`
 69        /// (key-value pair) tuples, or a list of `struct { []const u8 }`
 70        /// (only keys) tuples if `V` is `void`.
 71        pub inline fn initComptime(comptime kvs_list: anytype) Self {
 72            comptime {
 73                var self = Self{};
 74                if (kvs_list.len == 0)
 75                    return self;
 76
 77                // Since the KVs are sorted, a linearly-growing bound will never
 78                // be sufficient for extreme cases. So we grow proportional to
 79                // N*log2(N).
 80                @setEvalBranchQuota(10 * kvs_list.len * std.math.log2_int_ceil(usize, kvs_list.len));
 81
 82                var sorted_keys: [kvs_list.len][]const u8 = undefined;
 83                var sorted_vals: [kvs_list.len]V = undefined;
 84
 85                self.initSortedKVs(kvs_list, &sorted_keys, &sorted_vals);
 86                const final_keys = sorted_keys;
 87                const final_vals = sorted_vals;
 88                self.kvs = &.{
 89                    .keys = &final_keys,
 90                    .values = &final_vals,
 91                    .len = @intCast(kvs_list.len),
 92                };
 93
 94                var len_indexes: [self.max_len + 1]u32 = undefined;
 95                self.initLenIndexes(&len_indexes);
 96                const final_len_indexes = len_indexes;
 97                self.len_indexes = &final_len_indexes;
 98                self.len_indexes_len = @intCast(len_indexes.len);
 99                return self;
100            }
101        }
102
103        /// Returns a map backed by memory allocated with `allocator`.
104        ///
105        /// Handles `kvs_list` the same way as `initComptime()`.
106        pub fn init(kvs_list: anytype, allocator: mem.Allocator) !Self {
107            var self = Self{};
108            if (kvs_list.len == 0)
109                return self;
110
111            const sorted_keys = try allocator.alloc([]const u8, kvs_list.len);
112            errdefer allocator.free(sorted_keys);
113            const sorted_vals = try allocator.alloc(V, kvs_list.len);
114            errdefer allocator.free(sorted_vals);
115            const kvs = try allocator.create(KVs);
116            errdefer allocator.destroy(kvs);
117
118            self.initSortedKVs(kvs_list, sorted_keys, sorted_vals);
119            kvs.* = .{
120                .keys = sorted_keys.ptr,
121                .values = sorted_vals.ptr,
122                .len = @intCast(kvs_list.len),
123            };
124            self.kvs = kvs;
125
126            const len_indexes = try allocator.alloc(u32, self.max_len + 1);
127            self.initLenIndexes(len_indexes);
128            self.len_indexes = len_indexes.ptr;
129            self.len_indexes_len = @intCast(len_indexes.len);
130            return self;
131        }
132
133        /// this method should only be used with init() and not with initComptime().
134        pub fn deinit(self: Self, allocator: mem.Allocator) void {
135            allocator.free(self.len_indexes[0..self.len_indexes_len]);
136            allocator.free(self.kvs.keys[0..self.kvs.len]);
137            allocator.free(self.kvs.values[0..self.kvs.len]);
138            allocator.destroy(self.kvs);
139        }
140
141        const SortContext = struct {
142            keys: [][]const u8,
143            vals: []V,
144
145            pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
146                return ctx.keys[a].len < ctx.keys[b].len;
147            }
148
149            pub fn swap(ctx: @This(), a: usize, b: usize) void {
150                std.mem.swap([]const u8, &ctx.keys[a], &ctx.keys[b]);
151                std.mem.swap(V, &ctx.vals[a], &ctx.vals[b]);
152            }
153        };
154
155        fn initSortedKVs(
156            self: *Self,
157            kvs_list: anytype,
158            sorted_keys: [][]const u8,
159            sorted_vals: []V,
160        ) void {
161            for (kvs_list, 0..) |kv, i| {
162                sorted_keys[i] = kv.@"0";
163                sorted_vals[i] = if (V == void) {} else kv.@"1";
164                self.min_len = @intCast(@min(self.min_len, kv.@"0".len));
165                self.max_len = @intCast(@max(self.max_len, kv.@"0".len));
166            }
167            mem.sortUnstableContext(0, sorted_keys.len, SortContext{
168                .keys = sorted_keys,
169                .vals = sorted_vals,
170            });
171        }
172
173        fn initLenIndexes(self: Self, len_indexes: []u32) void {
174            var len: usize = 0;
175            var i: u32 = 0;
176            while (len <= self.max_len) : (len += 1) {
177                // find the first keyword len == len
178                while (len > self.kvs.keys[i].len) {
179                    i += 1;
180                }
181                len_indexes[len] = i;
182            }
183        }
184
185        /// Checks if the map has a value for the key.
186        pub fn has(self: Self, str: []const u8) bool {
187            return self.get(str) != null;
188        }
189
190        /// Returns the value for the key if any, else null.
191        pub fn get(self: Self, str: []const u8) ?V {
192            if (self.kvs.len == 0)
193                return null;
194
195            return self.kvs.values[self.getIndex(str) orelse return null];
196        }
197
198        pub fn getIndex(self: Self, str: []const u8) ?usize {
199            const kvs = self.kvs.*;
200            if (kvs.len == 0)
201                return null;
202
203            if (str.len < self.min_len or str.len > self.max_len)
204                return null;
205
206            var i = self.len_indexes[str.len];
207            while (true) {
208                const key = kvs.keys[i];
209                if (key.len != str.len)
210                    return null;
211                if (eql(key, str))
212                    return i;
213                i += 1;
214                if (i >= kvs.len)
215                    return null;
216            }
217        }
218
219        /// Returns the key-value pair where key is the longest prefix of `str`
220        /// else null.
221        ///
222        /// This is effectively an O(N) algorithm which loops from `max_len` to
223        /// `min_len` and calls `getIndex()` to check all keys with the given
224        /// len.
225        pub fn getLongestPrefix(self: Self, str: []const u8) ?KV {
226            if (self.kvs.len == 0)
227                return null;
228            const i = self.getLongestPrefixIndex(str) orelse return null;
229            const kvs = self.kvs.*;
230            return .{
231                .key = kvs.keys[i],
232                .value = kvs.values[i],
233            };
234        }
235
236        pub fn getLongestPrefixIndex(self: Self, str: []const u8) ?usize {
237            if (self.kvs.len == 0)
238                return null;
239
240            if (str.len < self.min_len)
241                return null;
242
243            var len = @min(self.max_len, str.len);
244            while (len >= self.min_len) : (len -= 1) {
245                if (self.getIndex(str[0..len])) |i|
246                    return i;
247            }
248            return null;
249        }
250
251        pub fn keys(self: Self) []const []const u8 {
252            const kvs = self.kvs.*;
253            return kvs.keys[0..kvs.len];
254        }
255
256        pub fn values(self: Self) []const V {
257            const kvs = self.kvs.*;
258            return kvs.values[0..kvs.len];
259        }
260    };
261}
262
263const TestEnum = enum { A, B, C, D, E };
264const TestMap = StaticStringMap(TestEnum);
265const TestKV = struct { []const u8, TestEnum };
266const TestMapVoid = StaticStringMap(void);
267const TestKVVoid = struct { []const u8 };
268const TestMapWithEql = StaticStringMapWithEql(TestEnum, eqlAsciiIgnoreCase);
269const testing = std.testing;
270const test_alloc = testing.allocator;
271
272test "list literal of list literals" {
273    const slice: []const TestKV = &.{
274        .{ "these", .D },
275        .{ "have", .A },
276        .{ "nothing", .B },
277        .{ "incommon", .C },
278        .{ "samelen", .E },
279    };
280
281    const map = TestMap.initComptime(slice);
282    try testMap(map);
283    // Default comparison is case sensitive
284    try testing.expect(null == map.get("NOTHING"));
285
286    // runtime init(), deinit()
287    const map_rt = try TestMap.init(slice, test_alloc);
288    defer map_rt.deinit(test_alloc);
289    try testMap(map_rt);
290    // Default comparison is case sensitive
291    try testing.expect(null == map_rt.get("NOTHING"));
292}
293
294test "array of structs" {
295    const slice = [_]TestKV{
296        .{ "these", .D },
297        .{ "have", .A },
298        .{ "nothing", .B },
299        .{ "incommon", .C },
300        .{ "samelen", .E },
301    };
302
303    try testMap(TestMap.initComptime(slice));
304}
305
306test "slice of structs" {
307    const slice = [_]TestKV{
308        .{ "these", .D },
309        .{ "have", .A },
310        .{ "nothing", .B },
311        .{ "incommon", .C },
312        .{ "samelen", .E },
313    };
314
315    try testMap(TestMap.initComptime(slice));
316}
317
318fn testMap(map: anytype) !void {
319    try testing.expectEqual(TestEnum.A, map.get("have").?);
320    try testing.expectEqual(TestEnum.B, map.get("nothing").?);
321    try testing.expect(null == map.get("missing"));
322    try testing.expectEqual(TestEnum.D, map.get("these").?);
323    try testing.expectEqual(TestEnum.E, map.get("samelen").?);
324
325    try testing.expect(!map.has("missing"));
326    try testing.expect(map.has("these"));
327
328    try testing.expect(null == map.get(""));
329    try testing.expect(null == map.get("averylongstringthathasnomatches"));
330}
331
332test "void value type, slice of structs" {
333    const slice = [_]TestKVVoid{
334        .{"these"},
335        .{"have"},
336        .{"nothing"},
337        .{"incommon"},
338        .{"samelen"},
339    };
340    const map = TestMapVoid.initComptime(slice);
341    try testSet(map);
342    // Default comparison is case sensitive
343    try testing.expect(null == map.get("NOTHING"));
344}
345
346test "void value type, list literal of list literals" {
347    const slice = [_]TestKVVoid{
348        .{"these"},
349        .{"have"},
350        .{"nothing"},
351        .{"incommon"},
352        .{"samelen"},
353    };
354
355    try testSet(TestMapVoid.initComptime(slice));
356}
357
358fn testSet(map: TestMapVoid) !void {
359    try testing.expectEqual({}, map.get("have").?);
360    try testing.expectEqual({}, map.get("nothing").?);
361    try testing.expect(null == map.get("missing"));
362    try testing.expectEqual({}, map.get("these").?);
363    try testing.expectEqual({}, map.get("samelen").?);
364
365    try testing.expect(!map.has("missing"));
366    try testing.expect(map.has("these"));
367
368    try testing.expect(null == map.get(""));
369    try testing.expect(null == map.get("averylongstringthathasnomatches"));
370}
371
372fn testStaticStringMapWithEql(map: TestMapWithEql) !void {
373    try testMap(map);
374    try testing.expectEqual(TestEnum.A, map.get("HAVE").?);
375    try testing.expectEqual(TestEnum.E, map.get("SameLen").?);
376    try testing.expect(null == map.get("SameLength"));
377    try testing.expect(map.has("ThESe"));
378}
379
380test "StaticStringMapWithEql" {
381    const slice = [_]TestKV{
382        .{ "these", .D },
383        .{ "have", .A },
384        .{ "nothing", .B },
385        .{ "incommon", .C },
386        .{ "samelen", .E },
387    };
388
389    try testStaticStringMapWithEql(TestMapWithEql.initComptime(slice));
390}
391
392test "empty" {
393    const m1 = StaticStringMap(usize).initComptime(.{});
394    try testing.expect(null == m1.get("anything"));
395
396    const m2 = StaticStringMapWithEql(usize, eqlAsciiIgnoreCase).initComptime(.{});
397    try testing.expect(null == m2.get("anything"));
398
399    const m3 = try StaticStringMap(usize).init(.{}, test_alloc);
400    try testing.expect(null == m3.get("anything"));
401
402    const m4 = try StaticStringMapWithEql(usize, eqlAsciiIgnoreCase).init(.{}, test_alloc);
403    try testing.expect(null == m4.get("anything"));
404}
405
406test "redundant entries" {
407    const slice = [_]TestKV{
408        .{ "redundant", .D },
409        .{ "theNeedle", .A },
410        .{ "redundant", .B },
411        .{ "re" ++ "dundant", .C },
412        .{ "redun" ++ "dant", .E },
413    };
414    const map = TestMap.initComptime(slice);
415
416    // No promises about which one you get:
417    try testing.expect(null != map.get("redundant"));
418
419    // Default map is not case sensitive:
420    try testing.expect(null == map.get("REDUNDANT"));
421
422    try testing.expectEqual(TestEnum.A, map.get("theNeedle").?);
423}
424
425test "redundant insensitive" {
426    const slice = [_]TestKV{
427        .{ "redundant", .D },
428        .{ "theNeedle", .A },
429        .{ "redundanT", .B },
430        .{ "RE" ++ "dundant", .C },
431        .{ "redun" ++ "DANT", .E },
432    };
433
434    const map = TestMapWithEql.initComptime(slice);
435
436    // No promises about which result you'll get ...
437    try testing.expect(null != map.get("REDUNDANT"));
438    try testing.expect(null != map.get("ReDuNdAnT"));
439    try testing.expectEqual(TestEnum.A, map.get("theNeedle").?);
440}
441
442test "comptime-only value" {
443    const map = StaticStringMap(type).initComptime(.{
444        .{ "a", struct {
445            pub const foo = 1;
446        } },
447        .{ "b", struct {
448            pub const foo = 2;
449        } },
450        .{ "c", struct {
451            pub const foo = 3;
452        } },
453    });
454
455    try testing.expect(map.get("a").?.foo == 1);
456    try testing.expect(map.get("b").?.foo == 2);
457    try testing.expect(map.get("c").?.foo == 3);
458    try testing.expect(map.get("d") == null);
459}
460
461test "getLongestPrefix" {
462    const slice = [_]TestKV{
463        .{ "a", .A },
464        .{ "aa", .B },
465        .{ "aaa", .C },
466        .{ "aaaa", .D },
467    };
468
469    const map = TestMap.initComptime(slice);
470
471    try testing.expectEqual(null, map.getLongestPrefix(""));
472    try testing.expectEqual(null, map.getLongestPrefix("bar"));
473    try testing.expectEqualStrings("aaaa", map.getLongestPrefix("aaaabar").?.key);
474    try testing.expectEqualStrings("aaa", map.getLongestPrefix("aaabar").?.key);
475}
476
477test "getLongestPrefix2" {
478    const slice = [_]struct { []const u8, u8 }{
479        .{ "one", 1 },
480        .{ "two", 2 },
481        .{ "three", 3 },
482        .{ "four", 4 },
483        .{ "five", 5 },
484        .{ "six", 6 },
485        .{ "seven", 7 },
486        .{ "eight", 8 },
487        .{ "nine", 9 },
488    };
489    const map = StaticStringMap(u8).initComptime(slice);
490
491    try testing.expectEqual(1, map.get("one"));
492    try testing.expectEqual(null, map.get("o"));
493    try testing.expectEqual(null, map.get("onexxx"));
494    try testing.expectEqual(9, map.get("nine"));
495    try testing.expectEqual(null, map.get("n"));
496    try testing.expectEqual(null, map.get("ninexxx"));
497    try testing.expectEqual(null, map.get("xxx"));
498
499    try testing.expectEqual(1, map.getLongestPrefix("one").?.value);
500    try testing.expectEqual(1, map.getLongestPrefix("onexxx").?.value);
501    try testing.expectEqual(null, map.getLongestPrefix("o"));
502    try testing.expectEqual(null, map.getLongestPrefix("on"));
503    try testing.expectEqual(9, map.getLongestPrefix("nine").?.value);
504    try testing.expectEqual(9, map.getLongestPrefix("ninexxx").?.value);
505    try testing.expectEqual(null, map.getLongestPrefix("n"));
506    try testing.expectEqual(null, map.getLongestPrefix("xxx"));
507}
508
509test "sorting kvs doesn't exceed eval branch quota" {
510    // from https://github.com/ziglang/zig/issues/19803
511    const TypeToByteSizeLUT = std.StaticStringMap(u32).initComptime(.{
512        .{ "bool", 0 },
513        .{ "c_int", 0 },
514        .{ "c_long", 0 },
515        .{ "c_longdouble", 0 },
516        .{ "t20", 0 },
517        .{ "t19", 0 },
518        .{ "t18", 0 },
519        .{ "t17", 0 },
520        .{ "t16", 0 },
521        .{ "t15", 0 },
522        .{ "t14", 0 },
523        .{ "t13", 0 },
524        .{ "t12", 0 },
525        .{ "t11", 0 },
526        .{ "t10", 0 },
527        .{ "t9", 0 },
528        .{ "t8", 0 },
529        .{ "t7", 0 },
530        .{ "t6", 0 },
531        .{ "t5", 0 },
532        .{ "t4", 0 },
533        .{ "t3", 0 },
534        .{ "t2", 0 },
535        .{ "t1", 1 },
536    });
537    try testing.expectEqual(1, TypeToByteSizeLUT.get("t1"));
538}