HashMap.an

import Std.C.putchar, calloc
import Std.Stream.Stream, iter, stream_fn, Emit, emit
import Std.Hash.Hash

export HashMap, Entry, Hash, empty, of, clear, resize, should_resize,
    insert, get_entry, get, contains, remove, len, is_empty, keys, values,
    get_or_insert, get_or_insert_with, retain, merge,
    print_hashmap, stream_map

// A linear-scan mutable HashMap using a resize factor of 2.0
type HashMap k v =
    len: Usz
    capacity: Usz
    entries: Ptr (Entry k v)

type Entry k v =
    key: k
    value: v
    occupied: Bool
    tombstone: Bool

/// Returns the length of this HashMap
HashMap.len (m: ref HashMap k v): Usz = m.len

/// Returns true if this HashMap is empty. Equivalent to `map.len () == 0`
HashMap.is_empty (m: ref HashMap k v): Bool = m.len == 0

/// Create a new, empty HashMap
/// This will not allocate until the first entry is inserted.
HashMap.empty () = HashMap 0 0 (null ())

/// Create a HashMap containing each `k, v` pair from the given stream.
HashMap.of (s: s) {Stream s (k, v)} {Hash k} {Eq k}: HashMap k v =
    var map = HashMap.empty ()
    iter s (fn (k, v) -> map.insert k v)
    map
 
HashMap.clear (map: mut HashMap k v) : Unit =
    if map.capacity != 0 then
        repeat map.capacity fn i ->
            entry = ptr_to_mut (offset map.entries i)
            entry.occupied := false
            entry.tombstone := false
        map.len := 0usz

HashMap.resize (map: mut HashMap k v) (new_capacity: Usz) {Hash k} {Eq k}: Unit =
    if new_capacity > map.capacity then
        new_memory = calloc new_capacity (size_of (MkType : Type (Entry k v)))

        if new_memory == null () then
            panic "Ran out of memory while resizing HashMap"

        var new_map = HashMap map.len new_capacity new_memory

        repeat map.capacity fn i ->
            entry: Entry k v = deref_ptr <| offset map.entries i
            if entry.occupied then
                new_map.insert entry.key entry.value

        map.capacity := new_capacity
        map.entries := new_memory

// Should we resize this map when pushing another element?
HashMap.should_resize (map: ref HashMap k v) : Bool =
    scale_factor = 2
    (map.len + 1) * scale_factor > map.capacity

/// Inserts a key-value pair into the map, removing any item that may previously
/// have been associated with the same key and returning the removed value if there was one.
HashMap.insert (map: mut HashMap k v) (key: k) (value: v) {Hash k} {Eq k}: Maybe v =
    if map.should_resize () then
        map.resize ((map.capacity + 1) * 2)

    var previous = Maybe.None
    iter_until (map: mut HashMap k v) (key: ref k) (value: ref v) (start: Usz) (end: Usz) : Bool =
        if start >= end then
            false
        else
            entry_ptr = offset map.entries start
            entry = ptr_to_mut entry_ptr
            if entry.occupied then
                if key == entry.key then
                    previous := Some (deref entry.value)
                    entry := Entry (deref key) (deref value) true false
                    true
                else
                    iter_until map key value (start + 1) end
            else
                entry := Entry (deref key) (deref value) true false
                true

    h = Usz (Hash.hash (deref (ref key))) % map.capacity

    if not iter_until map key value h map.capacity then
        if not iter_until map key value 0 h then
            panic "Failed to insert entry into map"

    if previous is None then
        map.len := map.len + 1usz
    previous

// TODO: Name resolution issue preventing this from being visible to iter_until
// if it is defined within get_entry
type EntryResult k v =
    | Found (Ptr (Entry k v))
    | NotFound
    | NotFinished

HashMap.get_entry (map: ref HashMap k v) (key: ref k) {Hash k} {eq: Eq k}: Maybe (ref (Entry k v)) =
    map = transmute map : mut HashMap k v
    map.get_entry_mut key

