Binary Tree as immutable state backing, with rebinding pattern for minimal copy-on-write modifications. Every merkle cached *by default*.

experimental.go

package main

import (
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"github.com/protolambda/zrnt/eth2/core"
	"log"
)

// Experimental code! Everything a tree and cached by default.
// Also spaghetti, iterating on the idea/approach first, then refactoring/polishing later.

const (
	mask0 = ^uint64((1 << (1 << iota)) - 1)
	mask1
	mask2
	mask3
	mask4
	mask5
)

const (
	bit0 = uint8(1 << iota)
	bit1
	bit2
	bit3
	bit4
	bit5
)

func GetDepth(v uint64) (out uint8) {
	// bitmagic: binary search through a uint64, offset down by 1 to not round powers of 2 up.
	// Then adding 1 to it to not get the index of the first bit, but the length of the bits (depth of tree)
	// Zero is a special case, it has a 0 depth.
	// Example:
	//  (in out): (0 0), (1 1), (2 1), (3 2), (4 2), (5 3), (6 3), (7 3), (8 3), (9 4)
	if v == 0 {
		return 0
	}
	v--
	if v&mask5 != 0 {
		v >>= bit5
		out |= bit5
	}
	if v&mask4 != 0 {
		v >>= bit4
		out |= bit4
	}
	if v&mask3 != 0 {
		v >>= bit3
		out |= bit3
	}
	if v&mask2 != 0 {
		v >>= bit2
		out |= bit2
	}
	if v&mask1 != 0 {
		v >>= bit1
		out |= bit1
	}
	if v&mask0 != 0 {
		out |= bit0
	}
	out++
	return
}

type HashFn func(a Root, b Root) Root

func (h *HashFn) MerkleizeBytes(b []byte) Root {
	// TODO
	return Root{}
}

// A link is called to rebind a value
type Link func(v Node)

type Root [32]byte

func (s *Root) ComputeRoot(h HashFn) Root {
	if s == nil {
		return Root{}
	}
	return *s
}

type GetterInteraction interface {
	Getter(target uint64, depth uint8) (Node, error)
}
type SetterInteraction interface {
	Setter(target uint64, depth uint8) (Link, error)
}
type ExpandIntoInteraction interface {
	ExpandInto(target uint64, depth uint8) (Link, error)
}

type NodeInteraction interface {
	GetterInteraction
	SetterInteraction
	ExpandIntoInteraction
	// TODO: refactor these to use generalized indices as tree position.
}

type Node interface {
	ComputeRoot(h HashFn) Root
}

type RebindableNode interface {
	Node
	Bind(bindingLink Link)
}

type ComplexNode interface {
	Node
	NodeInteraction
}

type SubtreeView struct {
	ComplexNode
	Link
	depth uint8
}

func (sv *SubtreeView) Bind(bindingLink Link) {
	sv.Link = bindingLink
}

func (sv *SubtreeView) RebindInner(v Node) {
	// the view keeps track of the latest node, but proxies the rebind up without stepping in-between.
	sv.ComplexNode = v.(ComplexNode)

	if sv.Link == nil {
		// If nil, the view is maintaining the latest root binding
		return
	}
	sv.Link(v)
}

// Result will be nil if an error occurred.
func (sv *SubtreeView) Get(i uint64) (Node, error) {
	return sv.Getter(i, sv.depth)
}

// Result will be nil if an error occurred.
func (sv *SubtreeView) Set(i uint64) (Link, error) {
	return sv.Setter(i, sv.depth)
}

type ListView struct {
	SubtreeView
	limit uint64
}

type ListLength uint64

func (ll ListLength) ComputeRoot(h HashFn) (out Root) {
	binary.LittleEndian.PutUint64(out[:], uint64(ll))
	return
}

func List(limit uint64, nodes ...Node) (out *ListView) {
	elementCount := uint64(len(nodes))
	if elementCount > limit {
		// TODO: or add error return?
		return nil
	}
	depth := GetDepth(limit)
	out = &ListView{
		SubtreeView: SubtreeView{
			depth: depth + 1,
		},
		limit: limit,
	} // 1 extra for length mix-in

	mixin := ListLength(len(nodes))

	contents := &Commit{}
	contents.expandInplaceTo(nodes, depth)

	root := &Commit{Link: out.RebindInner, Left: contents, Right: mixin}
	contents.Bind(root.RebindLeft)

	out.ComplexNode = root
	return out
}

// Takes the backing, binds the view as a proxy link to keep track of changes,
// and defines a List based on a limit and a dynamic number of number of elements.
// The link however is optional: if nil, no changes are propagated upwards,
// but the view itself still tracks the latest changes.
func NewListView(backing ComplexRebindableNode, link Link, limit uint64) *ListView {
	view := &ListView{
		SubtreeView: SubtreeView{
			ComplexNode: backing,
			Link:        link,
			depth:       GetDepth(limit),
		},
		limit: limit,
	}
	backing.Bind(view.RebindInner)
	return view
}

