github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/tree/tree_arenas_enabled.go (about)

     1  //go:build goexperiment.arenas
     2  
     3  package tree
     4  
     5  import (
     6  	"arena"
     7  	"bytes"
     8  	"github.com/pyroscope-io/pyroscope/pkg/util/arenahelper"
     9  	"sort"
    10  )
    11  
    12  func (t *Tree) InsertStackA(stack [][]byte, v uint64) {
    13  	a := t.arena
    14  	if a == nil {
    15  		t.InsertStack(stack, v)
    16  		return
    17  	}
    18  	n := t.root
    19  	for j := range stack {
    20  		n.Total += v
    21  		n = n.insertA(a, stack[j])
    22  	}
    23  	// Leaf.
    24  	n.Total += v
    25  	n.Self += v
    26  }
    27  
    28  func (n *treeNode) insertA(a arenahelper.ArenaWrapper, targetLabel []byte) *treeNode {
    29  	i := sort.Search(len(n.ChildrenNodes), func(i int) bool {
    30  		return bytes.Compare(n.ChildrenNodes[i].Name, targetLabel) >= 0
    31  	})
    32  	if i > len(n.ChildrenNodes)-1 || !bytes.Equal(n.ChildrenNodes[i].Name, targetLabel) {
    33  		l := arena.MakeSlice[byte](a, len(targetLabel), len(targetLabel))
    34  		copy(l, targetLabel)
    35  		child := newNodeA(a, l)
    36  		if len(n.ChildrenNodes) < cap(n.ChildrenNodes) {
    37  			n.ChildrenNodes = append(n.ChildrenNodes, child)
    38  		} else {
    39  			n.ChildrenNodes = arenahelper.AppendA(n.ChildrenNodes, child, a)
    40  		}
    41  		copy(n.ChildrenNodes[i+1:], n.ChildrenNodes[i:])
    42  		n.ChildrenNodes[i] = child
    43  	}
    44  	return n.ChildrenNodes[i]
    45  }
    46  
    47  func newNodeA(a arenahelper.ArenaWrapper, label []byte) *treeNode {
    48  	n := arena.New[treeNode](a)
    49  	n.Name = label
    50  	return n
    51  }
    52  
    53  func NewA(a arenahelper.ArenaWrapper) *Tree {
    54  	t := arena.New[Tree](a)
    55  	t.root = newNodeA(a, nil)
    56  	t.arena = a
    57  	return t
    58  }