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 static, comptime allocated memory.
104 ///
105 /// `V` must be an enum. The enum's tag names will be used as the keys.
106 pub inline fn initEnum() Self {
107 comptime {
108 var self: Self = .{};
109
110 const field_names = @typeInfo(V).@"enum".field_names;
111 if (field_names.len == 0) return self;
112
113 // Since the KVs are sorted, a linearly-growing bound will never
114 // be sufficient for extreme cases. So we grow proportional to
115 // N*log2(N).
116 @setEvalBranchQuota(10 * field_names.len * std.math.log2_int_ceil(usize, field_names.len));
117
118 var sorted_keys: [field_names.len][]const u8 = field_names[0..field_names.len].*;
119 var sorted_vals: [field_names.len]V = undefined;
120 for (&sorted_vals, @typeInfo(V).@"enum".field_values) |*x, i| x.* = @fromBackingInt(@intCast(i));
121
122 for (field_names) |field_name| {
123 self.min_len = @min(self.min_len, field_name.len);
124 self.max_len = @max(self.max_len, field_name.len);
125 }
126
127 mem.sortUnstableContext(0, sorted_keys.len, SortContext{
128 .keys = &sorted_keys,
129 .vals = &sorted_vals,
130 });
131
132 const final_keys = sorted_keys;
133 const final_vals = sorted_vals;
134 self.kvs = &.{
135 .keys = &final_keys,
136 .values = &final_vals,
137 .len = @intCast(field_names.len),
138 };
139
140 var len_indexes: [self.max_len + 1]u32 = undefined;
141 self.initLenIndexes(&len_indexes);
142 const final_len_indexes = len_indexes;
143 self.len_indexes = &final_len_indexes;
144 self.len_indexes_len = @intCast(len_indexes.len);
145 return self;
146 }
147 }
148
149 /// Returns a map backed by memory allocated with `allocator`.
150 ///
151 /// Handles `kvs_list` the same way as `initComptime()`.
152 pub fn init(kvs_list: anytype, allocator: mem.Allocator) !Self {
153 var self = Self{};
154 if (kvs_list.len == 0)
155 return self;
156
157 const sorted_keys = try allocator.alloc([]const u8, kvs_list.len);
158 errdefer allocator.free(sorted_keys);
159 const sorted_vals = try allocator.alloc(V, kvs_list.len);
160 errdefer allocator.free(sorted_vals);
161 const kvs = try allocator.create(KVs);
162 errdefer allocator.destroy(kvs);
163
164 self.initSortedKVs(kvs_list, sorted_keys, sorted_vals);
165 kvs.* = .{
166 .keys = sorted_keys.ptr,
167 .values = sorted_vals.ptr,
168 .len = @intCast(kvs_list.len),
169 };
170 self.kvs = kvs;
171
172 const len_indexes = try allocator.alloc(u32, self.max_len + 1);
173 self.initLenIndexes(len_indexes);
174 self.len_indexes = len_indexes.ptr;
175 self.len_indexes_len = @intCast(len_indexes.len);
176 return self;
177 }
178
179 /// this method should only be used with init() and not with initComptime().
180 pub fn deinit(self: Self, allocator: mem.Allocator) void {
181 allocator.free(self.len_indexes[0..self.len_indexes_len]);
182 allocator.free(self.kvs.keys[0..self.kvs.len]);
183 allocator.free(self.kvs.values[0..self.kvs.len]);
184 allocator.destroy(self.kvs);
185 }
186
187 const SortContext = struct {
188 keys: [][]const u8,
189 vals: []V,
190
191 pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
192 return ctx.keys[a].len < ctx.keys[b].len;
193 }
194
195 pub fn swap(ctx: @This(), a: usize, b: usize) void {
196 std.mem.swap([]const u8, &ctx.keys[a], &ctx.keys[b]);
197 std.mem.swap(V, &ctx.vals[a], &ctx.vals[b]);
198 }
199 };
200
201 fn initSortedKVs(
202 self: *Self,
203 kvs_list: anytype,
204 sorted_keys: [][]const u8,
205 sorted_vals: []V,
206 ) void {
207 for (kvs_list, 0..) |kv, i| {
208 sorted_keys[i] = kv.@"0";
209 sorted_vals[i] = if (V == void) {} else kv.@"1";
210 self.min_len = @intCast(@min(self.min_len, kv.@"0".len));
211 self.max_len = @intCast(@max(self.max_len, kv.@"0".len));
212 }
213 mem.sortUnstableContext(0, sorted_keys.len, SortContext{
214 .keys = sorted_keys,
215 .vals = sorted_vals,
216 });
217 }
218
219 fn initLenIndexes(self: Self, len_indexes: []u32) void {
220 var len: usize = 0;
221 var i: u32 = 0;
222 while (len <= self.max_len) : (len += 1) {
223 // find the first keyword len == len
224 while (len > self.kvs.keys[i].len) {
225 i += 1;
226 }
227 len_indexes[len] = i;
228 }
229 }
230
231 /// Checks if the map has a value for the key.
232 pub fn has(self: Self, str: []const u8) bool {
233 return self.get(str) != null;
234 }
235
236 /// Returns the value for the key if any, else null.
237 pub fn get(self: Self, str: []const u8) ?V {
238 if (self.kvs.len == 0)
239 return null;
240
241 return self.kvs.values[self.getIndex(str) orelse return null];
242 }
243
244 /// Returns the index corresponding to the `str` within the
245 /// generated `kvs`, or `null` if `str` was not found.
246 /// The returned index is unrelated to the input `kvs_list`.
247 pub fn getIndex(self: Self, str: []const u8) ?usize {
248 const kvs = self.kvs.*;
249 if (kvs.len == 0)
250 return null;
251
252 if (str.len < self.min_len or str.len > self.max_len)
253 return null;
254
255 var i = self.len_indexes[str.len];
256 while (true) {
257 const key = kvs.keys[i];
258 if (key.len != str.len)
259 return null;
260 if (eql(key, str))
261 return i;
262 i += 1;
263 if (i >= kvs.len)
264 return null;
265 }
266 }
267
268 /// Returns the key-value pair where key is the longest prefix of `str`
269 /// else null.
270 ///
271 /// This is effectively an O(N) algorithm which loops from `max_len` to
272 /// `min_len` and calls `getIndex()` to check all keys with the given
273 /// len.
274 pub fn getLongestPrefix(self: Self, str: []const u8) ?KV {
275 if (self.kvs.len == 0)
276 return null;
277 const i = self.getLongestPrefixIndex(str) orelse return null;
278 const kvs = self.kvs.*;
279 return .{
280 .key = kvs.keys[i],
281 .value = kvs.values[i],
282 };
283 }
284
285 /// Returns the index within the generated `kvs` corresponding to the
286 /// key-value pair where key is the longest prefix of `str`, or `null`
287 /// if no such key-value pair was found. The returned index is unrelated
288 /// to the input `kvs_list`.
289 ///
290 /// This is effectively an O(N) algorithm which loops from `max_len` to
291 /// `min_len` and calls `getIndex()` to check all keys with the given
292 /// len.
293 pub fn getLongestPrefixIndex(self: Self, str: []const u8) ?usize {
294 if (self.kvs.len == 0)
295 return null;
296
297 if (str.len < self.min_len)
298 return null;
299
300 var len = @min(self.max_len, str.len);
301 while (len >= self.min_len) : (len -= 1) {
302 if (self.getIndex(str[0..len])) |i|
303 return i;
304 }
305 return null;
306 }
307
308 /// Returns the slice of keys from the generated `kvs`, which may
309 /// be in a different order than the input `kvs_list`.
310 pub fn keys(self: Self) []const []const u8 {
311 const kvs = self.kvs.*;
312 return kvs.keys[0..kvs.len];
313 }
314
315 /// Returns the slice of values from the generated `kvs`, which may
316 /// be in a different order than the input `kvs_list`.
317 pub fn values(self: Self) []const V {
318 const kvs = self.kvs.*;
319 return kvs.values[0..kvs.len];
320 }
321 };
322}
323
324const TestEnum = enum { A, B, C, D, E };
325const TestMap = StaticStringMap(TestEnum);
326const TestKV = struct { []const u8, TestEnum };
327const TestMapVoid = StaticStringMap(void);
328const TestKVVoid = struct { []const u8 };
329const TestMapWithEql = StaticStringMapWithEql(TestEnum, eqlAsciiIgnoreCase);
330const testing = std.testing;
331const test_alloc = testing.allocator;
332
333test "list literal of list literals" {
334 const slice: []const TestKV = &.{
335 .{ "these", .D },
336 .{ "have", .A },
337 .{ "nothing", .B },
338 .{ "incommon", .C },
339 .{ "samelen", .E },
340 };
341
342 const map = TestMap.initComptime(slice);
343 try testMap(map);
344 // Default comparison is case sensitive
345 try testing.expect(null == map.get("NOTHING"));
346
347 // runtime init(), deinit()
348 const map_rt = try TestMap.init(slice, test_alloc);
349 defer map_rt.deinit(test_alloc);
350 try testMap(map_rt);
351 // Default comparison is case sensitive
352 try testing.expect(null == map_rt.get("NOTHING"));
353}
354
355test "array of structs" {
356 const slice = [_]TestKV{
357 .{ "these", .D },
358 .{ "have", .A },
359 .{ "nothing", .B },
360 .{ "incommon", .C },
361 .{ "samelen", .E },
362 };
363
364 try testMap(TestMap.initComptime(slice));
365}
366
367test "slice of structs" {
368 const slice = [_]TestKV{
369 .{ "these", .D },
370 .{ "have", .A },
371 .{ "nothing", .B },
372 .{ "incommon", .C },
373 .{ "samelen", .E },
374 };
375
376 try testMap(TestMap.initComptime(slice));
377}
378
379fn testMap(map: anytype) !void {
380 try testing.expectEqual(TestEnum.A, map.get("have").?);
381 try testing.expectEqual(TestEnum.B, map.get("nothing").?);
382 try testing.expect(null == map.get("missing"));
383 try testing.expectEqual(TestEnum.D, map.get("these").?);
384 try testing.expectEqual(TestEnum.E, map.get("samelen").?);
385
386 try testing.expect(!map.has("missing"));
387 try testing.expect(map.has("these"));
388
389 try testing.expect(null == map.get(""));
390 try testing.expect(null == map.get("averylongstringthathasnomatches"));
391}
392
393test "void value type, slice of structs" {
394 const slice = [_]TestKVVoid{
395 .{"these"},
396 .{"have"},
397 .{"nothing"},
398 .{"incommon"},
399 .{"samelen"},
400 };
401 const map = TestMapVoid.initComptime(slice);
402 try testSet(map);
403 // Default comparison is case sensitive
404 try testing.expect(null == map.get("NOTHING"));
405}
406
407test "void value type, list literal of list literals" {
408 const slice = [_]TestKVVoid{
409 .{"these"},
410 .{"have"},
411 .{"nothing"},
412 .{"incommon"},
413 .{"samelen"},
414 };
415
416 try testSet(TestMapVoid.initComptime(slice));
417}
418
419fn testSet(map: TestMapVoid) !void {
420 try testing.expectEqual({}, map.get("have").?);
421 try testing.expectEqual({}, map.get("nothing").?);
422 try testing.expect(null == map.get("missing"));
423 try testing.expectEqual({}, map.get("these").?);
424 try testing.expectEqual({}, map.get("samelen").?);
425
426 try testing.expect(!map.has("missing"));
427 try testing.expect(map.has("these"));
428
429 try testing.expect(null == map.get(""));
430 try testing.expect(null == map.get("averylongstringthathasnomatches"));
431}
432
433fn testStaticStringMapWithEql(map: TestMapWithEql) !void {
434 try testMap(map);
435 try testing.expectEqual(TestEnum.A, map.get("HAVE").?);
436 try testing.expectEqual(TestEnum.E, map.get("SameLen").?);
437 try testing.expect(null == map.get("SameLength"));
438 try testing.expect(map.has("ThESe"));
439}
440
441test "StaticStringMapWithEql" {
442 const slice = [_]TestKV{
443 .{ "these", .D },
444 .{ "have", .A },
445 .{ "nothing", .B },
446 .{ "incommon", .C },
447 .{ "samelen", .E },
448 };
449
450 try testStaticStringMapWithEql(TestMapWithEql.initComptime(slice));
451}
452
453test "empty" {
454 const m1 = StaticStringMap(usize).initComptime(.{});
455 try testing.expect(null == m1.get("anything"));
456
457 const m2 = StaticStringMapWithEql(usize, eqlAsciiIgnoreCase).initComptime(.{});
458 try testing.expect(null == m2.get("anything"));
459
460 const m3 = try StaticStringMap(usize).init(.{}, test_alloc);
461 try testing.expect(null == m3.get("anything"));
462
463 const m4 = try StaticStringMapWithEql(usize, eqlAsciiIgnoreCase).init(.{}, test_alloc);
464 try testing.expect(null == m4.get("anything"));
465}
466
467test "redundant entries" {
468 const slice = [_]TestKV{
469 .{ "redundant", .D },
470 .{ "theNeedle", .A },
471 .{ "redundant", .B },
472 .{ "re" ++ "dundant", .C },
473 .{ "redun" ++ "dant", .E },
474 };
475 const map = TestMap.initComptime(slice);
476
477 // No promises about which one you get:
478 try testing.expect(null != map.get("redundant"));
479
480 // Default map is not case sensitive:
481 try testing.expect(null == map.get("REDUNDANT"));
482
483 try testing.expectEqual(TestEnum.A, map.get("theNeedle").?);
484}
485
486test "redundant insensitive" {
487 const slice = [_]TestKV{
488 .{ "redundant", .D },
489 .{ "theNeedle", .A },
490 .{ "redundanT", .B },
491 .{ "RE" ++ "dundant", .C },
492 .{ "redun" ++ "DANT", .E },
493 };
494
495 const map = TestMapWithEql.initComptime(slice);
496
497 // No promises about which result you'll get ...
498 try testing.expect(null != map.get("REDUNDANT"));
499 try testing.expect(null != map.get("ReDuNdAnT"));
500 try testing.expectEqual(TestEnum.A, map.get("theNeedle").?);
501}
502
503test "comptime-only value" {
504 const map = StaticStringMap(type).initComptime(.{
505 .{ "a", struct {
506 pub const foo = 1;
507 } },
508 .{ "b", struct {
509 pub const foo = 2;
510 } },
511 .{ "c", struct {
512 pub const foo = 3;
513 } },
514 });
515
516 try testing.expect(map.get("a").?.foo == 1);
517 try testing.expect(map.get("b").?.foo == 2);
518 try testing.expect(map.get("c").?.foo == 3);
519 try testing.expect(map.get("d") == null);
520}
521
522test "getIndex" {
523 const slice = [_]TestKV{
524 .{ "longer", .A },
525 .{ "short", .B },
526 };
527 const map = TestMap.initComptime(slice);
528
529 // This index can be different than the index of the "short" KV in `slice`
530 const short_index = map.getIndex("short").?;
531 try testing.expectEqualStrings("short", map.keys()[short_index]);
532 try testing.expectEqual(.B, map.values()[short_index]);
533 try testing.expectEqual(null, map.getIndex("missing"));
534}
535
536test "getLongestPrefix" {
537 const slice = [_]TestKV{
538 .{ "a", .A },
539 .{ "aa", .B },
540 .{ "aaa", .C },
541 .{ "aaaa", .D },
542 };
543
544 const map = TestMap.initComptime(slice);
545
546 try testing.expectEqual(null, map.getLongestPrefix(""));
547 try testing.expectEqual(null, map.getLongestPrefix("bar"));
548 try testing.expectEqualStrings("aaaa", map.getLongestPrefix("aaaabar").?.key);
549 try testing.expectEqualStrings("aaa", map.getLongestPrefix("aaabar").?.key);
550}
551
552test "getLongestPrefix2" {
553 const slice = [_]struct { []const u8, u8 }{
554 .{ "one", 1 },
555 .{ "two", 2 },
556 .{ "three", 3 },
557 .{ "four", 4 },
558 .{ "five", 5 },
559 .{ "six", 6 },
560 .{ "seven", 7 },
561 .{ "eight", 8 },
562 .{ "nine", 9 },
563 };
564 const map = StaticStringMap(u8).initComptime(slice);
565
566 try testing.expectEqual(1, map.get("one"));
567 try testing.expectEqual(null, map.get("o"));
568 try testing.expectEqual(null, map.get("onexxx"));
569 try testing.expectEqual(9, map.get("nine"));
570 try testing.expectEqual(null, map.get("n"));
571 try testing.expectEqual(null, map.get("ninexxx"));
572 try testing.expectEqual(null, map.get("xxx"));
573
574 try testing.expectEqual(1, map.getLongestPrefix("one").?.value);
575 try testing.expectEqual(1, map.getLongestPrefix("onexxx").?.value);
576 try testing.expectEqual(null, map.getLongestPrefix("o"));
577 try testing.expectEqual(null, map.getLongestPrefix("on"));
578 try testing.expectEqual(9, map.getLongestPrefix("nine").?.value);
579 try testing.expectEqual(9, map.getLongestPrefix("ninexxx").?.value);
580 try testing.expectEqual(null, map.getLongestPrefix("n"));
581 try testing.expectEqual(null, map.getLongestPrefix("xxx"));
582}
583
584test "sorting kvs doesn't exceed eval branch quota" {
585 // from https://github.com/ziglang/zig/issues/19803
586 const TypeToByteSizeLUT = std.StaticStringMap(u32).initComptime(.{
587 .{ "bool", 0 },
588 .{ "c_int", 0 },
589 .{ "c_long", 0 },
590 .{ "c_longdouble", 0 },
591 .{ "t20", 0 },
592 .{ "t19", 0 },
593 .{ "t18", 0 },
594 .{ "t17", 0 },
595 .{ "t16", 0 },
596 .{ "t15", 0 },
597 .{ "t14", 0 },
598 .{ "t13", 0 },
599 .{ "t12", 0 },
600 .{ "t11", 0 },
601 .{ "t10", 0 },
602 .{ "t9", 0 },
603 .{ "t8", 0 },
604 .{ "t7", 0 },
605 .{ "t6", 0 },
606 .{ "t5", 0 },
607 .{ "t4", 0 },
608 .{ "t3", 0 },
609 .{ "t2", 0 },
610 .{ "t1", 1 },
611 });
612 try testing.expectEqual(1, TypeToByteSizeLUT.get("t1"));
613}
614
615test "initEnum" {
616 const UnsortedEnum = enum { BB, A, CCC, DDD };
617 const map = StaticStringMap(UnsortedEnum).initEnum();
618 try testing.expect(map.has("A"));
619 try testing.expect(!map.has("a"));
620 try testing.expectEqual(.BB, map.get("BB"));
621 try testing.expectEqual(.A, map.get("A"));
622 try testing.expectEqual(.CCC, map.get("CCC"));
623 try testing.expectEqual(.DDD, map.get("DDD"));
624 try testing.expectEqual(null, map.getIndex("F"));
625}