btf: change Spec.Add semantics WIP

Signed-off-by: Lorenz Bauer <i@lmb.io>

btf/btf.go

@@ -510,7 +510,7 @@ func (sw sliceWriter) Write(p []byte) (int, error) {
 	return copy(sw, p), nil
 }
 
-// Add a Type.
+// Add a Type to the Spec, making it queryable via [TypeByName], etc.
 //
 // Adding the identical Type multiple times is valid and will return a stable ID.
 //
@@ -520,29 +520,24 @@ func (s *Spec) Add(typ Type) (TypeID, error) {
 		return 0, fmt.Errorf("can't add nil Type")
 	}
 
-	hasID := func(t Type) (skip bool) {
-		_, isVoid := t.(*Void)
-		_, alreadyEncoded := s.typeIDs[t]
-		return isVoid || alreadyEncoded
+	if id, err := s.TypeID(typ); err == nil {
+		return id, nil
 	}
 
-	iter := postorderTraversal(typ, hasID)
-	for iter.Next() {
-		id := s.lastTypeID + 1
-		if id < s.lastTypeID {
-			return 0, fmt.Errorf("type ID overflow")
-		}
+	id := s.lastTypeID + 1
+	if id < s.lastTypeID {
+		return 0, fmt.Errorf("type ID overflow")
+	}
 
-		s.typeIDs[iter.Type] = id
-		s.types = append(s.types, iter.Type)
-		s.lastTypeID = id
+	s.typeIDs[typ] = id
+	s.types = append(s.types, typ)
+	s.lastTypeID = id
 
-		if name := newEssentialName(iter.Type.TypeName()); name != "" {
-			s.namedTypes[name] = append(s.namedTypes[name], iter.Type)
-		}
+	if name := newEssentialName(typ.TypeName()); name != "" {
+		s.namedTypes[name] = append(s.namedTypes[name], typ)
 	}
 
-	return s.TypeID(typ)
+	return id, nil
 }
 
 // TypeByID returns the BTF Type with the given type ID.

btf/btf_test.go

@@ -224,31 +224,36 @@ func TestSpecAdd(t *testing.T) {
 		Size:     2,
 		Encoding: Signed | Char,
 	}
+	pi := &Pointer{i}
 
 	s := NewSpec()
-	id, err := s.Add(i)
+	id, err := s.Add(pi)
 	qt.Assert(t, err, qt.IsNil)
 	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("First non-void type doesn't get id 1"))
-	id, err = s.Add(i)
-	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(1), qt.Commentf("Adding a type twice returns different ids"))
 
-	id, err = s.TypeID(i)
+	id, err = s.Add(pi)
 	qt.Assert(t, err, qt.IsNil)
 	qt.Assert(t, id, qt.Equals, TypeID(1))
 
-	id, err = s.Add(&Pointer{i})
+	_, err = s.TypeID(i)
+	qt.Assert(t, err, qt.IsNotNil, qt.Commentf("Children mustn't be added"))
+
+	id, err = s.Add(i)
 	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(2))
+	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Second type doesn't get id 2"))
 
-	id, err = s.Add(&Typedef{"baz", i})
+	id, err = s.Add(i)
 	qt.Assert(t, err, qt.IsNil)
-	qt.Assert(t, id, qt.Equals, TypeID(3))
+	qt.Assert(t, id, qt.Equals, TypeID(2), qt.Commentf("Adding a type twice returns different ids"))
 
 	typ, err := s.AnyTypeByName("foo")
-	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, err, qt.IsNil, qt.Commentf("Add doesn't make named type queryable"))
 	qt.Assert(t, typ, qt.Equals, i)
 
+	id, err = s.Add(&Typedef{"baz", i})
+	qt.Assert(t, err, qt.IsNil)
+	qt.Assert(t, id, qt.Equals, TypeID(3))
+
 	_, err = s.AnyTypeByName("baz")
 	qt.Assert(t, err, qt.IsNil)
 }
