package pmap

type InplaceMap interface {
	Map
	SetInplace(Key, Value) InplaceMap
}

func (p pmap) SetInplace(k Key, v Value) InplaceMap {
	root, added := inplace_insert(p.root, k, v, p.hash)
	p.root = root
	if added {
		p.len++
	}
	return p
}

func inplace_insert(n interface{}, key Key, val Value, hashFn HashFunc) (newNode interface{}, added bool) {
	if n == nil {
		return leaf{key, val}, true
	}
	hash := hashFn(key)
	var insert func(n interface{}, shift uint32) interface{}
	insert = func(n interface{}, shift uint32) interface{} {
		// insert returns a new node (if the node should be replaced)
		// or nil (if the node is unchanged or modified in place)

		//fmt.Printf("insert %v %x %#v\n", shift, hash, n)
		switch n := n.(type) {
		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 {
				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
				x := newnode(n, h, shift)
				if x0 := insert(x, shift); x0 != nil {
					return x0
				}
				return x
			}
		case *node:
			c := n.getNode(shift, hash, key)
			if c == nil {
				// new node
				c = leaf{key, val}
				added = true

				m := bitmask(hash >> shift)
				n.bitmap |= m
				i := n.index(m)
				n.child = append(n.child, nil)   // allocate space
				copy(n.child[i+1:], n.child[i:]) // shift array up
				n.child[i] = c                   // mutate node
			} else {
				c = insert(c, shift+nodeShift)
				if c == nil {
					return nil
				}
				m := bitmask(hash >> shift)
				i := n.index(m)
				n.child[i] = c // mutate node
			}
			n.check()
			return nil
		case *collision:
			if n.hash != hash {
				// not a collision, so we must still have some hash bits left
				// split the trie
				x := newnode(n, n.hash, shift)
				if x0 := insert(x, shift); x0 != nil {
					return x0
				}
				return x
			}
			for i := range n.leaf {
				if key == n.leaf[i].k {
					// replace existing entry
					n.leaf[i] = leaf{key, val} // mutate
					added = false
					return nil
				}
			}
			// new collision
			added = true
			n.leaf = append(n.leaf, leaf{key, val}) // mutate
			return nil
		default:
			panic("pmap: unhandled case in insert")
		}
	}
	newNode = insert(n, 0)
	if newNode == nil {
		newNode = n
	}
	return
}