From d27d5348f942ccce0cd90cc4879d2f336f902e41 Mon Sep 17 00:00:00 2001
From: magical <magical@tilde.town>
Date: Sat, 22 Jan 2022 21:04:29 +0000
Subject: [PATCH] booyah

---
 pmap.go      | 50 ++++++++++++++++++++++++++++++++++++--------------
 pmap_test.go |  2 +-
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/pmap.go b/pmap.go
index 8a3e4b0..eedef9c 100644
--- a/pmap.go
+++ b/pmap.go
@@ -50,11 +50,20 @@ func (p pmap) Get(k Key) (Value, bool) {
 }
 
 func (p pmap) Set(k Key, v Value) Map {
+	h := hash(k) // TODO: p.hash
+	root, added := insert(p.root, 0, h, k, v, hash)
+	p.root = root
+	if added {
+		p.len++
+	}
+	//pretty.Println(p)
+	return p
+
 	//p.root = leaf{k, v}
 	//p.len = 1
 	//return p
 
-	h := hash(k) // TODO: p.hash
+	h = hash(k) // TODO: p.hash
 	n, _ := p.root.(*node)
 	if n == nil {
 		n = &node{}
@@ -116,9 +125,13 @@ func lookup(root interface{}, shift, hash uint32, key Key) (Value, bool) {
 }
 
 func singleton(key Key, val Value, hash, shift uint32) *node {
+	return newnode(leaf{key, val}, hash, shift)
+}
+
+func newnode(child interface{}, hash, shift uint32) *node {
 	n := &node{}
 	idx := hash >> shift & mask
-	n.child[idx] = leaf{key, val}
+	n.child[idx] = child
 	n.bitmap = 1 << idx
 	return n
 }
@@ -129,8 +142,9 @@ func insert(n interface{}, shift, hash uint32, key Key, val Value, hashFn hashFu
 	if n == nil {
 		return leaf{key, val}, true
 	}
-	var insert func(n interface{}, shift uint32) interface{}
-	insert = func(n interface{}, shift uint32) interface{} {
+	var _insert func(n interface{}, shift uint32) interface{}
+	_insert = func(n interface{}, shift uint32) interface{} {
+		//fmt.Printf("insert %v %x %#v\n", shift, hash, n)
 		switch n := n.(type) {
 		//case nil:
 		//	added = true
@@ -145,26 +159,34 @@ func insert(n interface{}, shift, hash uint32, key Key, val Value, hashFn hashFu
 				added = true
 				return collision{hash, []leaf{{key, val}, n}}
 			} else {
+				if h>>shift == hash>>shift {
+					panic("pmap: infinite loop in insert")
+				}
+				// not a collision, so we must still have some hash bits left
 				// split the trie
-				m := singleton(key, val, hash, shift)
-				return insert(m, shift+log2deg)
+				m := newnode(n, h, shift)
+				return _insert(m, shift)
 			}
 		case *node:
 			c := n.getNode(shift, hash, key)
 			if c == nil {
 				// new node
-				// TODO
+				c = leaf{key, val}
 				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)
+				c = _insert(c, shift+log2deg)
 			}
+			idx := hash >> shift & mask
+			x := &node{child: n.child, bitmap: n.bitmap}
+			x.child[idx] = c
+			x.bitmap |= 1 << idx
+			return x
 		case *collision:
 			if n.hash != hash {
-				panic("hash mismatch") // TODO: can this happen? do we even need collision.hash?
+				// not a collision, so we must still have some hash bits left
+				// split the trie
+				m := newnode(n, n.hash, shift)
+				return _insert(m, shift)
 			}
 			for i := range n.leaf {
 				if key == n.leaf[i].k {
@@ -184,6 +206,6 @@ func insert(n interface{}, shift, hash uint32, key Key, val Value, hashFn hashFu
 			panic("pmap: unhandled case in insert")
 		}
 	}
-	newNode = insert(n, shift)
+	newNode = _insert(n, shift)
 	return
 }
diff --git a/pmap_test.go b/pmap_test.go
index 791a695..7f25dd3 100644
--- a/pmap_test.go
+++ b/pmap_test.go
@@ -4,7 +4,7 @@ import "testing"
 
 func TestPmap(t *testing.T) {
 	p := New()
-	const numElems = 5
+	const numElems = 100
 	for i := range make([]int, numElems) {
 		p = p.Set(i, i)
 	}