

record HashMap[K,V] (
    buckets : Vector[NodePtr[K,V]],
    size : SizeT,
);

alias NodePtr[K,V] = UniqueValue[Node[K,V]];
alias WeakNodePtr[K,V] = Pointer[Node[K,V]];

private record Node[K,V] (
    key : K,
    value : V,
    next : NodePtr[K,V],
);



//
// default constructor
//

[K,V]
overload HashMap[K,V]() --> returned:HashMap[K,V] {
    returned.buckets <-- Vector[NodePtr[K,V]]();
    returned.size = 0;

    try {
        resize(returned.buckets, 32);
    } catch (ex) {
        destroy(returned.buckets);
        throw ex;
    }
}



//
// clear
//

[K,V]
inline overload clear(a:HashMap[K,V]) {
    a.size = 0;
    for (bucket in a.buckets)
        bucket = nullUniqueValue(Node[K,V]);
}



//
// size
//

[K,V]
overload size(a:HashMap[K,V]) = a.size;



//
// lookup
//

[K,V]
overload lookup(a:HashMap[K,V], key:K) {
    var i = bucketId(a, key);
    return findInBucket(a.buckets[i], key);
}



//
// put
//

[K,V]
overload put(a:HashMap[K,V], key:K, forward value:V) {
    var i = bucketId(a, key);
    var ptr = findInBucket(a.buckets[i], key);
    if (not null?(ptr)) {
        ptr^ = value;
        return;
    }
    ref bucket = a.buckets[i];
    bucket = UniqueValue(Node(key, value, move(bucket)));
    a.size += 1;
    if (a.size > (size(a.buckets) * 3 / 4)) {
        grow(a);
    }
}



//
// index
//

[K,V]
overload index(a:HashMap[K,V], key:K) {
    var i = bucketId(a, key);
    var ptr = findInBucket(a.buckets[i], key);
    if (not null?(ptr))
        return ref ptr^;
    if (a.size >= (size(a.buckets) * 3 / 4)) {
        grow(a);
        i = bucketId(a, key); // recalculate index
    }
    ref bucket = a.buckets[i];
    bucket = UniqueValue(Node(key, V(), move(bucket)));
    a.size += 1;
    return ref bucket^.value;
}



//
// remove
//

[K,V]
overload remove(a:HashMap[K,V], key:K) {
    var i = bucketId(a, key);
    var prevNode = null(Node[K,V]);
    var node = &a.buckets[i]^;
    while (not null?(node)) {
        if (node^.key == key)
            break;
        prevNode = node;
        node = &node^.next^;
    }
    if (null?(node))
        return false;
    if (null?(prevNode))
        a.buckets[i] = node^.next;
    else
        prevNode^.next = node^.next;
    a.size -= 1;
    return true;
}



//
// items
//

[K,V]
overload items(forward a:HashMap[K,V]) {
    var x = HashMapItems(captureValue(a), K, V);
    while (null?(x.node) and hasNext?(x.bucketsIter))
        x.node = &next(x.bucketsIter)^;
    return move(x);
}

private record HashMapItems[CapturedMap,K,V] =
    map: CapturedMap,
    bucketsIter : SequenceIteratorType(Vector[NodePtr[K,V]]),
    node : WeakNodePtr[K,V];

[CapturedMap,K,V]
overload HashMapItems(forward captured:CapturedMap, static K, static V)
    --> returned:HashMapItems[CapturedMap,K,V]
{
    returned.map <-- captured;
    try {
        returned.bucketsIter <-- iterator(capturedRef(returned.map).buckets);
        returned.node <-- null(Node[K,V]);
    } catch (ex) {
        destroy(returned.map);
        throw ex;
    }
}

[C,K,V]
overload iterator(x:HashMapItems[C,K,V]) = x;

[C,K,V]
overload hasNext?(x:HashMapItems[C,K,V]) = not null?(x.node);

[C,K,V]
overload next(x:HashMapItems[C,K,V]) {
    assert(not null?(x.node));
    ref key, value = x.node^.key, x.node^.value;
    x.node = &x.node^.next^;
    while (null?(x.node) and hasNext?(x.bucketsIter))
        x.node = &next(x.bucketsIter)^;
    return ref key, value;
}



//
// assignment
//

[K,V]
overload assign(ref to: HashMap[K,V], ref from: HashMap[K,V]) {
    var tmp = from;
    destroy(to);
    to <-- move(tmp);
}



//
// helper procs
//

[K,V]
private bucketId(a:HashMap[K,V], key:K) =
    hash(key) % size(a.buckets);

[K,V]
private findInBucket(_node:NodePtr[K,V], key:K) {
    var node = &_node^;
    while (not null?(node)) {
        if (node^.key == key)
            return &node^.value;
        node = &node^.next^;
    }
    return null(V);
}

[K,V]
private grow(a:HashMap[K,V]) {
    var oldbuckets = move(a.buckets);
    a.size = 0;

    resize(a.buckets, size(oldbuckets) * 2);

    for(i in range(0, size(oldbuckets))) {
        var node = &oldbuckets[i]^;
        while (not null?(node)) {
            put(a, move(node^.key), move(node^.value));
            node = &node^.next^;
        }
    }
}
