pmap/pvec.go

package pmap

import (
	"fmt"
	"math/bits"
	"unsafe"
)

const vecChunk = 16

type pvec struct {
	head *vnode
	tail []Value
	len  int
}

type vnode struct {
	car, cdr *vnode
}

type vleaf struct {
	//vnode // for debugging

	elem [vecChunk]Value
}

func getleaf(n *vnode, index, size int) (*vleaf, int) {
	// precondition: 0 <= index < size
	// cut is the largest power of 2 less than size
	// change the loop condition as follows for trees where the leaves are always at the same depth (incl the unbalanced ones)
	// 	for cut := floor(size); cut >= vecChunk; cut >>= 1
	for size > vecChunk {
		cut := floor(size)
		if index < cut {
			n = n.car
			//size = cut

			// optimization: size is now a power of two, so this subtree is complete and
			// the loop can be a lot simpler. small thing, but makes a difference when
			// you're doing a lot of gets
			for cut >>= 1; cut >= vecChunk; cut >>= 1 {
				if index < cut {
					n = n.car
				} else {
					n = n.cdr
					index -= cut
				}
			}
			break
		} else {
			n = n.cdr
			index -= cut
			size -= cut
		}
	}
	return (*vleaf)(unsafe.Pointer(n)), index
}

func list2pvec(elem []Value) *pvec {
	// 0, 8 == depth 0
	// 9, 16 == depth 1
	// 17, 32 == depth 2
	// etc
	depth := 0
	if len(elem) > 0 {
		depth = bits.Len((uint(len(elem)) - 1) / vecChunk)
	}
	_ = depth
	return &pvec{
		head: mkvnode(elem, len(elem)),
		tail: nil,
		len:  len(elem),
	}
}

func mkvnode(elem []Value, size int) *vnode {
	if len(elem) == 0 {
		return nil
	}
	if size <= vecChunk {
		leaf := new(vleaf)
		copy(leaf.elem[:], elem)
		return (*vnode)(unsafe.Pointer(leaf))
	}
	cut := floor(size)
	n := new(vnode)
	n.car = mkvnode(elem[:], cut)
	if cut < size {
		n.cdr = mkvnode(elem[cut:], size-cut)
	}
	return n
}

func pvec2list(p *pvec) []Value {
	var l = make([]Value, p.len)
	for i := range l {
		leaf, j := getleaf(p.head, i, p.len)
		if leaf == nil {
			panic(fmt.Sprint("nil leaf [index=", i, ", len=", p.len, "]"))
		}
		if j != i%vecChunk {
			panic("index mismatch")
		}
		l[i] = leaf.elem[j]
	}
	return l
}

// floor returns the largest power of 2 less than n
// returns 0 if n <= 0
func floor(n int) int {
	if n <= 1 {
		return 0
	}
	return 1 << bits.Len(uint(n-1)) >> 1
}