HashMap.get_entry_mut (map: mut HashMap k v) (key: ref k) {Hash k} {eq: Eq k}: Maybe (mut (Entry k v)) =
    if map.capacity == 0 then return None

    // Return whether to continue by wrapping around the end, and the value if found
    iter_until (start: Usz) (end: Usz) : EntryResult k v =
        if start >= end then
            NotFinished
        else
            entry_ptr = offset map.entries start
            entry = deref_ptr entry_ptr
            if entry.occupied and ref entry.key == key then
                Found entry_ptr
            else if entry.tombstone then
                iter_until (start + 1) end
            else
                NotFound

    h = Usz (Hash.hash (deref key)) % map.capacity
    match iter_until h map.capacity
    | Found value -> Some (ptr_to_mut value)
    | NotFound -> None
    | NotFinished ->
        match iter_until 0 h
        | Found value -> Some (ptr_to_mut value)
        | _ -> None

HashMap.get (map: ref HashMap k v) (key: ref k) {Hash k} {Eq k} : Maybe v =
    match map.get_entry key
    | Some entry -> Some (deref entry.value)
    | None -> None

HashMap.contains (map: ref HashMap k v) (key: ref k) {Hash k} {Eq k}: Bool =
    map.get key is Some _

HashMap.remove (map: mut HashMap k v) (key: k) {Hash k} {Eq k} : Maybe v =
    match map.get_entry_mut key
    | None -> None
    | Some entry ->
        entry.occupied := false
        entry.tombstone := true
        map.len := map.len - 1
        Some (deref entry.value)

impl print_hashmap {Print k} {Print v}: Print (ref HashMap k v) with
    print map =
        putchar '['
        var printed_first = false
        for i in 0 .. map.capacity do
            entry: Entry k v = map.entries.[i]
            if entry.occupied then
                if printed_first then
                    print ", "

                print entry.key
                print " -> "
                print entry.value
                printed_first := true

        putchar ']'

implicit stream_map: Stream (HashMap k v) (k, v) =
    Stream fn map e ->
        for i in 0 .. map.capacity do
            entry: Entry k v = map.entries.[i]
            if entry.occupied then
                e.emit (entry.key, entry.value)

/// Emit each key. Bit-copies the key; safe for `Copy` keys.
HashMap.keys (m: ref HashMap k v) = fn (e: Emit k) ->
    for i in 0 .. m.capacity do
        entry: Entry k v = m.entries.[i]
        if entry.occupied then e.emit entry.key

/// Emit each value. Bit-copies the value; safe for `Copy` values.
HashMap.values (m: ref HashMap k v) = fn (e: Emit v) ->
    for i in 0 .. m.capacity do
        entry: Entry k v = m.entries.[i]
        if entry.occupied then e.emit entry.value

/// Return the existing value at `key`, or insert `default` and return it.
/// Requires `Copy k` so we can probe the map twice with the same key.
HashMap.get_or_insert (m: mut HashMap k v) (key: k) (default: v) {Hash k} {Eq k} {Copy k}: v =
    match m.get (ref key)
    | Some existing -> existing
    | None ->
        ignore (m.insert key default)
        match m.get (ref key)
        | Some v -> v
        | None -> panic "HashMap.get_or_insert: lookup after insert failed"

/// Same as `get_or_insert` but builds the default lazily.
HashMap.get_or_insert_with (m: mut HashMap k v) (key: k) (mk: fn Unit => v) {Hash k} {Eq k} {Copy k}: v =
    match m.get (ref key)
    | Some existing -> existing
    | None ->
        ignore (m.insert key (mk ()))
        match m.get (ref key)
        | Some v -> v
        | None -> panic "HashMap.get_or_insert_with: lookup after insert failed"

/// Remove every entry for which `pred key value` returns false.
HashMap.retain (m: mut HashMap k v) (pred: fn k v => Bool) {Copy k} {Copy v}: Unit =
    for i in 0 .. m.capacity do
        entry_ptr = offset m.entries i
        entry = ptr_to_mut entry_ptr
        if entry.occupied then
            if not (pred entry.key entry.value) then
                entry.occupied := false
                entry.tombstone := true
                m.len -= 1

/// Insert every entry from `src` into `dst`, overwriting on conflict.
HashMap.merge (dst: mut HashMap k v) (src: HashMap k v) {Hash k} {Eq k}: Unit =
    iter src (fn (k, v) -> ignore (HashMap.insert dst k v))