diff --git a/btf/types.go b/btf/types.go index 70123cf9e..a168eeca8 100644 --- a/btf/types.go +++ b/btf/types.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "math" + "math/bits" "reflect" "strings" @@ -748,77 +749,91 @@ func (c copier) copy(typ *Type, transform Transformer) { } } -// typeDeque keeps track of pointers to types which still -// need to be visited. -type typeDeque struct { - types []*Type +type typeDeque = deque[*Type] + +// deque implements a double ended queue. +type deque[T any] struct { + elems []T read, write uint64 mask uint64 } -func (dq *typeDeque) empty() bool { +func (dq *deque[T]) empty() bool { return dq.read == dq.write } -// push adds a type to the stack. -func (dq *typeDeque) push(t *Type) { - if dq.write-dq.read < uint64(len(dq.types)) { - dq.types[dq.write&dq.mask] = t +func (dq *deque[T]) remainingCap() int { + return len(dq.elems) - int(dq.write-dq.read) +} + +// push adds an element to the end. +func (dq *deque[T]) push(e T) { + if dq.remainingCap() >= 1 { + dq.elems[dq.write&dq.mask] = e dq.write++ return } - new := len(dq.types) * 2 - if new == 0 { - new = 8 - } + elems := dq.linearise(1) + elems = append(elems, e) - types := make([]*Type, new) - pivot := dq.read & dq.mask - n := copy(types, dq.types[pivot:]) - n += copy(types[n:], dq.types[:pivot]) - types[n] = t - - dq.types = types - dq.mask = uint64(new) - 1 - dq.read, dq.write = 0, uint64(n+1) + dq.elems = elems[:cap(elems)] + dq.mask = uint64(cap(elems)) - 1 + dq.read, dq.write = 0, uint64(len(elems)) } -// shift returns the first element or null. -func (dq *typeDeque) shift() *Type { +// shift returns the first element or the zero value. +func (dq *deque[T]) shift() T { + var zero T + if dq.empty() { - return nil + return zero } index := dq.read & dq.mask - t := dq.types[index] - dq.types[index] = nil + t := dq.elems[index] + dq.elems[index] = zero dq.read++ return t } -// pop returns the last element or null. -func (dq *typeDeque) pop() *Type { +// pop returns the last element or the zero value. +func (dq *deque[T]) pop() T { + var zero T + if dq.empty() { - return nil + return zero } dq.write-- index := dq.write & dq.mask - t := dq.types[index] - dq.types[index] = nil + t := dq.elems[index] + dq.elems[index] = zero return t } -// all returns all elements. +// linearise the contents of the deque. // -// The deque is empty after calling this method. -func (dq *typeDeque) all() []*Type { +// The returned slice has space for at least n more elements and has power +// of two capacity. +func (dq *deque[T]) linearise(n int) []T { length := dq.write - dq.read - types := make([]*Type, 0, length) - for t := dq.shift(); t != nil; t = dq.shift() { - types = append(types, t) + need := length + uint64(n) + if need < length { + panic("overflow") } + + // Round up to the new power of two which is at least 8. + // See https://jameshfisher.com/2018/03/30/round-up-power-2/ + capacity := 1 << (64 - bits.LeadingZeros64(need-1)) + if capacity < 8 { + capacity = 8 + } + + types := make([]T, length, capacity) + pivot := dq.read & dq.mask + copied := copy(types, dq.elems[pivot:]) + copy(types[copied:], dq.elems[:pivot]) return types } diff --git a/btf/types_test.go b/btf/types_test.go index 23b6fa86b..3d6076df6 100644 --- a/btf/types_test.go +++ b/btf/types_test.go @@ -160,77 +160,78 @@ func TestType(t *testing.T) { } } -func TestTypeDeque(t *testing.T) { - a, b := new(Type), new(Type) - +func TestDeque(t *testing.T) { t.Run("pop", func(t *testing.T) { - var td typeDeque - td.push(a) - td.push(b) + var dq deque[int] + dq.push(1) + dq.push(2) - if td.pop() != b { - t.Error("Didn't pop b first") + if dq.pop() != 2 { + t.Error("Didn't pop 2 first") } - if td.pop() != a { - t.Error("Didn't pop a second") + if dq.pop() != 1 { + t.Error("Didn't pop 1 second") } - if td.pop() != nil { - t.Error("Didn't pop nil") + if dq.pop() != 0 { + t.Error("Didn't pop zero") } }) t.Run("shift", func(t *testing.T) { - var td typeDeque - td.push(a) - td.push(b) + var td deque[int] + td.push(1) + td.push(2) - if td.shift() != a { - t.Error("Didn't shift a second") + if td.shift() != 1 { + t.Error("Didn't shift 1 first") } - if td.shift() != b { - t.Error("Didn't shift b first") + if td.shift() != 2 { + t.Error("Didn't shift b second") } - if td.shift() != nil { - t.Error("Didn't shift nil") + if td.shift() != 0 { + t.Error("Didn't shift zero") } }) t.Run("push", func(t *testing.T) { - var td typeDeque - td.push(a) - td.push(b) + var td deque[int] + td.push(1) + td.push(2) td.shift() - ts := make([]Type, 12) - for i := range ts { - td.push(&ts[i]) + for i := 1; i <= 12; i++ { + td.push(i) } - if td.shift() != b { - t.Error("Didn't shift b first") + if td.shift() != 2 { + t.Error("Didn't shift 2 first") } - for i := range ts { - if td.shift() != &ts[i] { - t.Fatal("Shifted wrong Type at pos", i) + for i := 1; i <= 12; i++ { + if v := td.shift(); v != i { + t.Fatalf("Shifted %d at pos %d", v, i) } } }) - t.Run("all", func(t *testing.T) { - var td typeDeque - td.push(a) - td.push(b) + t.Run("linearise", func(t *testing.T) { + var td deque[int] + td.push(1) + td.push(2) - all := td.all() + all := td.linearise(0) if len(all) != 2 { t.Fatal("Expected 2 elements, got", len(all)) } - if all[0] != a || all[1] != b { + if cap(all)&(cap(all)-1) != 0 { + t.Fatalf("Capacity %d is not a power of two", cap(all)) + } + + if all[0] != 1 || all[1] != 2 { t.Fatal("Elements don't match") } })