//! Reference counted, owned, and interned strings

const std = @import("std");
const Allocator = std.mem.Allocator;

bytes: std.ArrayListUnmanaged(u8),
map: std.StringHashMapUnmanaged(void),

// Construct new empty string intern pool
pub const empty: @This() = .{ .bytes = .empty, .map = .empty };

// Free all memory allocated by the intern pool. Additionally, ensure
// all memory allocated by the pool in it's lifetime is also freed.
pub fn deinit(self: *@This(), gpa: Allocator) void {
    self.bytes.deinit(gpa);
    self.map.deinit(gpa);
    self.* = undefined;
}

// Add a string to the intern pool, returning a stable pointer to the string.
// The pointer will last as far as the string is not removed with `remove()`.
pub fn add(self: *@This(), gpa: Allocator, str: []const u8) ![]const u8 {
    if (self.map.getKey(str)) |entry| {
        // Return pre-existing
        return entry;
    }

    // Allocate memory for & copy the string
    const str_start_idx = self.bytes.items.len;
    try self.bytes.appendSlice(gpa, str);
    const entry_str = self.bytes.items[str_start_idx..];
    errdefer self.bytes.items.len = str_start_idx;

    // Add the entry to our map, use the owned string
    try self.map.putNoClobber(gpa, entry_str, void{});

    // Return the stable pointer
    return entry_str;
}

test "Intern Pool" {
    const gpa = std.testing.allocator;
    try generalWorkload(gpa);

    try std.testing.checkAllAllocationFailures(gpa, generalWorkload, .{});
}

fn generalWorkload(gpa: Allocator) !void {
    var pool: @This() = .empty;
    defer pool.deinit(gpa);

    // Inserting elements
    const a = try pool.add(gpa, "x");
    const b = try pool.add(gpa, "y");
    const c = try pool.add(gpa, "z");
    const d = try pool.add(gpa, "z");
    const e = try pool.add(gpa, "y");

    try std.testing.expectEqualSlices(u8, "x", a);
    try std.testing.expectEqualSlices(u8, "y", b);
    try std.testing.expectEqualSlices(u8, "z", c);
    try std.testing.expectEqualSlices(u8, "z", d);
    try std.testing.expectEqualSlices(u8, "y", e);

    try std.testing.expectEqual(b.ptr, e.ptr);
    try std.testing.expectEqual(c.ptr, d.ptr);

    try std.testing.expectEqual(pool.map.size, 3);
    try std.testing.expectEqualSlices(u8, "xyz", pool.bytes.items);
}