154 lines
5.0 KiB
Zig
154 lines
5.0 KiB
Zig
const std = @import("std");
|
|
const Allocator = std.mem.Allocator;
|
|
const PriorityDequeue = std.PriorityDequeue;
|
|
const StringArrayHashMap = std.StringArrayHashMap;
|
|
const StringHashMap = std.StringHashMap;
|
|
const BoundedArray = std.BoundedArray;
|
|
const testing = std.testing;
|
|
|
|
// ShellIndex is an index to the shell strings. As shell can be up to 64 bytes
|
|
// (1<<6), maximum number of shells is 63 (1<<6-1), the maximum location offset
|
|
// is 1<<12. To make location resolvable in 10 bits, all shells will be padded
|
|
// to 4 bytes.
|
|
const ShellIndex = struct {
|
|
offset: u10,
|
|
len: u6,
|
|
};
|
|
|
|
// MaxShells is the maximum number of "popular" shells.
|
|
const MaxShells = 63;
|
|
|
|
// ShellPopcon is a shell popularity contest: collect shells and return the
|
|
// popular ones, sorted by score. score := len(shell) * number_of_shells.
|
|
// String values are copied, the returned slice of shells is allocated
|
|
// using an allocator.
|
|
const ShellPopcon = struct {
|
|
counts: std.StringHashMap(u32),
|
|
allocator: Allocator,
|
|
const Self = @This();
|
|
const KV = struct {
|
|
shell: []const u8,
|
|
score: u32,
|
|
};
|
|
|
|
const ShellSections = struct {
|
|
index: []ShellIndex,
|
|
blob: []const u8,
|
|
|
|
offsets: StringHashMap(u10),
|
|
|
|
pub fn getOffset(self: *ShellSections, shell: []const u8) ?u10 {
|
|
return self.offsets.get(shell);
|
|
}
|
|
|
|
// initializes ShellSections. All strings are copied, nothing is owned.
|
|
pub fn init(allocator: Allocator, shells: BoundedArray([]const u8, MaxShells)) ShellSections {
|
|
self.offsets = StringHashMap(u10).init(allocator);
|
|
_ = allocator;
|
|
_ = shells;
|
|
}
|
|
};
|
|
|
|
pub fn init(allocator: Allocator) Self {
|
|
return Self{
|
|
.counts = std.StringHashMap(u32).init(allocator),
|
|
.allocator = allocator,
|
|
};
|
|
}
|
|
|
|
pub fn deinit(self: *Self) void {
|
|
var it = self.counts.keyIterator();
|
|
while (it.next()) |key_ptr| {
|
|
self.counts.allocator.free(key_ptr.*);
|
|
}
|
|
self.counts.deinit();
|
|
self.* = undefined;
|
|
}
|
|
|
|
pub fn put(self: *Self, shell: []const u8) !void {
|
|
// TODO getOrPutAdapted may be more elegant, not sure which
|
|
// context to pass.
|
|
if (self.counts.getPtr(shell)) |ptr| {
|
|
ptr.* += 1;
|
|
} else {
|
|
var ourShell = try self.allocator.alloc(u8, shell.len);
|
|
std.mem.copy(u8, ourShell, shell);
|
|
try self.counts.put(ourShell, 1);
|
|
}
|
|
}
|
|
|
|
fn cmpShells(context: void, a: KV, b: KV) std.math.Order {
|
|
_ = context;
|
|
return std.math.order(a.score, b.score);
|
|
}
|
|
|
|
pub fn getSections(self: *Self, limit: u32) ShellSections {
|
|
const stderr = std.io.getStdErr().writer();
|
|
_ = stderr;
|
|
|
|
var deque = PriorityDequeue(KV, void, cmpShells).init(self.allocator, {});
|
|
defer deque.deinit();
|
|
|
|
var it = self.counts.iterator();
|
|
while (it.next()) |entry| {
|
|
if (entry.value_ptr.* == 1) {
|
|
continue;
|
|
}
|
|
const score = @truncate(u32, entry.key_ptr.*.len) * entry.value_ptr.*;
|
|
try deque.add(KV{ .shell = entry.key_ptr.*, .score = score });
|
|
}
|
|
|
|
const total = std.math.min(deque.count(), limit);
|
|
var strSlice = self.allocator.alloc([]u8, total);
|
|
defer strSlice.deinit();
|
|
|
|
var i: u32 = 0;
|
|
while (i < total) {
|
|
strSlice[i] = deque.removeMax();
|
|
i += 1;
|
|
}
|
|
|
|
return ShellSections.init(self.allocator, strSlice);
|
|
}
|
|
};
|
|
|
|
test "[]u8 comparison" {
|
|
var s1: []const u8 = "/bin/bash";
|
|
var s2: []const u8 = "/bin/bash";
|
|
try testing.expectEqual(s1, s2);
|
|
}
|
|
|
|
test "basic shellpop" {
|
|
var popcon = ShellPopcon.init(testing.allocator);
|
|
defer popcon.deinit();
|
|
|
|
try popcon.put("/bin/bash");
|
|
try popcon.put("/bin/bash");
|
|
try popcon.put("/bin/bash");
|
|
try popcon.put("/bin/zsh");
|
|
try popcon.put("/bin/zsh");
|
|
try popcon.put("/bin/zsh");
|
|
try popcon.put("/bin/zsh");
|
|
try popcon.put("/bin/nobody");
|
|
try popcon.put("/bin/very-long-shell-name-ought-to-be-first");
|
|
try popcon.put("/bin/very-long-shell-name-ought-to-be-first");
|
|
|
|
const stderr = std.io.getStdErr().writer();
|
|
|
|
var topshells = try popcon.top(2);
|
|
defer topshells.deinit();
|
|
var shellStrings = topshells.keys();
|
|
try testing.expectEqual(shellStrings.len, 2);
|
|
|
|
try stderr.print("\n", .{});
|
|
try stderr.print("0th type: {s}\n", .{@typeName(@TypeOf(shellStrings[0]))});
|
|
try stderr.print("1st type: {s}\n", .{@typeName(@TypeOf(shellStrings[1]))});
|
|
try stderr.print("0th: {s}, len: {d}\n", .{ shellStrings[0], shellStrings[0].len });
|
|
try stderr.print("0ww: /bin/very-long-shell-name-ought-to-be-first\n", .{});
|
|
try stderr.print("1st: {s}, len: {d}\n", .{ shellStrings[1], shellStrings[1].len });
|
|
try stderr.print("1ww: /bin/zsh\n", .{});
|
|
|
|
try testing.expectEqual(shellStrings[0], "/bin/very-long-shell-name-ought-to-be-first");
|
|
try testing.expectEqual(shellStrings[1], "/bin/zsh");
|
|
}
|