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}