Rewrite Map using a hash-balanced tree The gob-encoded deep copy approach was an easy hack but it failed in important ways. This complete rewrite uses a binary search tree where the tree's keys are hashes of the user's desired keys. This makes us pretty certain the tree will stay reasonably balanced without having to write a bunch of tree balancing code. This passes all tests and is about 200x faster than the previous implementation. I hacked this code together pretty quickly and there's lots of ugliness and repetitive code in there. I hope to clean it up soon.
diff --git a/map.go b/map.go index bc52931..0dc6d09 100644 --- a/map.go +++ b/map.go
@@ -1,65 +1,283 @@ // Immutable (i.e. persistent) data structures package ps -import "bytes" -import "encoding/gob" +import . "fmt" + +import "hash/fnv" type Any interface{} -type Map map[string]Any +type Map struct { + count int + hash uint64 // hash of the key (used for tree balancing) + key string + value Any + left *Map + right *Map +} // NewMap allocates a new, persistent map from strings to any value -func NewMap() Map { - m := make(map[string]Any) - return m +func NewMap() *Map { + var m Map + m.count = 0 // 0 count indicates an empty tree + return &m +} + +// IsNil returns true if the Map is empty +func (self *Map) IsNil() bool { + return self.count == 0 +} + +// clone returns an exact duplicate of a tree node +func (self *Map) clone() *Map { + var m Map + m.count = self.count + m.hash = self.hash + m.key = self.key + m.value = self.value + m.left = self.left + m.right = self.right + return &m +} + +// hashKey returns a hash code for a given string +func hashKey(key string) uint64 { + hasher := fnv.New64() + Fprint(hasher, key) + return hasher.Sum64() } // Set returns a new map similar to this one but with key and value // associated. If the key didn't exist, it's created; otherwise, the // associated value is changed. -func (prev Map) Set(key string, value Any) Map { - next := prev.Clone() - next[key] = value - return next +func (self *Map) Set(key string, value Any) *Map { + hash := hashKey(key) + return setLowLevel(self, hash, key, value) +} + +func setLowLevel(self *Map, hash uint64, key string, value Any) *Map { + if self == nil || self.IsNil() { // an empty tree is easy + var m Map + m.count = 1 + m.hash = hash + m.key = key + m.value = value + return &m + } + + if hash < self.hash { // insert into left subtree + m := self.clone() + m.left = setLowLevel(self.left, hash, key, value) + recalculateCount(m) + return m + } + if hash > self.hash { // insert into right subtree + m := self.clone() + m.right = setLowLevel(self.right, hash, key, value) + recalculateCount(m) + return m + } + + // replacing a key's previous value + m := self.clone() + m.value = value + return m +} + +// modifies a map by recalculating its key count based on the counts +// of its subtrees +func recalculateCount(m *Map) { + count := 0 + if m.left != nil { + count += m.left.Size() + } + if m.right != nil { + count += m.right.Size() + } + m.count = count + 1 // add one to count ourself } // Delete returns a new map with the association for key, if any, removed -func (prev Map) Delete(key string) Map { - next := prev.Clone() - delete(next, key) - return next +func (m *Map) Delete(key string) *Map { + hash := hashKey(key) + newMap, _ := deleteLowLevel(m, hash) + return newMap +} + +func deleteLowLevel(self *Map, hash uint64) (*Map, bool) { + // empty trees are easy + if self == nil { + return NewMap(), false + } + if self.IsNil() { + return self, false + } + + if hash < self.hash { // look in the left subtree + newLeft, found := deleteLowLevel(self.left, hash) + if !found { + return self, false + } + newMap := self.clone() + newMap.left = newLeft + recalculateCount(newMap) + } + if hash > self.hash { // look in the right subtree + newRight, found := deleteLowLevel(self.right, hash) + if !found { + return self, false + } + newMap := self.clone() + newMap.right = newRight + recalculateCount(newMap) + } + + // we must delete our own node + if self.isLeaf() { // we have no children + return NewMap(), true + } + if self.subtreeCount() == 1 { // only one subtree + if self.hasLeft() { // it's the left one + return self.left, true + } + return self.right, true // it's the right one + } + + // find a node to replace us + if self.left.Size() > self.right.Size() { // make left side smaller + replacement, newLeft := self.left.deleteRightmost() + newMap := replacement.clone() + newMap.left = newLeft + newMap.right = self.right + recalculateCount(newMap) + return newMap, true + } + + // make right side smaller + replacement, newRight := self.right.deleteLeftmost() + newMap := replacement.clone() + newMap.right = newRight + newMap.left = self.left + recalculateCount(newMap) + return newMap, true +} + +// delete the left or rightmost node in a tree returning the node that +// was deleted and the tree left over after its deletion +func (m *Map) deleteRightmost() (*Map, *Map) { + if m.isLeaf() { + return m, NewMap() + } + if m.hasRight() { + deleted, newRight := m.right.deleteRightmost() + newMap := m.clone() + newMap.right = newRight + recalculateCount(newMap) + return deleted, newMap + } + + deleted := m.clone() + deleted.left = nil + return deleted, m.left +} +func (m *Map) deleteLeftmost() (*Map, *Map) { + if m.isLeaf() { + return m, NewMap() + } + if m.hasLeft() { + deleted, newLeft := m.left.deleteLeftmost() + newMap := m.clone() + newMap.left = newLeft + recalculateCount(newMap) + return deleted, newMap + } + + deleted := m.clone() + deleted.count = 1 + deleted.right = nil + return deleted, m.right +} + +// hasLeft and hasRight return true if this tree has a left or right subtree +func (m *Map) hasLeft() bool { + return m.left != nil && !m.left.IsNil() +} +func (m *Map) hasRight() bool { + return m.right != nil && !m.right.IsNil() +} + +// isLeaf returns true if this is a leaf node +func (m *Map) isLeaf() bool { + return m.Size() == 1 +} + +// returns the number of child subtrees we have +func (m *Map) subtreeCount() int { + count := 0 + if m.hasLeft() { + count++ + } + if m.hasRight() { + count++ + } + return count } // Lookup returns a pair of values. The first is the value, the second is // true if the value exists. -func (m Map) Lookup(key string) (Any, bool) { - v, ok := m[key] - return v, ok +func (m *Map) Lookup(key string) (Any, bool) { + hash := hashKey(key) + return lookupLowLevel(m, hash) +} + +func lookupLowLevel(self *Map, hash uint64) (Any, bool) { + if self == nil || self.IsNil() { // an empty tree is easy + return nil, false + } + + if hash < self.hash { // look in the left subtree + return lookupLowLevel(self.left, hash) + } + if hash > self.hash { // look in the right subtree + return lookupLowLevel(self.right, hash) + } + + // we found it + return self.value, true } // Size returns the number of key value pairs in the map -func (m Map) Size() int { - return len(m) +func (m *Map) Size() int { + return m.count +} + +// ForEach executes a callback on each key value pair in the map +func (m *Map) ForEach(f func(key string, val Any)) { + if m.IsNil() { + return + } + + // left branch + if m.left != nil { + m.left.ForEach(f) + } + + // ourself + f(m.key, m.value) + + // right branch + if m.right != nil { + m.right.ForEach(f) + } } // Keys returns a slice containing all keys in this map -func (m Map) Keys() []string { - keys := make([]string, len(m)) +func (m *Map) Keys() []string { + keys := make([]string, m.Size()) i := 0 - for key := range m { - keys[i] = key + m.ForEach( func (k string, v Any) { + keys[i] = k i++ - } + }) return keys } - -// Clone deep copies `src` into `dst` -func (prev Map) Clone() Map { - next := NewMap() - var buf bytes.Buffer - enc := gob.NewEncoder(&buf) - dec := gob.NewDecoder(&buf) - enc.Encode(prev) - dec.Decode(&next) - return next -}