1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
const std = @import("std");
const mem = std.mem;
pub const Color = enum { red, blue };
pub const Node = struct {
children: [256]?*Node = [_]?*Node{null} ** 256,
value: ?Color = null,
};
pub const Trie = struct {
allocator: mem.Allocator,
root: Node,
node_count: usize = 0,
fn asBytes(cp: u24) []const u8 {
const bytes: [3]u8 = @bitCast(cp);
return if (bytes[0] < 128)
bytes[0..1]
else if (bytes[1] == 0)
bytes[0..1]
else if (bytes[2] == 0)
bytes[0..2]
else
bytes[0..];
}
pub fn put(self: *Trie, cp: u24, v: Color) !void {
const s = asBytes(cp);
var current: *Node = &self.root;
for (s, 0..) |b, i| {
if (current.children[b]) |n| {
if (i == s.len - 1) {
n.value = v;
return;
}
current = n;
continue;
}
self.node_count += 1;
const new_node = try self.allocator.create(Node);
new_node.* = .{ .value = if (i == s.len - 1) v else null };
current.children[b] = new_node;
current = new_node;
}
}
pub fn get(self: Trie, cp: u24) ?Color {
const s = asBytes(cp);
var current = &self.root;
return for (s, 0..) |b, i| {
if (current.children[b]) |n| {
if (i == s.len - 1) break if (n.value) |v| v else null;
current = n;
}
} else null;
}
};
test "Trie works" {
var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
defer arena.deinit();
const allocator = arena.allocator();
var trie = Trie{ .allocator = allocator, .root = .{} };
const cp_1: u21 = '\u{10ffff}';
const cp_2: u21 = '\u{10ff}';
const cp_3: u21 = '\u{10}';
try trie.put(cp_1, .red);
try trie.put(cp_3, .blue);
try std.testing.expectEqual(@as(?Color, .red), trie.get(cp_1));
try std.testing.expectEqual(@as(?Color, null), trie.get(cp_2));
try std.testing.expectEqual(@as(?Color, .blue), trie.get(cp_3));
}
|
