Shrinking map using go generics

shrinkmap.go2

// Just an experiment to build a minimal "fix" for
// https://github.com/golang/go/issues/20135
// using go generics.
//
// It's actually a better "go" map than the builtin
// ones, as the zero value is a valid map. :)
//
// WARNING: ONLY MINIMAL TESTING HAS BEEN PERFORMED.

package main

import "fmt"

const (
	defaultShrinkWork = 4 // how many entries to migrate for each modification
)

type Map(type K comparable, V interface{}) struct {
	cur        map[K]V
	next       map[K]V
	max        int
	modCount   int
	holdShrink bool
}

func (m *Map(K, V)) init() {
	if m.cur == nil {
		m.cur = make(map[K]V)
	}
}

func (m *Map(K, V)) Put(key K, value V) {
	m.init()
	if m.next == nil {
		m.cur[key] = value
	} else {
		delete(m.cur, key)
		m.next[key] = value
	}
	if l := m.Len(); m.max < l {
		m.max = l
	}
	m.doShrinkWork(defaultShrinkWork)
}

func (m *Map(K, V)) Get(key K) (V, bool) {
	m.init()
	v, ok := m.cur[key]
	if m.next != nil && !ok {
		v, ok = m.next[key]
	}
	return v, ok
	/*
		var v V
		var ok bool
		if m.next == nil {
			v, ok = m.cur[key]
		} else if len(m.cur) >= len(m.next) {
			v, ok = m.cur[key]
			if !ok {
				v, ok = m.next[key]
			}
		} else {
			v, ok = m.next[key]
			if !ok {
				v, ok = m.cur[key]
			}
		}
		return v, ok
	*/
}

func (m *Map(K, V)) Delete(key K) bool {
	m.init()
	if _, ok := m.cur[key]; ok {
		delete(m.cur, key)
		m.doShrinkWork(defaultShrinkWork)
		m.checkShrink() // after doShinkWork so that we do either that or checkShrink but not both
		return true
	}
	if m.next != nil {
		if _, ok := m.next[key]; ok {
			delete(m.next, key)
			m.doShrinkWork(defaultShrinkWork)
			return true
		}
	}
	return false
}

func (m *Map(K, V)) checkShrink() {
	if m.next == nil && m.max > 10 && m.max > len(m.cur)*3 && !m.holdShrink {
		m.next = make(map[K]V, len(m.cur)*3/2)
	}
}

func (m *Map(K, V)) ForEach(f func(K, V) bool) {
	m.init()

	if m.holdShrink {
		panic("ForEach called concurrently")
	}
	m.holdShrink = true
	m.modCount = 0
	defer func() {
		m.holdShrink = false
		m.checkShrink()
		m.doShrinkWork(m.modCount * defaultShrinkWork)
	}()

	for k, v := range m.cur {
		if f(k, v) {
			return
		}
	}
	for k, v := range m.next {
		if f(k, v) {
			return
		}
	}
}

func (m *Map(K, V)) doShrinkWork(work int) {
	m.modCount++
	if m.next != nil && !m.holdShrink && work > 0 {
		m.doShrinkWorkSlow(work)
	}
}

func (m *Map(K, V)) doShrinkWorkSlow(work int) {
	for k, v := range m.cur {
		if _, ok := m.next[k]; ok {
			panic(fmt.Sprintf("invariant violation: key exists in both maps: %q", key))
		}
		m.next[k] = v
		delete(m.cur, k)
		work--
		if work <= 0 {
			break
		}
	}
	if len(m.cur) == 0 {
		m.cur, m.next, m.max = m.next, nil, len(m.next)
	}
}

func (m *Map(K, V)) Shrink() {
	m.init()
	m.doShrinkWork(m.Len())
}

func (m *Map(K, V)) Len() int {
	m.init()
	return len(m.cur) + len(m.next)
}

func main() {
	{
		m := Map(int, int){}
		for i := 0; i < 100000; i++ {
			m.Put(i, i)
		}
		for i := 0; i < 100000; i++ {
			m.Delete(i)
		}
		fmt.Println(m.Len())
	}
	{
		m := Map(int, int){}
		for i := 0; i < 100000; i++ {
			m.Put(i, i)
		}
		m.ForEach(func(k, _ int) bool {
			m.Delete(k)  // you can call any function(s) apart from m.ForEach
			return false // false->continue, true->break
		})
		fmt.Println(m.Len())
	}
}