// TODO: append/pop modify both contents as list-length: batching the changes would be good.

func (cv *ListView) Append(v Node) error {
	ll, err := cv.Length()
	if err != nil {
		return err
	}
	if ll >= cv.limit {
		return fmt.Errorf("list length is %d and appending would exceed the list limit %d", ll, cv.limit)
	}
	// Appending is done by setting the node at the index list-length. And expanding where necessary as it is being set.
	setLast, err := cv.ExpandInto(ll, cv.depth)
	if err != nil {
		return fmt.Errorf("failed to get a setter to append an item")
	}
	// Append the item by setting the newly allocated last item to it.
	setLast(v)
	// And update the list length
	setLength, err := cv.SubtreeView.Setter(1, 1)
	if err != nil {
		return err
	}
	setLength(ListLength(ll + 1))
	return nil
}

func (cv *ListView) Pop() error {
	ll, err := cv.Length()
	if err != nil {
		return err
	}
	if ll == 0 {
		return fmt.Errorf("list length is 0 and no item can be popped")
	}
	// Popping is done by setting the node at the index list-length. And expanding where necessary as it is being set.
	setLast, err := cv.ExpandInto(ll, cv.depth)
	if err != nil {
		return fmt.Errorf("failed to get a setter to append an item")
	}
	// Pop the item by setting it to the zero hash
	setLast(&ZeroHashes[0])
	// And update the list length
	setLength, err := cv.SubtreeView.Setter(1, 1)
	if err != nil {
		return err
	}
	setLength(ListLength(ll - 1))
	return nil
}

// Use .SubtreeView.Get(i) to work with the tree and get explicit tree errors instead of nil result.
func (cv *ListView) Get(i uint64) Node {
	ll, err := cv.Length()
	if err != nil || i >= ll {
		return nil
	}
	v, _ := cv.SubtreeView.Get(i)
	return v
}

// Use .SubtreeView.Set(i) to work with the tree and get explicit tree errors instead of nil result.
func (cv *ListView) Set(i uint64) Link {
	ll, err := cv.Length()
	if err != nil || i >= ll {
		return nil
	}
	v, _ := cv.SubtreeView.Set(i)
	return v
}

func (cv *ListView) Length() (uint64, error) {
	v, err := cv.SubtreeView.Getter(1, 1)
	if err != nil {
		return 0, err
	}
	ll, ok := v.(ListLength)
	if !ok {
		return 0, fmt.Errorf("cannot read node %v as list-length", v)
	}
	if uint64(ll) > cv.limit {
		return 0, fmt.Errorf("cannot read list length, length appears to be bigger than limit allows")
	}
	return uint64(ll), nil
}

type ContainerView struct {
	SubtreeView
	fieldCount uint64
}

func Container(nodes ...Node) (out *ContainerView) {
	elementCount := uint64(len(nodes))
	depth := GetDepth(elementCount)
	out = &ContainerView{
		SubtreeView: SubtreeView{
			depth: depth,
		},
		fieldCount: elementCount,
	}
	inner := &Commit{Link: out.RebindInner}
	inner.expandInplaceTo(nodes, depth)
	out.ComplexNode = inner
	return out
}

type ComplexRebindableNode interface {
	RebindableNode
	NodeInteraction
}

// Takes the backing, binds the view as a proxy link to keep track of changes,
// and defines a Container based on a fixed number of elements.
// The link however is optional: if nil, no changes are propagated upwards,
// but the view itself still tracks the latest changes.
func NewContainerView(backing ComplexRebindableNode, link Link, elementCount uint64) *ContainerView {
	view := &ContainerView{
		SubtreeView: SubtreeView{
			ComplexNode: backing,
			Link:        link,
			depth:       GetDepth(elementCount),
		},
		fieldCount: elementCount,
	}
	backing.Bind(view.RebindInner)
	return view
}

// Use .SubtreeView.Get(i) to work with the tree and get explicit tree errors instead of nil result.
func (cv *ContainerView) Get(i uint64) Node {
	if i >= cv.fieldCount {
		return nil
	}
	v, _ := cv.SubtreeView.Get(i)
	return v
}

// Use .SubtreeView.Set(i) to work with the tree and get explicit tree errors instead of nil result.
func (cv *ContainerView) Set(i uint64) Link {
	if i >= cv.fieldCount {
		return nil
	}
	v, _ := cv.SubtreeView.Set(i)
	return v
}

func (cv *ContainerView) Length() uint64 {
	return cv.fieldCount
}

