import "@std/mem"
import "@std/meta"

PutError :: enum { key_exists }

Entry proc($K, $V type) type {
    return struct {
        # hash = 0 means empty
        hash usize = 0
        key K
        value V
    }
}

#! owns entry storage and retains its allocator; release storage with deinit.
#! keys and values are stored by value; any resources they reference remain caller-owned.
#! null callbacks select automatic operations; equal keys must always produce equal hashes.
HashMap proc(
    $K, $V type,
    $hash_key ?proc(key K) usize,
    $keys_eql ?proc(a, b K) bool,
) type {
    return struct {
        entries []mut Entry(K, V)
        count usize
        allocator mem.Allocator
    }
}

StringHashMap proc($V type) type {
    return HashMap([]u8, V, str_hash, str_eql)
}

#! automatic value hashing for integers, booleans, enums, distinct types, arrays and records.
#! slices and floating-point keys require explicit hash/equality operations through HashMap.
#! record padding is ignored; pointer, union and optional keys are not supported.
AutoHashMap proc($K, $V type) type {
    return HashMap(K, V, null, null)
}

#! retain `allocator` for later growth and cleanup; this call does not allocate.
init proc(
    $K, $V type,
    $hash_key ?proc(key K) usize,
    $keys_eql ?proc(a, b K) bool,
    allocator mem.Allocator,
) HashMap(K, V, hash_key, keys_eql) {
    return HashMap(K, V, hash_key, keys_eql){
        entries = &[],
        count = 0,
        allocator = allocator,
    }
}

#! free owned entry storage using the retained allocator, not resources referenced by keys or values.
#! note: this operation invalidates the map.
deinit proc(
    $K, $V type,
    $hash_key ?proc(key K) usize,
    $keys_eql ?proc(a, b K) bool,
    map @HashMap(K, V, hash_key, keys_eql),
) void { mem.free(map.allocator, map.entries) }

get proc(
    $K, $V type,
    $hash_key ?proc(key K) usize,
    $keys_eql ?proc(a, b K) bool,
    map @HashMap(K, V, hash_key, keys_eql),
    key K,
) ?V {
    if (map.count == 0) return null

    hash usize := 0
    if ($(hash_key == null)) hash = normalize(auto_hash(K, key))
    else hash = normalize(call_hash(K, hash_key?, key))
    idx := hash & (map.entries.len - 1)

    while true {
        entry :: map.entries[idx]
        if (entry.hash == 0) return null

        if (entry.hash == hash) {
            if $(keys_eql == null) {
                if (auto_eql(K, entry.key, key)) return entry.value
            } else {
                if (call_eql(K, keys_eql?, entry.key, key)) return entry.value
            }
        }

        idx = (idx + 1) & (map.entries.len - 1)
    }
}

put proc(
    $K, $V type,
    $hash_key ?proc(key K) usize,
    $keys_eql ?proc(a, b K) bool,
    map @mut HashMap(K, V, hash_key, keys_eql),
    key K,
    value V,
) void ! (PutError | mem.AllocError) {
    threshold :: map.entries.len - divtrunc!(map.entries.len, 4)
    if (map.entries.len == 0 or map.count + 1 > threshold) {
        old_entries :: map.entries
        new_size :: if (old_entries.len > 0) old_entries.len * 2 else 8
        new_entries :: try mem.alloc(Entry(K, V), map.allocator, new_size)

        for (0..new_entries.len) |i| new_entries[i].hash = 0

        for old_entries |entry| {
            if (entry.hash == 0) continue

            rehash_index := entry.hash & (new_entries.len - 1)
            while (new_entries[rehash_index].hash != 0) rehash_index = (rehash_index + 1) & (new_entries.len - 1)

            new_entries[rehash_index] = entry
        }

        map.entries = new_entries
        mem.free(map.allocator, old_entries)
    }

    hash usize := 0
    if ($(hash_key == null)) hash = normalize(auto_hash(K, key))
    else hash = normalize(call_hash(K, hash_key?, key))
    idx := hash & (map.entries.len - 1)

    while true {
        entry :: map.entries[idx]
        if entry.hash == 0 {
            map.entries[idx] = Entry(K, V){
                hash = hash,
                key = key,
                value = value,
            }
            map.count += 1
            return
        }

        if (entry.hash == hash) {
            if $(keys_eql == null) {
                if (auto_eql(K, entry.key, key)) return .key_exists
            } else {
                if (call_eql(K, keys_eql?, entry.key, key)) return .key_exists
            }
        }

        idx = (idx + 1) & (map.entries.len - 1)
    }
}

# concrete comptime callbacks use named calls so their specializations are registered.
@hide:file
call_hash proc($K type, $callback proc(key K) usize, key K) usize {
    return callback(key)
}

@hide:file
call_eql proc($K type, $callback proc(a, b K) bool, a, b K) bool {
    return callback(a, b)
}

@hide:file
normalize proc(hash usize) usize {
    # mapping both 0 and 1 to 1 is safe because equality resolves
    # collisions (since hash and key must both be equal).
    if (hash == 0) return 1
    return hash
}

#! FNV-1a hash implementation.
#! note: vulnerable to collision attacks.
@hide:file
str_hash proc(key []u8) usize {
    hash u32 := 2166136261 # offset basis
    prime u32 :: 16777619

    for key |byte| {
        product u64 :: u64(hash xor u32(byte)) * u64(prime)
        hash = u32(product & u64(maxval!(u32)))
    }

    return usize(hash)
}

@hide:file
str_eql proc(a, b []u8) bool { return mem.eql(a, b) }

@hide:file
auto_hash proc($K type, key K) usize {
    match typeinfo!(K) {
        .integer: {
            bytes @[sizeof!(K)]u8 :: ptrcast!(&key)
            return str_hash(bytes[..])
        }
        .bool: return if (key) 1 else 0
        .enum |info|: {
            inline for (info.fields) |field| {
                if (key == field!(K, field.name)) return field.index
            }
            unreachable
        }
        .distinct |backing|: return auto_hash(ptrcast!(@backing, &key)^)
        .array: {
            hash usize := 0
            for (key) |element| hash = combine_hash(hash, auto_hash(element))
            return hash
        }
        .record |info|: {
            hash usize := 0
            inline for (info.fields) |field| {
                hash = combine_hash(hash, auto_hash(field!(key, field.name)))
            }
            return hash
        }
        else: compile_error!("AutoHashMap key requires integers, booleans, enums, distinct types, arrays or records; use HashMap for custom hashing")
    }
}

@hide:file
combine_hash proc(left, right usize) usize {
    # rotate before mixing so ordered fields do not cancel like a plain xor.
    return ((left << 5) xor (left >> (sizeof!(usize) * 8 - 5))) xor right
}

@hide:file
auto_eql proc($K type, left, right K) bool {
    match typeinfo!(K) {
        .integer, .bool, .enum: return left == right
        .distinct |backing|: return auto_eql(ptrcast!(@backing, &left)^, ptrcast!(@backing, &right)^)
        .array: {
            for (left) |element, index| {
                if (!auto_eql(element, right[index])) return false
            }
            return true
        }
        .record |info|: {
            inline for (info.fields) |field| {
                if (!auto_eql(field!(left, field.name), field!(right, field.name))) return false
            }
            return true
        }
        else: compile_error!("AutoHashMap key requires integers, booleans, enums, distinct types, arrays or records; use HashMap for custom equality")
    }
}