compact_map.oc

//! An order-preserving hash map.
//!
//! Uses open addressing with linear probing. Inspired by Python's dict.

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

const INDEX_FREE: i32 = -1
const INDEX_DELETED: i32 = -2

struct Item<K, V> {
    hash: u32
    key: K
    value: V
}

struct Map<K, V> {
    items: &Vector<Item<K, V>>
    indices: &i32
    capacity: u32
    num_tombstones: u32
}

def Map::new(capacity: u32 = 16): &Map<K, V> {
    let items = Vector<Item<K, V>>::new(capacity)
    let indices = mem::alloc<i32>(capacity)
    for let i = 0; i < capacity; i++ {
        indices[i] = INDEX_FREE
    }
    let map = mem::alloc<Map<K, V>>()
    map.items = items
    map.indices = indices
    map.capacity = capacity
    map.num_tombstones = 0
    return map
}

def Map::free(&this) {
    if not this? return
    mem::free(.indices)
    .items.free()
    mem::free(this)
}

def Map::get_index(&this, key: K, hash: u32): u32 {
    let perturb = hash
    let j = hash % this.capacity
    let i = j

    let first_deleted = -1
    while .indices[i] != INDEX_FREE {
        if .indices[i] == INDEX_DELETED {
            if first_deleted < 0 {
                first_deleted = i as i32
            }
        } else {
            let item = .items.data[.indices[i]]  // Should be safe
            if item.hash == hash and item.key.eq(key) {
                return i
            }
        }
        j = 5 * j + perturb + 1
        i = j % .capacity
        perturb = perturb >> 5
    }

    if first_deleted < 0 {
        return i
    }
    return first_deleted as u32
}


// TODO: Don't mess up the insertion order when removing an item. Maybe use a linked list?
def Map::remove(&this, key: K) {
    let hash = key.hash()
    let index = .get_index(key, hash)
    if .indices[index] == INDEX_FREE {
        return
    }
    assert index >= 0

    let item_index = .indices[index] as u32
    let item = .items.at(item_index)

    let last_item = .items.at(.items.size - 1)
    let last_index = .get_index(last_item.key, last_item.hash)

    assert last_index >= 0
    assert .indices[last_index] >= 0

    let last_item_index = .indices[last_index] as u32
    assert last_item_index == .items.size - 1 // Sanity check

    .indices[last_index] = item_index as i32

    .indices[index] = INDEX_DELETED
    .items.data[item_index] = last_item
    .items.pop()
    .num_tombstones += 1

    .resize_if_necessary()
}

def Map::resize(&this, new_capacity: u32) {
    let old_indices = .indices
    .indices = mem::alloc<i32>(new_capacity)
    .capacity = new_capacity
    for let i = 0; i < new_capacity; i++ {
        .indices[i] = INDEX_FREE
    }
    for let i = 0; i < .items.size; i++ {
        let item = .items.at(i)
        let index = .get_index(item.key, item.hash)

        if .indices[index] == INDEX_FREE {
            .indices[index] = i as i32
        }
    }
    .num_tombstones = 0
    mem::free(old_indices)
}

def Map::resize_if_necessary(&this) {
    if .num_tombstones + .items.size as u32 >= .capacity * 3 / 4 {
        .resize(.capacity * 2)
    }
}

[operator "[]="]
def Map::insert(&this, key: K, value: V) {
    let hash = key.hash()
    let index = .get_index(key, hash)

    if .indices[index] < 0 {
        .indices[index] = .items.size as i32
        .items.push(Item<K, V>(hash, key, value))
        .resize_if_necessary()

    } else {
        if .indices[index] == INDEX_DELETED {
            .num_tombstones -= 1
        }
        let item_index = .indices[index] as u32
        .items.data[item_index].value = value   // Should be safe
    }
}

def Map::get_item(&this, key: K): &Item<K, V> {
    let hash = key.hash()
    let index = .get_index(key, hash)
    if .indices[index] < 0 {
        return null
    }

    let idx = .indices[index] as u32
    return &.items.data[idx] // Should be safe
}

def Map::get(&this, key: K, defolt: V): V {
    let item = .get_item(key)
    if not item? return defolt
    return item.value
}

[operator "[]"]
def Map::at(&this, key: K): V {
    let item = .get_item(key)
    if not item? {
        assert false, "Key not found in Map::at()"
    }
    return item.value
}

def Map::size(&this): u32 => .items.size

[operator "in"]
def Map::contains(&this, key: K): bool {
    let hash = key.hash()
    let index = .get_index(key, hash)
    return .indices[index] >= 0
}

def Map::clear(&this) {
    .items.clear()
    std::libc::memset(.indices, INDEX_FREE as u8, .capacity * sizeof(i32))
}

def Map::is_empty(&this): bool => .items.size == 0

[operator "+="]
def Map::extend(&this, other: &Map<K, V>) {
    for item in other.iter() {
        .insert(item.key, item.value)
    }
}

def Map::iter(&this): Iterator<K, V> => Iterator<K, V>(.items.iter())
def Map::iter_keys(&this): KeyIterator<K, V> => KeyIterator<K, V>(.items.iter())
def Map::iter_values(&this): ValueIterator<K, V> => ValueIterator<K, V>(.items.iter())

struct Iterator<K, V> {
    iter: VectorIterator<Item<K, V>>
}
def Iterator::has_value(&this): bool => .iter.has_value()
def Iterator::next(&this) { .iter.next() }
def Iterator::cur(&this): Item<K, V> => .iter.cur()

struct KeyIterator<K, V> {
    iter: VectorIterator<Item<K, V>>
}
def KeyIterator::has_value(&this): bool => .iter.has_value()
def KeyIterator::next(&this) { .iter.next() }
def KeyIterator::cur(&this): K => .iter.cur().key


struct ValueIterator<K, V> {
    iter: VectorIterator<Item<K, V>>
}
def ValueIterator::has_value(&this): bool => .iter.has_value()
def ValueIterator::next(&this) { .iter.next() }
def ValueIterator::cur(&this): V => .iter.cur().value