// An immutable (L, R) pair with a link to the holding node.
// If L or R changes, the link is used to bind a new (L, *R) or (*L, R) pair in the holding value.
type Commit struct {
	// TODO: instead of value + bool, it could also be a pointer (nil case == computed false).
	//  But more objects/indirection/allocations.
	Value    Root
	computed bool // true if Value is set to H(L, R)
	Link     Link
	Left     Node
	Right    Node
}

func (c *Commit) Bind(bindingLink Link) {
	c.Link = bindingLink
}

func (c *Commit) ComputeRoot(h HashFn) Root {
	if c.computed {
		return c.Value
	}
	if c.Left == nil || c.Right == nil {
		panic("invalid state, cannot have left without right")
	}
	c.Value = h(c.Left.ComputeRoot(h), c.Right.ComputeRoot(h))
	c.computed = true
	return c.Value
}

func (c *Commit) RebindLeft(v Node) {
	if c.Link == nil {
		log.Println("cannot rebind from unbound node!")
		return
	}
	c.Link(&Commit{
		Value:    Root{},
		computed: false,
		Link:     c.Link,
		Left:     v,
		Right:    c.Right,
	})
}

func (c *Commit) RebindRight(v Node) {
	if c.Link == nil {
		log.Println("cannot rebind from unbound node!")
		return
	}
	c.Link(&Commit{
		Value:    Root{},
		computed: false,
		Link:     c.Link,
		Left:     c.Left,
		Right:    v,
	})
}

func (c *Commit) Expand() {
	next := &Commit{
		Value:    Root{},
		computed: false,
		Link:     c.Link,
		Left:     nil,
		Right:    nil,
	}
	left := &Commit{
		Value:    Root{},
		computed: false,
		Link:     next.RebindLeft,
		Left:     nil,
		Right:    nil,
	}
	right := &Commit{
		Value:    Root{},
		computed: false,
		Link:     next.RebindRight,
		Left:     nil,
		Right:    nil,
	}
	next.Left = left
	next.Right = right
	c.Link(next)
}

func sha256Combi(a Root, b Root) Root {
	v := [64]byte{}
	copy(v[:32], a[:])
	copy(v[32:], b[:])
	return sha256.Sum256(v[:])
}

const zeroHashesLevels = 64

var ZeroHashes []Root

// initialize the zero-hashes pre-computed data with the given hash-function.
func InitZeroHashes(h HashFn) {
	ZeroHashes = make([]Root, zeroHashesLevels+1)
	for i := 0; i < zeroHashesLevels; i++ {
		ZeroHashes[i+1] = h(ZeroHashes[i], ZeroHashes[i])
	}
}

func init() {
	InitZeroHashes(sha256Combi)
}

// Unsafe! Modifies L and R, without triggering a rebind in the parent
func (c *Commit) expandInplaceTo(nodes []Node, depth uint8) {
	c.computed = false
	if depth == 0 {
		panic("invalid usage")
	}
	if depth == 1 {
		c.Left = nodes[0]
		if r, ok := nodes[0].(RebindableNode); ok {
			r.Bind(c.RebindLeft)
		}
		if len(nodes) > 1 {
			c.Right = nodes[1]
			if r, ok := nodes[1].(RebindableNode); ok {
				r.Bind(c.RebindRight)
			}
		} else {
			c.Right = &ZeroHashes[0]
		}
	} else {
		pivot := uint64(1) << depth
		c.Left = &Commit{
			Value:    Root{},
			computed: false,
			Link:     c.RebindLeft,
			Left:     nil,
			Right:    nil,
		}
		if uint64(len(nodes)) <= pivot {
			c.Left.(*Commit).expandInplaceTo(nodes, depth-1)
			c.Right = &ZeroHashes[depth]
		} else {
			c.Left.(*Commit).expandInplaceTo(nodes[:pivot], depth-1)
			c.Right = &Commit{
				Value:    Root{},
				computed: false,
				Link:     c.RebindRight,
				Left:     nil,
				Right:    nil,
			}
			c.Right.(*Commit).expandInplaceTo(nodes[pivot:], depth-1)
		}
	}
}

func (c *Commit) Getter(target uint64, depth uint8) (Node, error) {
	if depth == 0 {
		return c, nil
	}
	if depth == 1 {
		if target == 0 {
			return c.Left, nil
		}
		if target == 1 {
			return c.Right, nil
		}
	}
	if pivot := uint64(1) << depth; target < pivot {
		if c.Left == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no left node", target, depth)
		}
		if left, ok := c.Left.(GetterInteraction); ok {
			return left.Getter(target, depth-1)
		} else {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: left node has no GetterInteraction", target, depth)
		}
	} else {
		if c.Right == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no right node", target, depth)
		}
		if right, ok := c.Right.(GetterInteraction); ok {
			return right.Getter(target&^pivot, depth-1)
		} else {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: right node has no GetterInteraction", target, depth)
		}
	}
}

