From a121f47873172b28eb5b4789f91d34bea34d0a3b Mon Sep 17 00:00:00 2001
From: magical <magical@tilde.town>
Date: Sat, 22 Jan 2022 09:49:46 +0000
Subject: [PATCH] initial commit

---
 pmap.go      | 189 +++++++++++++++++++++++++++++++++++++++++++++++++++
 pmap_test.go |  20 ++++++
 2 files changed, 209 insertions(+)
 create mode 100644 pmap.go
 create mode 100644 pmap_test.go

diff --git a/pmap.go b/pmap.go
new file mode 100644
index 0000000..8a3e4b0
--- /dev/null
+++ b/pmap.go
@@ -0,0 +1,189 @@
+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
+}
diff --git a/pmap_test.go b/pmap_test.go
new file mode 100644
index 0000000..791a695
--- /dev/null
+++ b/pmap_test.go
@@ -0,0 +1,20 @@
+package pmap
+
+import "testing"
+
+func TestPmap(t *testing.T) {
+	p := New()
+	const numElems = 5
+	for i := range make([]int, numElems) {
+		p = p.Set(i, i)
+	}
+	if p.Len() != numElems {
+		t.Fatalf("Len() = %v, want %v", p.Len(), numElems)
+	}
+	for i := range make([]int, numElems) {
+		v, ok := p.Get(i)
+		if v != i || !ok {
+			t.Errorf("Get(%d) = %v %v, want %v %v", i, v, ok, i, true)
+		}
+	}
+}