pmap/pmap.go

package pmap

const deg = 8 // branch factor of nodes
const log2deg = 3
const mask = 0b111

type Key = interface{}
type Value = interface{}

type Map interface {
	Get(Key) (Value, bool)
	Set(Key, Value) Map
	Del(Key) Map
	Len() int
}

type pmap struct {
	root interface{}
	len  int
	hash func(Key) uint32
}

// A Map implemented as a hashed trie
type node struct {
	child  [deg]interface{}
	bitmap uint32
}

type collision struct {
	hash uint32
	leaf []leaf
}

type leaf struct {
	k Key
	v Value
}

func New() Map {
	return pmap{}
}

func (p pmap) Len() int {
	return p.len
}

func (p pmap) Get(k Key) (Value, bool) {
	h := hash(k) // TODO: p.hash
	return lookup(p.root, 0, h, k)
}

func (p pmap) Set(k Key, v Value) Map {
	//p.root = leaf{k, v}
	//p.len = 1
	//return p

	h := hash(k) // TODO: p.hash
	n, _ := p.root.(*node)
	if n == nil {
		n = &node{}
	}
	n.child[h&mask] = leaf{k, v}
	n.bitmap = 1 << (h & mask)
	p.root = n
	p.len++
	return p
}

func (p pmap) Del(k Key) Map {
	return p
}

func hash(k Key) uint32 {
	u := k.(int)
	return uint32(uint(u) + uint(u)>>32)
}

func (m *node) getNode(shift, hash uint32, key Key) interface{} {
	i := hash >> shift & mask
	return m.child[i]
}

func (m *collision) getNode(hash uint32, key Key) interface{} {
	if hash != m.hash {
		return nil
	}
	for i := range m.leaf {
		if key == m.leaf[i].k {
			return m.leaf[i].v
		}
	}
	return nil
}

func lookup(root interface{}, shift, hash uint32, key Key) (Value, bool) {
	cur := root
	for {
		switch n := cur.(type) {
		case nil:
			return nil, false
		case leaf:
			if n.k == key {
				return n.v, true
			} else {
				return nil, false
			}
		case *node:
			cur = n.getNode(shift, hash, key)
			shift += log2deg
		case *collision:
			cur = n.getNode(hash, key)
		default:
			panic("pmap: unhandled case in lookup")
		}
	}
}

func singleton(key Key, val Value, hash, shift uint32) *node {
	n := &node{}
	idx := hash >> shift & mask
	n.child[idx] = leaf{key, val}
	n.bitmap = 1 << idx
	return n
}

type hashFunc = func(Key) uint32

func insert(n interface{}, shift, hash uint32, key Key, val Value, hashFn hashFunc) (newNode interface{}, added bool) {
	if n == nil {
		return leaf{key, val}, true
	}
	var insert func(n interface{}, shift uint32) interface{}
	insert = func(n interface{}, shift uint32) interface{} {
		switch n := n.(type) {
		//case nil:
		//	added = true
		//	return leaf{key, val}
		case leaf:
			if n.k == key {
				// replace existing entry
				added = false
				return leaf{key, val}
			} else if h := hashFn(n.k); h == hash {
				// collision
				added = true
				return collision{hash, []leaf{{key, val}, n}}
			} else {
				// split the trie
				m := singleton(key, val, hash, shift)
				return insert(m, shift+log2deg)
			}
		case *node:
			c := n.getNode(shift, hash, key)
			if c == nil {
				// new node
				// TODO
				added = true
				return singleton(key, val, shift, hash)
			} else {
				// if it's a leaf and the hashes match, form a collision
				// if we run out of hash, form a leaf
				//
				return insert(c, shift+log2deg)
			}
		case *collision:
			if n.hash != hash {
				panic("hash mismatch") // TODO: can this happen? do we even need collision.hash?
			}
			for i := range n.leaf {
				if key == n.leaf[i].k {
					// replace existing entry
					l := make([]leaf, 1, len(n.leaf))
					l[0] = leaf{key, val}
					l = append(l, n.leaf[:i]...)
					l = append(l, n.leaf[i+1:]...)
					added = false
					return collision{hash, l}
				}
			}
			// new collision
			added = true
			return collision{hash, append([]leaf{{key, val}}, n.leaf...)}
		default:
			panic("pmap: unhandled case in insert")
		}
	}
	newNode = insert(n, shift)
	return
}