map.oc

//! A simple hash-map
//!
//! This hash map uses closed hashing with linked lists for collision resolution.
//! The hash map will automatically resize itself when the number of items
//! exceeds the number of buckets.
//!
//! Sample usage:
//!
//!     import std::map::Map
//!
//!     let map = Map<u32, str>::new()
//!     map.insert(1, "one")
//!     map.insert(2, "two")
//!     for it in map.iter() {
//!         println(`{it.key} = {it.value}`)
//!     }
//!     map.remove(1)

import std::vector::Vector
import std::traits::{ hash, eq }
import std::mem

//* Linked list node for the hash map
struct Item<K, V> {
    key: K
    value: V
    next: &Item<K, V>
}

//* Creates a new node
def Item::new(key: K, value: V, next: &Item<K, V> = null): &Item<K, V> {
    let node = mem::alloc<Item<K, V>>()
    node.key = key
    node.value = value
    node.next = next
    return node
}

//* Free the linked list starting at this node
def Item::free_list(&this) {
    let cur = this
    while cur? {
        let next = cur.next
        std::mem::free(cur)
        cur = next
    }
}

//* A Hash-map
//*
//* Preconditions:
//*  - `K` type must implement the `hash` method: `K::hash(this): u32`
//*  - `K` type must implement the `eq` method: `K::eq(this, b: K): bool`
struct Map<K, V> {
    buckets: &&Item<K, V>
    size: u32
    num_buckets: u32
    num_collisions: u32
}

//* Creates a new hash map
def Map::new(capacity: u32 = 8): &Map<K, V> {
    let map = mem::alloc<Map<K, V>>()
    map.num_buckets = capacity
    map.buckets = mem::alloc<&Item<K, V>>(map.num_buckets)
    return map
}

//* Hashes a key and returns the bucket index
def Map::hash(&this, key: K): u32 {
    let hash = key.hash()
    hash = hash % .num_buckets
    if hash < 0 {
        hash += .num_buckets
    }
    return hash
}

//* Returns the node with the given key, or null
def Map::get_item(&this, key: K): &Item<K, V> {
    let hash = .hash(key)
    let node = .buckets[hash]
    while node? {
        if node.key.eq(key) {
            return node
        }
        node = node.next
    }
    return null
}

[operator "[]"]
//* Returns the value for the given key, or fails
def Map::at(&this, key: K): V {
    let node = .get_item(key)
    assert node?, "Key not found"
    return node.value
}

//* Returns the value for the given key, or the default value
def Map::get(&this, key: K, defolt: V): V {
    let node = .get_item(key)
    if not node? then return defolt
    return node.value
}

[operator "in"]
//* Checks if the map contains the given key
def Map::contains(&this, key: K): bool {
    return .get_item(key)?
}

[operator "[]="]
//* Inserts / updates the given key-value pair
def Map::insert(&this, key: K, value: V) {
    let node = .get_item(key)
    if node? {
        node.value = value
    } else {
        let hash = .hash(key)
        let new_node = Item<K, V>::new(key, value, .buckets[hash])
        if .buckets[hash]? {
            .num_collisions += 1
        }
        .buckets[hash] = new_node
        .size += 1
        if .size > .num_buckets {
            .resize()
        }
    }
}

[operator "+="]
//* Inserts all key-value pairs from the other map into
def Map::extend(&this, other: &Map<K, V>) {
    for iter in other.iter() {
        .insert(iter.key, iter.value)
    }
}

//* Removes the given key from the map
def Map::remove(&this, key: K) {
    let node = .get_item(key)
    if node? {
        let hash = .hash(key)
        if .buckets[hash] == node {
            .buckets[hash] = node.next
        } else {
            let prev = .buckets[hash]
            while prev.next != node {
                prev = prev.next
            }
            prev.next = node.next
        }
        mem::free(node)
        .size -= 1
    }
}

//* Resizes the map to double the number of buckets
def Map::resize(&this) {
    let old_buckets = .buckets
    let old_num_buckets = .num_buckets
    let old_size = .size
    .num_collisions = 0
    .num_buckets *= 2
    .num_buckets = .num_buckets.max(16)
    .buckets = mem::alloc<&Item<K, V>>(.num_buckets)
    for let i = 0; i < old_num_buckets; i += 1 {
        let node = old_buckets[i]
        while node? {
            let new_hash = .hash(node.key)
            let new_node = Item<K, V>::new(node.key, node.value, .buckets[new_hash])
            if .buckets[new_hash]? {
                .num_collisions += 1
            }
            .buckets[new_hash] = new_node
            node = node.next
        }
    }
    for let i = 0; i < old_num_buckets; i += 1 {
        old_buckets[i].free_list()
    }
    mem::free(old_buckets)
}

//* Checks if the map is empty
def Map::is_empty(&this): bool => .size == 0

//* Frees the map and all its nodes
def Map::free(&this) {
    for let i = 0; i < .num_buckets; i += 1 {
        .buckets[i].free_list()
    }
    mem::free(.buckets)
    mem::free(this)
}

//* Clears the map
def Map::clear(&this) {
    for let i = 0; i < .num_buckets; i += 1 {
        .buckets[i].free_list()
        .buckets[i] = null
    }
    .size = 0
    .num_collisions = 0
}

//* Returns an iterator over the map
def Map::iter(&this): Iterator<K, V> {
    return Iterator<K, V>::make(this)
}

//* Returns an iterator over the keys of the map
def Map::iter_keys(&this): KeyIterator<K, V> => KeyIterator<K, V>(.iter())

//* Returns an iterator over the values of the map
def Map::iter_values(&this): ValueIterator<K, V> => ValueIterator<K, V>(.iter())

//* Iterator over the items of a map
struct Iterator<K, V> {
    idx: i32
    node: &Item<K, V>
    map: &Map<K, V>
}

def Iterator::key(&this): K {
    return .node.key
}

def Iterator::value(&this): V {
    return .node.value
}

def Iterator::make(map: &Map<K, V>): Iterator<K, V> {
    let it = Iterator<K, V>(idx: -1, node: null, map)
    it.next()
    return it
}

def Iterator::has_value(&this): bool => .node?

def Iterator::cur(&this): &Item<K, V> => .node

def Iterator::next(&this) {
    while .idx < .map.num_buckets as i32 {
        while .node? {
            .node = .node.next
            if .node? return
        }
        .idx += 1
        .node = if .idx < .map.num_buckets as i32 {
            yield .map.buckets[.idx]
        } else {
            yield null
        }
        if .node? return
    }
}

//* Iterator over the keys of a map
struct KeyIterator<K, V> {
    map_iter: Iterator<K, V>
}
def KeyIterator::has_value(&this): bool => .map_iter.has_value()
def KeyIterator::cur(&this): K => .map_iter.cur().key
def KeyIterator::next(&this) {
    .map_iter.next()
}

//* Iterator over the values of a map
struct ValueIterator<K, V> {
    map_iter: Iterator<K, V>
}
def ValueIterator::has_value(&this): bool => .map_iter.has_value()
def ValueIterator::cur(&this): V => .map_iter.cur().value
def ValueIterator::next(&this) {
    .map_iter.next()
}