Make stack implementation more generic and move it to pkg/util

pkg/util/stack.go

@@ -0,0 +1,44 @@
+package util
+
+// Stack is a generic stack implementation that supports len() and cap().
+type Stack[T any] []T
+
+// Push pushes a new element on the stack.
+func (s *Stack[T]) Push(element T) {
+ *s = append(*s, element)
+}
+
+// Peek returns the top element from the stack without removing it. If the stack is
+// empty the second return value is false.
+func (s *Stack[T]) Peek() (T, bool) {
+ if len(*s) == 0 {
+ var zero T
+ return zero, false
+ }
+ return (*s)[len(*s)-1], true
+}
+
+// Pop returns the top element from the stack and removes it. If the stack is empty
+// the second return value is false.
+func (s *Stack[T]) Pop() (T, bool) {
+ if len(*s) == 0 {
+ var zero T
+ return zero, false
+ }
+ i := len(*s) - 1
+
+ element := (*s)[i]
+ *s = (*s)[:i]
+
+ return element, true
+}
+
+// IsEmpty returns whether the stack contains elements
+func (s *Stack[T]) IsEmpty() bool {
+ return len(*s) == 0
+}
+
+// Reset empties the stack while maintaining its capacity
+func (s *Stack[T]) Reset() {
+ *s = (*s)[:0]
+}

pkg/util/stack_test.go

@@ -0,0 +1,55 @@
+package util
+
+import (
+ "testing"
+
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestStack(t *testing.T) {
+ type testVal struct {
+ val int
+ }
+
+ var (
+ testStack Stack[*testVal]
+ val *testVal
+ )
+
+ assert.Len(t, testStack, 0)
+ assert.True(t, testStack.IsEmpty(), "testStack expected to be empty")
+
+ _, ok := testStack.Peek()
+ assert.False(t, ok, "testStack.Peek() expected to return false")
+ _, ok = testStack.Pop()
+ assert.False(t, ok, "testStack.Pop() expected to return false")
+
+ testStack.Push(&testVal{1})
+ val, ok = testStack.Peek()
+ require.True(t, ok)
+ assert.Equal(t, &testVal{1}, val)
+
+ testStack.Push(&testVal{2})
+ val, ok = testStack.Peek()
+ require.True(t, ok)
+ assert.Equal(t, &testVal{2}, val)
+
+ val, ok = testStack.Pop()
+ require.True(t, ok)
+ assert.Equal(t, &testVal{2}, val)
+
+ val, ok = testStack.Pop()
+ require.True(t, ok)
+ assert.Equal(t, &testVal{1}, val)
+
+ _, ok = testStack.Pop()
+ assert.True(t, testStack.IsEmpty(), "testStack expected to be empty")
+ assert.False(t, ok, "testStack.Pop() expected to return false")
+
+ testStack.Push(&testVal{1})
+ testStack.Reset()
+ assert.True(t, testStack.IsEmpty(), "testStack expected to be empty")
+ assert.Len(t, testStack, 0)
+ assert.Greater(t, cap(testStack), 0)
+}

tempodb/encoding/vparquet2/nested_set_model.go

@@ -38,19 +38,19 @@ func assignNestedSetModelBounds(trace *Trace, forceAssignment bool) {
  // Traverse the tree depth first. When traversing down into the tree, assign NestedSetLeft
  // and assign NestedSetRight when going up.
  var (
- ancestors stack[wrappedSpan]
+ ancestors util.Stack[*wrappedSpan]
  nestedSetBound int32 = 1
  )
 
  for _, root := range rootSpans {
  root.span.NestedSetLeft = nestedSetBound
  nestedSetBound++
 
- ancestors.reset()
- ancestors.push(root)
+ ancestors.Reset()
+ ancestors.Push(root)
 
- for !ancestors.isEmpty() {
- parent := ancestors.peek()
+ for !ancestors.IsEmpty() {
+ parent, _ := ancestors.Peek()
  children := spanChildren[parent.id]
 
  if parent.nextChild < len(children) {
@@ -63,52 +63,23 @@ func assignNestedSetModelBounds(trace *Trace, forceAssignment bool) {
  child.NestedSetLeft = nestedSetBound
  nestedSetBound++
 
- ancestors.push(&wrappedSpan{span: child, id: util.SpanIDToArray(child.SpanID)})
+ ancestors.Push(&wrappedSpan{span: child, id: util.SpanIDToArray(child.SpanID)})
  } else {
  // All children of the current node were visited: go up
 
  parent.span.NestedSetRight = nestedSetBound
  nestedSetBound++
 
- ancestors.pop()
+ ancestors.Pop()
  }
  }
  }
 }
 
+// wrappedSpan is used to remember the converted span ID and position of the child that
+// needs to be visited next
 type wrappedSpan struct {
  span *Span
  id [8]byte
  nextChild int
 }
-
-type stack[T any] []*T
-
-func (ss *stack[T]) push(element *T) {
- *ss = append(*ss, element)
-}
-
-func (ss *stack[T]) peek() *T {
- if len(*ss) == 0 {
- return nil
- }
- return (*ss)[len(*ss)-1]
-}
-
-func (ss *stack[T]) pop() *T {
- if len(*ss) == 0 {
- return nil
- }
- i := len(*ss) - 1
- s := (*ss)[i]
- *ss = (*ss)[:i]
- return s
-}
-
-func (ss *stack[T]) isEmpty() bool {
- return len(*ss) == 0
-}
-
-func (ss *stack[T]) reset() {
- *ss = (*ss)[:0]
-}

tempodb/encoding/vparquet2/nested_set_model_test.go

@@ -234,51 +234,3 @@ func assertEqualNestedSetModelBounds(t testing.TB, actual, expected *Trace) {
 
  assert.Equalf(t, expectedCount, actualCount, "expected %d spans but found %d instead", expectedCount, actualCount)
 }
-
-func TestStack(t *testing.T) {
- type testVal struct {
- val int
- }
-
- var (
- testStack stack[testVal]
- val *testVal
- )
-
- assert.Len(t, testStack, 0)
- assert.True(t, testStack.isEmpty(), "testStack expected to be empty")
-
- val = testStack.peek()
- assert.Nil(t, val, "testStack.peek() expected to be nil")
- val = testStack.pop()
- assert.Nil(t, val, "testStack.pop() expected to be nil")
-
- testStack.push(&testVal{1})
- val = testStack.peek()
- require.NotNil(t, val)
- assert.Equal(t, &testVal{1}, val)
-
- testStack.push(&testVal{2})
- val = testStack.peek()
- require.NotNil(t, val)
- assert.Equal(t, &testVal{2}, val)
-
- val = testStack.pop()
- require.NotNil(t, val)
- assert.Equal(t, &testVal{2}, val)
-
- val = testStack.pop()
- require.NotNil(t, val)
- assert.Equal(t, &testVal{1}, val)
-
- val = testStack.pop()
- assert.True(t, testStack.isEmpty(), "testStack expected to be empty")
- assert.Nil(t, val, "testStack.peek() expected to be nil")
- assert.Nil(t, val, "testStack.pop() expected to be nil")
-
- testStack.push(&testVal{1})
- testStack.reset()
- assert.True(t, testStack.isEmpty(), "testStack expected to be empty")
- assert.Len(t, testStack, 0)
- assert.Greater(t, cap(testStack), 0)
-}