func (c *Commit) ExpandInto(target uint64, depth uint8) (Link, error) {
	if depth == 0 {
		return c.Link, nil
	}
	if depth == 1 {
		if target == 0 {
			return c.RebindLeft, nil
		}
		if target == 1 {
			return c.RebindRight, nil
		}
	}
	if pivot := uint64(1) << depth; target < pivot {
		if c.Left == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no left node", target, depth)
		}
		if left, ok := c.Left.(ExpandIntoInteraction); ok {
			return left.ExpandInto(target, depth-1)
		} else {
			// stop immediate propagation of rebinds during this Set call.
			var tmp Node

			startC := &Commit{
				Link: func(v Node) {
					tmp = v
				},
				Left:     &ZeroHashes[depth-2],
				Right:    &ZeroHashes[depth-2],
			}
			tmp = startC
			// Get the setter, recurse into the new node
			l, err := startC.ExpandInto(target, depth-1)

			newLeftC := tmp.(*Commit)
			// Now update the link to attach the updates of the new left node
			newLeftC.Link = c.RebindLeft
			// And attach as the left node
			c.RebindLeft(newLeftC)
			return l, err
		}
	} else {
		if c.Right == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no right node", target, depth)
		}
		if right, ok := c.Right.(ExpandIntoInteraction); ok {
			return right.ExpandInto(target&^pivot, depth-1)
		} else {
			// stop immediate propagation of rebinds during this Set call.
			var tmp Node

			startC := &Commit{
				Link: func(v Node) {
					tmp = v
				},
				Left:  &ZeroHashes[depth-1],
				Right: &ZeroHashes[depth-1],
			}
			tmp = startC
			// Get the setter, recurse into the new node
			l, err := startC.ExpandInto(target&^pivot, depth-1)

			newRightC := tmp.(*Commit)
			// Now update the link to attach the updates of the new right node
			newRightC.Link = c.RebindRight
			// And attach as the right node
			c.RebindRight(newRightC)
			return l, err
		}
	}
}

func (c *Commit) Setter(target uint64, depth uint8) (Link, error) {
	if depth == 0 {
		return c.Link, nil
	}
	if depth == 1 {
		if target == 0 {
			return c.RebindLeft, nil
		}
		if target == 1 {
			return c.RebindRight, nil
		}
	}
	if pivot := uint64(1) << depth; target < pivot {
		if c.Left == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no left node", target, depth)
		}
		if left, ok := c.Left.(SetterInteraction); ok {
			return left.Setter(target, depth-1)
		} else {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: left node has no SetterInteraction", target, depth)
		}
	} else {
		if c.Right == nil {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: no right node", target, depth)
		}
		if right, ok := c.Right.(SetterInteraction); ok {
			return right.Setter(target&^pivot, depth-1)
		} else {
			return nil, fmt.Errorf("cannot find node at target %v in depth %v: right node has no SetterInteraction", target, depth)
		}
	}
}

// Temporarily decouples the commit from its parent to scope modifications within the "work" callback.
// Then rebinds to the parent, and propagates changes if there were any.
func (c *Commit) Batch(work func()) {
	link := c.Link
	var next Node
	c.Link = func(v Node) {
		next = v
	}
	work()
	if next != nil {
		link(next)
	}
	c.Link = link
}

// Example usage:
//  - anything can be a node
//  - commits are wrapped with navigation to fetch getter/setter pairs of containers.
//  - commits can be made read-only
//  - modifications can be batched

type BLSSignature [96]byte

func (v BLSSignature) ComputeRoot(h HashFn) (out Root) {
	return h.MerkleizeBytes(v[:])
}

type Slot uint64

func (s Slot) ComputeRoot(h HashFn) (out Root) {
	binary.LittleEndian.PutUint64(out[:], uint64(s))
	return
}

type Block struct {
	*ContainerView
	// Not included in default hash-tree-root
	Signature core.BLSSignature
}

func NewBlock() (b *Block) {
	return &Block{ContainerView: Container(
		Slot(0),
		&Root{},
		&Root{},
		NewBlockBody(),
	)}
}
func (b *Block) Slot() Slot { return b.Get(0).(Slot) }

type BlockBody struct {
	*ContainerView
}

func NewBlockBody() (b *BlockBody) {
	return &BlockBody{Container(
		Slot(0),
		&Root{},
		// TODO operations lists
	)}
}

func main() {
	b := NewBlock()
	b.Set(0)(&Root{1})
	fmt.Printf("%x", b.ComputeRoot(sha256Combi))
}

//
//func (b *Block) Body() (*Body, Link) {
//	getter, setter := b.Navigate(1, 8)
//	return (*Body)(getter.(*Commit)), setter
//}