import "@std/mem"
import "@std/arraylist"
import "@std/hashmap"

InternError :: enum { out_of_space }

StringId :: alias u32

no_str :: maxval!(StringId)

#! global pool owned by main.hon; initialize before scanning and release with deinit.
strings StringPool := undefined

#! retains its allocator and owns interned copies; release all storage with deinit.
StringPool :: struct {
    ids hashmap.StringHashMap(StringId) # string → id
    strings arraylist.ArrayList([]u8)   # id → owned string
    allocator mem.Allocator
}

#! retains `allocator` for interned copies and indexes; the caller owns the pool.
init proc(allocator mem.Allocator) StringPool {
    return StringPool{
        ids = hashmap.init(allocator),
        strings = arraylist.init(allocator),
        allocator = allocator,
    }
}

#! releases interned strings and indexes with the retained allocator.
deinit proc(pool @mut StringPool) void {
    hashmap.deinit(&pool.ids)
    for (pool.strings.items) |str| mem.free(pool.allocator, str)
    arraylist.deinit(&pool.strings)
}

#! copies new strings into pool-owned storage; IDs and borrowed strings last until deinit.
intern proc(pool @mut StringPool, str []u8) StringId ! (mem.AllocError | InternError) {
    if hashmap.get(&pool.ids, str) |id| return id

    if (pool.strings.items.len >= usize(no_str)) return .out_of_space
    id :: StringId(pool.strings.items.len)

    owned_str []mut u8 :: try mem.alloc(u8, pool.allocator, str.len)
    errdefer mem.free(pool.allocator, owned_str)
    memcopy!(owned_str, str)

    try arraylist.append(&pool.strings, owned_str)
    errdefer _ = arraylist.pop(&pool.strings)

    hashmap.put(&pool.ids, owned_str, id) catch |err| {
        match err {
            .key_exists: unreachable
            else: return err
        }
    }

    return id
}

#! returns pool-owned text borrowed until deinit; the caller must not free it.
get_str proc(pool @StringPool, id StringId) ?[]u8 {
    if (usize(id) >= pool.strings.items.len) return null
    return pool.strings.items[usize(id)]
}

get_id proc(pool @StringPool, str []u8) ?StringId {
    return hashmap.get(&pool.ids, str)
}