@@ -365,7 +370,7 @@ func TestLoadSpecFromElf(t *testing.T) {
 
 func TestVerifierError(t *testing.T) {
 	var buf bytes.Buffer
-	if err := marshalSpec(&buf, NewSpec(), nil, nil); err != nil {
+	if err := marshalTypes(&buf, []Type{&Void{}}, nil, nil); err != nil {
 		t.Fatal(err)
 	}
 

btf/ext_info.go

@@ -185,7 +185,7 @@ marshal:
 	buf := getBuffer()
 	defer putBuffer(buf)
 
-	if err := marshalSpec(buf, spec, stb, kernelMarshalOptions); err != nil {
+	if err := marshalTypes(buf, spec.types, stb, kernelMarshalOptions); err != nil {
 		return nil, nil, nil, fmt.Errorf("marshal BTF: %w", err)
 	}
 

btf/handle.go

@@ -45,7 +45,7 @@ func NewHandle(spec *Spec) (*Handle, error) {
 		stb = newStringTableBuilder(spec.strings.Num())
 	}
 
-	err := marshalSpec(buf, spec, stb, kernelMarshalOptions)
+	err := marshalTypes(buf, spec.types, stb, kernelMarshalOptions)
 	if err != nil {
 		return nil, fmt.Errorf("marshal BTF: %w", err)
 	}

btf/marshal.go

@@ -3,6 +3,7 @@ package btf
 import (
 	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"math"
 	"sync"
@@ -25,9 +26,11 @@ type encoder struct {
 	marshalOptions
 
 	byteOrder binary.ByteOrder
+	pending   internal.Deque[Type]
 	buf       *bytes.Buffer
 	strings   *stringTableBuilder
 	ids       map[Type]TypeID
+	lastID    TypeID
 }
 
 var emptyBTFHeader = make([]byte, btfHeaderLen)
@@ -48,26 +51,25 @@ func putBuffer(buf *bytes.Buffer) {
 	bufferPool.Put(buf)
 }
 
-// marshalSpec encodes a spec into BTF wire format.
+// marshalTypes encodes a slice of types into BTF wire format.
+//
+// types are guaranteed to be written in the order they are passed to this
+// function. The first type must always be Void.
 //
 // Doesn't support encoding split BTF since it's not possible to load
 // that into the kernel and we don't have a use case for writing BTF
 // out again.
 //
 // w should be retrieved from bufferPool. opts may be nil.
-func marshalSpec(w *bytes.Buffer, s *Spec, stb *stringTableBuilder, opts *marshalOptions) (err error) {
-	defer func() {
-		r := recover()
-		if r == nil {
-			return
-		}
+func marshalTypes(w *bytes.Buffer, types []Type, stb *stringTableBuilder, opts *marshalOptions) error {
+	if len(types) < 1 {
+		return errors.New("types must contain at least Void")
+	}
 
-		var ok bool
-		err, ok = r.(error)
-		if !ok {
-			panic(r)
-		}
-	}()
+	if _, ok := types[0].(*Void); !ok {
+		return fmt.Errorf("first type is %s, not Void", types[0])
+	}
+	types = types[1:]
 
 	if stb == nil {
 		stb = newStringTableBuilder(0)
@@ -77,20 +79,27 @@ func marshalSpec(w *bytes.Buffer, s *Spec, stb *stringTableBuilder, opts *marsha
 		byteOrder: internal.NativeEndian,
 		buf:       w,
 		strings:   stb,
-		ids:       s.typeIDs,
+		ids:       make(map[Type]TypeID, len(types)),
 	}
 
 	if opts != nil {
 		e.marshalOptions = *opts
 	}
 
+	// Ensure that passed types are marshaled in the exact order they were
+	// passed.
+	e.pending.Grow(len(types))
+	for _, typ := range types {
+		if err := e.allocateID(typ); err != nil {
+			return err
+		}
+	}
+
 	// Reserve space for the BTF header.
 	_, _ = e.buf.Write(emptyBTFHeader)
 
-	for _, typ := range s.types {
-		if err := e.deflateType(typ); err != nil {
-			return fmt.Errorf("deflate %s: %w", typ, err)
-		}
+	if err := e.deflatePending(); err != nil {
+		return err
 	}
 
 	length := e.buf.Len()
@@ -118,14 +127,26 @@ func marshalSpec(w *bytes.Buffer, s *Spec, stb *stringTableBuilder, opts *marsha
 		StringLen: uint32(stringLen),
 	}
 
-	err = binary.Write(sliceWriter(buf[:btfHeaderLen]), e.byteOrder, header)
+	err := binary.Write(sliceWriter(buf[:btfHeaderLen]), e.byteOrder, header)
 	if err != nil {
 		return fmt.Errorf("write header: %v", err)
 	}
 
 	return nil
 }
 
+func (e *encoder) allocateID(typ Type) error {
+	id := e.lastID + 1
+	if id < e.lastID {
+		return errors.New("type ID overflow")
+	}
+
+	e.pending.Push(typ)
+	e.ids[typ] = id
+	e.lastID = id
+	return nil
+}
+
 // id returns the ID for the given type or panics with an error.
 func (e *encoder) id(typ Type) TypeID {
 	if _, ok := typ.(*Void); ok {
@@ -134,13 +155,64 @@ func (e *encoder) id(typ Type) TypeID {
 
 	id, ok := e.ids[typ]
 	if !ok {
-		panic(fmt.Errorf("no ID for type %s", typ))
+		panic(fmt.Errorf("no ID for type %v", typ))
 	}
 
 	return id
 }
 
+func (e *encoder) deflatePending() error {
+	// Declare root outside of the loop to avoid repeated heap allocations.
+	var root Type
+	skip := func(t Type) (skip bool) {
+		if t == root {
+			// Force descending into the current root type even if it already
+			// has an ID. Otherwise we miss children of types that have their
+			// ID pre-allocated in marshalTypes.
+			return false
+		}
+
+		_, isVoid := t.(*Void)
+		_, alreadyEncoded := e.ids[t]
+		return isVoid || alreadyEncoded
+	}
+
+	for !e.pending.Empty() {
+		root = e.pending.Shift()
+
+		// Allocate IDs for all children of typ, including transitive dependencies.
+		iter := postorderTraversal(root, skip)
+		for iter.Next() {
+			if iter.Type == root {
+				// The iterator yields root at the end, do not allocate another ID.
+				break
+			}
+
+			if err := e.allocateID(iter.Type); err != nil {
+				return err
+			}
+		}
+
+		if err := e.deflateType(root); err != nil {
+			id := e.ids[root]
+			return fmt.Errorf("deflate %v with ID %d: %w", root, id, err)
+		}
+	}
+
+	return nil
+}
+
 func (e *encoder) deflateType(typ Type) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			var ok bool
+			err, ok = r.(error)
+			if !ok {
+				panic(r)
+			}
+		}
+	}()
+
 	var raw rawType
 	raw.NameOff, err = e.strings.Add(typ.TypeName())
 	if err != nil {
@@ -149,7 +221,7 @@ func (e *encoder) deflateType(typ Type) (err error) {
 
 	switch v := typ.(type) {
 	case *Void:
-		return nil
+		return errors.New("Void is implicit in BTF wire format")
 
 	case *Int:
 		raw.SetKind(kindInt)