Repository: cilium/ebpf
Pull request #1397 (2024-03-26 11:08:56)

btf: use recursion
Reduce iterator complexity in the btf package by using recursion. If you have played with the rangefunc proposal you will realise that `visitPostorder,` etc. can be turned into an `iter.Seq` with a few lines of code. The reason to make this change now (rather than waiting) is that the recursive implementation is easier to reason about and about as fast:

```
core: 1
goos: linux
goarch: amd64
pkg: github.com/cilium/ebpf/btf
cpu: 12th Gen Intel(R) Core(TM) i7-1260P
                                       │  base.txt   │            recursive.txt             │
                                       │   sec/op    │   sec/op     vs base                 │
ParseVmlinux                             64.06m ± 2%   64.43m ± 2%        ~ (p=0.394 n=6)
SpecCopy                                 353.2n ± 1%   270.3n ± 2%  -23.46% (p=0.002 n=6)
SpecTypeByID                             35.99n ± 0%   35.98n ± 0%   -0.03% (p=0.032 n=6)
CORESkBuff/byte_off                      1.663µ ± 0%   1.670µ ± 1%        ~ (p=0.255 n=6)
CORESkBuff/byte_sz                       1.684µ ± 0%   1.699µ ± 0%   +0.89% (p=0.002 n=6)
CORESkBuff/field_exists                  1.660µ ± 0%   1.673µ ± 0%   +0.75% (p=0.002 n=6)
CORESkBuff/signed                        1.676µ ± 0%   1.684µ ± 0%   +0.45% (p=0.002 n=6)
CORESkBuff/lshift_u64                    1.683µ ± 1%   1.691µ ± 1%        ~ (p=0.104 n=6)
CORESkBuff/rshift_u64                    1.682µ ± 0%   1.686µ ± 1%        ~ (p=0.058 n=6)
CORESkBuff/local_type_id                 123.2n ± 1%   122.6n ± 1%   -0.53% (p=0.013 n=6)
CORESkBuff/target_type_id                750.8n ± 1%   776.9n ± 0%   +3.49% (p=0.002 n=6)
CORESkBuff/type_exists                   755.5n ± 1%   779.5n ± 1%   +3.18% (p=0.002 n=6)
CORESkBuff/type_size                     777.9n ± 1%   814.1n ± 1%   +4.65% (p=0.002 n=6)
CORESkBuff/enumval_exists                744.4n ± 1%   780.1n ± 1%   +4.80% (p=0.002 n=6)
CORESkBuff/enumval_value                 744.3n ± 1%   774.4n ± 2%   +4.04% (p=0.002 n=6)
Marshaler                                13.64m ± 0%   11.73m ± 0%  -14.01% (p=0.002 n=6)
BuildVmlinux                             122.3m ± 2%   129.9m ± 1%   +6.25% (p=0.002 n=6)
StringTableZeroLookup                    2.885n ± 0%   2.885n ± 0%        ~ (p=0.574 n=6)
PostorderTraversal/single_type           32.20n ± 1%   39.69n ± 0%  +23.24% (p=0.002 n=6)
PostorderTraversal/cycle(1)              565.6n ± 1%   112.3n ± 0%  -80.14% (p=0.002 n=6)
PostorderTraversal/cycle(10)             2.730µ ± 1%   1.578µ ± 1%  -42.22% (p=0.002 n=6)
PostorderTraversal/gov_update_cpu_data   2.001m ± 1%   1.795m ± 1%  -10.32% (p=0.002 n=6)
PreorderTraversal/single_type            36.87n ± 0%
PreorderTraversal/cycle(1)               103.8n ± 1%
PreorderTraversal/cycle(10)              1.926µ ± 0%
PreorderTraversal/gov_update_cpu_data    2.181m ± 1%
Copy                                     5.084µ ± 1%   3.291µ ± 2%  -35.28% (p=0.002 n=6)
Walk/Void                                4.327n ± 1%   4.337n ± 0%        ~ (p=0.061 n=6)
Walk/Int[unsigned_size=0]                4.327n ± 0%   4.340n ± 0%   +0.31% (p=0.002 n=6)
Walk/Pointer[target=<nil>]               109.9n ± 2%   110.7n ± 2%        ~ (p=0.327 n=6)
Walk/Array[index=<nil>_type=<nil>_n=0]   119.2n ± 1%   118.0n ± 1%        ~ (p=0.054 n=6)
Walk/Struct[fields=2]                    122.5n ± 1%   122.4n ± 2%        ~ (p=0.970 n=6)
Walk/Union[fields=2]                     121.1n ± 1%   123.0n ± 2%   +1.57% (p=0.050 n=6)
Walk/Enum[size=0_values=0]               4.327n ± 0%   4.343n ± 0%   +0.35% (p=0.006 n=6)
Walk/Fwd[struct]                         4.340n ± 0%   4.340n ± 1%        ~ (p=0.333 n=6)
Walk/Typedef[<nil>]                      111.9n ± 4%   111.3n ± 1%        ~ (p=0.699 n=6)
Walk/Volatile[<nil>]                     110.8n ± 3%   110.5n ± 2%        ~ (p=0.974 n=6)
Walk/Const[<nil>]                        108.8n ± 2%   110.4n ± 2%        ~ (p=0.121 n=6)
Walk/Restrict[<nil>]                     110.9n ± 1%   110.4n ± 2%        ~ (p=0.671 n=6)
Walk/Func[static_proto=<nil>]            112.7n ± 2%   109.7n ± 2%   -2.66% (p=0.015 n=6)
Walk/FuncProto[args=2_return=<nil>]      136.9n ± 3%   128.6n ± 1%   -6.10% (p=0.002 n=6)
Walk/Var[static]                         114.4n ± 2%   107.5n ± 3%   -5.99% (p=0.002 n=6)
Walk/Datasec                             127.7n ± 2%   118.6n ± 1%   -7.09% (p=0.002 n=6)
UnderlyingType/no_unwrapping             3.384n ± 1%   3.380n ± 0%        ~ (p=0.221 n=6)
UnderlyingType/single_unwrapping         4.817n ± 1%   4.812n ± 0%   -0.08% (p=0.006 n=6)
geomean                                  538.4n        442.0n        -6.67%               ¹
¹ benchmark set differs from baseline; geomeans may not be comparable
```

Commit messages below.

---

btf: export As and remove Transformer

    BTF encodes C language constructs such as const, restrict, etc. as separate
    types. In most cases we don't care about these when traversing a type graph.
    Users can strip these qualifiers by invoking Copy() with a
    "transformer" aka UnderlyingType(). Unfortunately, copying is quite 
    expensive.

    For this reason we changed CO-RE relocations to use an unexported as() 
    function that can be used to "unwrap" a type instead of copying it. Think of
    it as a generic version of UnderlyingType. This is much faster than copying
    and turns out to be quite ergonomic as well, since we often have to assert a
    type anyways.

    Export As() so that external users can benefit from it. Usage is like so:

        foo, ok := As[*Int](typ)
       if !ok {
           panic("not an Int")
       }

    Remove the Transformer type and the extra argument from Copy, since As is
    the better way to deal with qualifiers and typedefs. This allows further
    simplifying Copy in a follow up commit.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: rename walkType to children

    children is a much better name than walkType, so let's use that.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: replace modifyGraphPreorder with recursion

    Rewrite Copy() to use recursion instead of a manually mantained stack.

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>

btf: replace postorderIterator with recursion

    Throw out the painstakingly hand optimized postorder iterator in favour of a
    simple recursive function. Turns out this can be a lot faster for larger 
    types, probably because the visited map can be allocated on the stack.

    There is a small hit to very simple types, but it doesn't seem to affect 
    overall benchmarks too much.

        core: 1
       goos: linux
       goarch: amd64
       pkg: github.com/cilium/ebpf/btf
       cpu: 12th Gen Intel(R) Core(TM) i7-1260P
                                           │ base-po.txt │          
    recursive.txt            │
                                           │   sec/op    │   sec/op     vs base 
                 │
       PostorderTraversal/single_type           32.22n ± 0%   39.69n ± 0% 
    +23.18% (p=0.002 n=6)
       PostorderTraversal/cycle(1)              576.3n ± 3%   112.3n ± 0% 
    -80.51% (p=0.002 n=6)
       PostorderTraversal/cycle(10)             2.807µ ± 1%   1.578µ ± 1% 
    -43.80% (p=0.002 n=6)
       PostorderTraversal/gov_update_cpu_data   2.039m ± 1%   1.795m ± 1% 
    -11.97% (p=0.002 n=6)
       geomean                                  3.211µ        1.885µ      
    -41.30%

                                            │  base-po.txt   │             
    recursive.txt               │
                                           │      B/op      │     B/op      vs
    base                    │
       PostorderTraversal/single_type             0.000 ± 0%       0.000 ± 0%   
         ~ (p=1.000 n=6) ¹
       PostorderTraversal/cycle(1)                264.0 ± 0%         0.0 ± 0% 
    -100.00% (p=0.002 n=6)
       PostorderTraversal/cycle(10)               716.5 ± 0%       326.0 ± 0%  
    -54.50% (p=0.002 n=6)
       PostorderTraversal/gov_update_cpu_data   345.1Ki ± 0%     334.8Ki ± 0%   
    -2.98% (p=0.002 n=6)
       geomean                                               ²                 ?
                         ² ³
       ¹ all samples are equal
       ² summaries must be >0 to compute geomean
       ³ ratios must be >0 to compute geomean

                                            │ base-po.txt  │            
    recursive.txt              │
                                           │  allocs/op   │ allocs/op   vs base 
                      │
       PostorderTraversal/single_type           0.000 ± 0%     0.000 ± 0%       
     ~ (p=1.000 n=6) ¹
       PostorderTraversal/cycle(1)              4.000 ± 0%     0.000 ± 0% 
    -100.00% (p=0.002 n=6)
       PostorderTraversal/cycle(10)             7.000 ± 0%     1.000 ± 0%  
    -85.71% (p=0.002 n=6)
       PostorderTraversal/gov_update_cpu_data   132.0 ± 1%     115.0 ± 1%  
    -12.88% (p=0.002 n=6)
       geomean                                             ²               ?    
                     ² ³
       ¹ all samples are equal
       ² summaries must be >0 to compute geomean
       ³ ratios must be >0 to compute geomean

    Signed-off-by: Lorenz Bauer <lmb@isovalent.com>


lmb: TODO:

- Check for users of btf.Copy with non-`nil` parameter.
- Cut release before merging?
lmb: Looking at Sourcegraph makes me think nobody calls `Copy`: https://sourcegraph.com/search?q=context:global+lang:go+%22btf.Copy%22&patternType=keyword&sm=0

Patch:
diff --git a/btf/btf.go b/btf/btf.go
index 3d92467..c5def03 100644
--- a/btf/btf.go
+++ b/btf/btf.go
@@ -90,23 +90,7 @@ func (mt *mutableTypes) add(typ Type, typeIDs map[Type]TypeID) Type {
 	mt.mu.Lock()
 	defer mt.mu.Unlock()
 
-	return modifyGraphPreorder(typ, func(t Type) (Type, bool) {
-		cpy, ok := mt.copies[t]
-		if ok {
-			// This has been copied previously, no need to continue.
-			return cpy, false
-		}
-
-		cpy = t.copy()
-		mt.copies[t] = cpy
-
-		if id, ok := typeIDs[t]; ok {
-			mt.copiedTypeIDs[cpy] = id
-		}
-
-		// This is a new copy, keep copying children.
-		return cpy, true
-	})
+	return copyType(typ, typeIDs, mt.copies, mt.copiedTypeIDs)
 }
 
 // copy a set of mutable types.
@@ -122,17 +106,14 @@ func (mt *mutableTypes) copy() mutableTypes {
 	mt.mu.RLock()
 	defer mt.mu.RUnlock()
 
-	copies := make(map[Type]Type, len(mt.copies))
+	copiesOfCopies := make(map[Type]Type, len(mt.copies))
 	for orig, copy := range mt.copies {
 		// NB: We make a copy of copy, not orig, so that changes to mutable types
 		// are preserved.
-		copyOfCopy := mtCopy.add(copy, mt.copiedTypeIDs)
-		copies[orig] = copyOfCopy
+		copyOfCopy := copyType(copy, mt.copiedTypeIDs, copiesOfCopies, mtCopy.copiedTypeIDs)
+		mtCopy.copies[orig] = copyOfCopy
 	}
 
-	// mtCopy.copies is currently map[copy]copyOfCopy, replace it with
-	// map[orig]copyOfCopy.
-	mtCopy.copies = copies
 	return mtCopy
 }
 
diff --git a/btf/core.go b/btf/core.go
index a3d311a..dc5baf5 100644
--- a/btf/core.go
+++ b/btf/core.go
@@ -421,7 +421,7 @@ func coreCalculateFixup(relo *CORERelocation, target Type, targetID TypeID, bo b
 		}
 
 	case reloFieldByteOffset, reloFieldByteSize, reloFieldExists, reloFieldLShiftU64, reloFieldRShiftU64, reloFieldSigned:
-		if _, ok := as[*Fwd](target); ok {
+		if _, ok := As[*Fwd](target); ok {
 			// We can't relocate fields using a forward declaration, so
 			// skip it. If a non-forward declaration is present in the BTF
 			// we'll find it in one of the other iterations.
@@ -490,14 +490,14 @@ func coreCalculateFixup(relo *CORERelocation, target Type, targetID TypeID, bo b
 		case reloFieldSigned:
 			switch local := UnderlyingType(localField.Type).(type) {
 			case *Enum:
-				target, ok := as[*Enum](targetField.Type)
+				target, ok := As[*Enum](targetField.Type)
 				if !ok {
 					return zero, fmt.Errorf("target isn't *Enum but %T", targetField.Type)
 				}
 
 				return fixup(boolToUint64(local.Signed), boolToUint64(target.Signed))
 			case *Int:
-				target, ok := as[*Int](targetField.Type)
+				target, ok := As[*Int](targetField.Type)
 				if !ok {
 					return zero, fmt.Errorf("target isn't *Int but %T", targetField.Type)
 				}
@@ -581,7 +581,7 @@ func (ca coreAccessor) String() string {
 }
 
 func (ca coreAccessor) enumValue(t Type) (*EnumValue, error) {
-	e, ok := as[*Enum](t)
+	e, ok := As[*Enum](t)
 	if !ok {
 		return nil, fmt.Errorf("not an enum: %s", t)
 	}
@@ -707,7 +707,7 @@ func coreFindField(localT Type, localAcc coreAccessor, targetT Type) (coreField,
 
 			localMember := localMembers[acc]
 			if localMember.Name == "" {
-				localMemberType, ok := as[composite](localMember.Type)
+				localMemberType, ok := As[composite](localMember.Type)
 				if !ok {
 					return coreField{}, coreField{}, fmt.Errorf("unnamed field with type %s: %s", localMember.Type, ErrNotSupported)
 				}
@@ -721,7 +721,7 @@ func coreFindField(localT Type, localAcc coreAccessor, targetT Type) (coreField,
 				continue
 			}
 
-			targetType, ok := as[composite](target.Type)
+			targetType, ok := As[composite](target.Type)
 			if !ok {
 				return coreField{}, coreField{}, fmt.Errorf("target not composite: %w", errImpossibleRelocation)
 			}
@@ -767,7 +767,7 @@ func coreFindField(localT Type, localAcc coreAccessor, targetT Type) (coreField,
 
 		case *Array:
 			// For arrays, acc is the index in the target.
-			targetType, ok := as[*Array](target.Type)
+			targetType, ok := As[*Array](target.Type)
 			if !ok {
 				return coreField{}, coreField{}, fmt.Errorf("target not array: %w", errImpossibleRelocation)
 			}
@@ -861,7 +861,7 @@ func coreFindMember(typ composite, name string) (Member, bool, error) {
 				continue
 			}
 
-			comp, ok := as[composite](member.Type)
+			comp, ok := As[composite](member.Type)
 			if !ok {
 				return Member{}, false, fmt.Errorf("anonymous non-composite type %T not allowed", member.Type)
 			}
@@ -880,7 +880,7 @@ func coreFindEnumValue(local Type, localAcc coreAccessor, target Type) (localVal
 		return nil, nil, err
 	}
 
-	targetEnum, ok := as[*Enum](target)
+	targetEnum, ok := As[*Enum](target)
 	if !ok {
 		return nil, nil, errImpossibleRelocation
 	}
diff --git a/btf/marshal.go b/btf/marshal.go
index 744c84c..bbf2cfd 100644
--- a/btf/marshal.go
+++ b/btf/marshal.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
+	"maps"
 	"math"
 	"slices"
 	"sync"
@@ -38,6 +39,7 @@ type encoder struct {
 	buf     *bytes.Buffer
 	strings *stringTableBuilder
 	ids     map[Type]TypeID
+	visited map[Type]struct{}
 	lastID  TypeID
 }
 
@@ -163,7 +165,8 @@ func (b *Builder) Marshal(buf []byte, opts *MarshalOptions) ([]byte, error) {
 		buf:            w,
 		strings:        stb,
 		lastID:         TypeID(len(b.types)),
-		ids:            make(map[Type]TypeID, len(b.types)),
+		visited:        make(map[Type]struct{}, len(b.types)),
+		ids:            maps.Clone(b.stableIDs),
 	}
 
 	// Ensure that types are marshaled in the exact order they were Add()ed.
@@ -171,7 +174,6 @@ func (b *Builder) Marshal(buf []byte, opts *MarshalOptions) ([]byte, error) {
 	e.pending.Grow(len(b.types))
 	for _, typ := range b.types {
 		e.pending.Push(typ)
-		e.ids[typ] = b.stableIDs[typ]
 	}
 
 	if err := e.deflatePending(); err != nil {
@@ -218,16 +220,28 @@ func (b *Builder) addString(str string) (uint32, error) {
 	return b.strings.Add(str)
 }
 
-func (e *encoder) allocateID(typ Type) error {
-	id := e.lastID + 1
-	if id < e.lastID {
-		return errors.New("type ID overflow")
-	}
+func (e *encoder) allocateIDs(root Type) (err error) {
+	visitInPostorder(root, e.visited, func(typ Type) bool {
+		if _, ok := typ.(*Void); ok {
+			return true
+		}
 
-	e.pending.Push(typ)
-	e.ids[typ] = id
-	e.lastID = id
-	return nil
+		if _, ok := e.ids[typ]; ok {
+			return true
+		}
+
+		id := e.lastID + 1
+		if id < e.lastID {
+			err = errors.New("type ID overflow")
+			return false
+		}
+
+		e.pending.Push(typ)
+		e.ids[typ] = id
+		e.lastID = id
+		return true
+	})
+	return
 }
 
 // id returns the ID for the given type or panics with an error.
@@ -247,33 +261,13 @@ func (e *encoder) id(typ Type) TypeID {
 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 via Add.
-			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.allocateIDs(root); err != nil {
+			return err
 		}
 
 		if err := e.deflateType(root); err != nil {
@@ -498,7 +492,7 @@ func (e *encoder) deflateEnum64(raw *rawType, enum *Enum) (err error) {
 		if enum.Signed {
 			placeholder.Encoding = Signed
 		}
-		if err := e.allocateID(placeholder); err != nil {
+		if err := e.allocateIDs(placeholder); err != nil {
 			return fmt.Errorf("add enum64 placeholder: %w", err)
 		}
 
diff --git a/btf/traversal.go b/btf/traversal.go
index 5a7387b..c39dc66 100644
--- a/btf/traversal.go
+++ b/btf/traversal.go
@@ -2,130 +2,41 @@ package btf
 
 import (
 	"fmt"
-
-	"github.com/cilium/ebpf/internal"
 )
 
 // Functions to traverse a cyclic graph of types. The below was very useful:
 // https://eli.thegreenplace.net/2015/directed-graph-traversal-orderings-and-applications-to-data-flow-analysis/#post-order-and-reverse-post-order
 
-type postorderIterator struct {
-	// Iteration skips types for which this function returns true.
-	skip func(Type) bool
-	// The root type. May be nil if skip(root) is true.
-	root Type
-
-	// Contains types which need to be either walked or yielded.
-	types typeDeque
-	// Contains a boolean whether the type has been walked or not.
-	walked internal.Deque[bool]
-	// The set of types which has been pushed onto types.
-	pushed map[Type]struct{}
-
-	// The current type. Only valid after a call to Next().
-	Type Type
-}
-
-// postorderTraversal iterates all types reachable from root by visiting the
-// leaves of the graph first.
+// Visit all types reachable from root in postorder.
 //
-// Types for which skip returns true are ignored. skip may be nil.
-func postorderTraversal(root Type, skip func(Type) (skip bool)) postorderIterator {
-	// Avoid allocations for the common case of a skipped root.
-	if skip != nil && skip(root) {
-		return postorderIterator{}
-	}
-
-	po := postorderIterator{root: root, skip: skip}
-	walkType(root, po.push)
-
-	return po
-}
-
-func (po *postorderIterator) push(t *Type) {
-	if _, ok := po.pushed[*t]; ok || *t == po.root {
-		return
-	}
-
-	if po.skip != nil && po.skip(*t) {
-		return
+// Traversal stops if yield returns false.
+//
+// Returns false if traversal was aborted.
+func visitInPostorder(root Type, visited map[Type]struct{}, yield func(typ Type) bool) bool {
+	if _, ok := visited[root]; ok {
+		return true
 	}
-
-	if po.pushed == nil {
-		// Lazily allocate pushed to avoid an allocation for Types without children.
-		po.pushed = make(map[Type]struct{})
+	if visited == nil {
+		visited = make(map[Type]struct{})
 	}
+	visited[root] = struct{}{}
 
-	po.pushed[*t] = struct{}{}
-	po.types.Push(t)
-	po.walked.Push(false)
-}
-
-// Next returns true if there is another Type to traverse.
-func (po *postorderIterator) Next() bool {
-	for !po.types.Empty() {
-		t := po.types.Pop()
-
-		if !po.walked.Pop() {
-			// Push the type again, so that we re-evaluate it in done state
-			// after all children have been handled.
-			po.types.Push(t)
-			po.walked.Push(true)
-
-			// Add all direct children to todo.
-			walkType(*t, po.push)
-		} else {
-			// We've walked this type previously, so we now know that all
-			// children have been handled.
-			po.Type = *t
-			return true
-		}
+	cont := children(root, func(child *Type) bool {
+		return visitInPostorder(*child, visited, yield)
+	})
+	if !cont {
+		return false
 	}
 
-	// Only return root once.
-	po.Type, po.root = po.root, nil
-	return po.Type != nil
+	return yield(root)
 }
 
-// modifyGraphPreorder allows modifying every Type in a graph.
-//
-// fn is invoked in preorder for every unique Type in a graph. See [Type] for the definition
-// of equality. Every occurrence of node is substituted with its replacement.
+// children calls yield on each child of typ.
 //
-// If cont is true, fn is invoked for every child of replacement. Otherwise
-// traversal stops.
+// Traversal stops if yield returns false.
 //
-// Returns the substitution of the root node.
-func modifyGraphPreorder(root Type, fn func(node Type) (replacement Type, cont bool)) Type {
-	sub, cont := fn(root)
-	replacements := map[Type]Type{root: sub}
-
-	// This is a preorder traversal.
-	var walk func(*Type)
-	walk = func(node *Type) {
-		sub, visited := replacements[*node]
-		if visited {
-			*node = sub
-			return
-		}
-
-		sub, cont := fn(*node)
-		replacements[*node] = sub
-		*node = sub
-
-		if cont {
-			walkType(*node, walk)
-		}
-	}
-
-	if cont {
-		walkType(sub, walk)
-	}
-	return sub
-}
-
-// walkType calls fn on each child of typ.
-func walkType(typ Type, fn func(*Type)) {
+// Returns false if traversal was aborted.
+func children(typ Type, yield func(child *Type) bool) bool {
 	// Explicitly type switch on the most common types to allow the inliner to
 	// do its work. This avoids allocating intermediate slices from walk() on
 	// the heap.
@@ -133,46 +44,80 @@ func walkType(typ Type, fn func(*Type)) {
 	case *Void, *Int, *Enum, *Fwd, *Float:
 		// No children to traverse.
 	case *Pointer:
-		fn(&v.Target)
+		if !yield(&v.Target) {
+			return false
+		}
 	case *Array:
-		fn(&v.Index)
-		fn(&v.Type)
+		if !yield(&v.Index) {
+			return false
+		}
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Struct:
 		for i := range v.Members {
-			fn(&v.Members[i].Type)
+			if !yield(&v.Members[i].Type) {
+				return false
+			}
 		}
 	case *Union:
 		for i := range v.Members {
-			fn(&v.Members[i].Type)
+			if !yield(&v.Members[i].Type) {
+				return false
+			}
 		}
 	case *Typedef:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Volatile:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Const:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Restrict:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Func:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *FuncProto:
-		fn(&v.Return)
+		if !yield(&v.Return) {
+			return false
+		}
 		for i := range v.Params {
-			fn(&v.Params[i].Type)
+			if !yield(&v.Params[i].Type) {
+				return false
+			}
 		}
 	case *Var:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *Datasec:
 		for i := range v.Vars {
-			fn(&v.Vars[i].Type)
+			if !yield(&v.Vars[i].Type) {
+				return false
+			}
 		}
 	case *declTag:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *typeTag:
-		fn(&v.Type)
+		if !yield(&v.Type) {
+			return false
+		}
 	case *cycle:
 		// cycle has children, but we ignore them deliberately.
 	default:
 		panic(fmt.Sprintf("don't know how to walk Type %T", v))
 	}
+
+	return true
 }
diff --git a/btf/types.go b/btf/types.go
index 1e57a97..31acbf9 100644
--- a/btf/types.go
+++ b/btf/types.go
@@ -667,42 +667,33 @@ func pow(n int) bool {
 	return n != 0 && (n&(n-1)) == 0
 }
 
-// Transformer modifies a given Type and returns the result.
-//
-// For example, UnderlyingType removes any qualifiers or typedefs from a type.
-// See the example on Copy for how to use a transform.
-type Transformer func(Type) Type
-
 // Copy a Type recursively.
 //
-// typ may form a cycle. If transform is not nil, it is called with the
-// to be copied type, and the returned value is copied instead.
-func Copy(typ Type, transform Transformer) Type {
-	copies := make(copier)
-	return copies.copy(typ, transform)
+// typ may form a cycle.
+func Copy(typ Type) Type {
+	return copyType(typ, nil, make(map[Type]Type), nil)
 }
 
-// A map of a type to its copy.
-type copier map[Type]Type
+func copyType(typ Type, ids map[Type]TypeID, copies map[Type]Type, copiedIDs map[Type]TypeID) Type {
+	cpy, ok := copies[typ]
+	if ok {
+		// This has been copied previously, no need to continue.
+		return cpy
+	}
 
-func (c copier) copy(typ Type, transform Transformer) Type {
-	return modifyGraphPreorder(typ, func(t Type) (Type, bool) {
-		cpy, ok := c[t]
-		if ok {
-			// This has been copied previously, no need to continue.
-			return cpy, false
-		}
+	cpy = typ.copy()
+	copies[typ] = cpy
 
-		if transform != nil {
-			cpy = transform(t).copy()
-		} else {
-			cpy = t.copy()
-		}
-		c[t] = cpy
+	if id, ok := ids[typ]; ok {
+		copiedIDs[cpy] = id
+	}
 
-		// This is a new copy, keep copying children.
-		return cpy, true
+	children(cpy, func(child *Type) bool {
+		*child = copyType(*child, ids, copies, copiedIDs)
+		return true
 	})
+
+	return cpy
 }
 
 type typeDeque = internal.Deque[*Type]
@@ -1205,12 +1196,15 @@ func UnderlyingType(typ Type) Type {
 	return &cycle{typ}
 }
 
-// as returns typ if is of type T. Otherwise it peels qualifiers and Typedefs
+// As returns typ if is of type T. Otherwise it peels qualifiers and Typedefs
 // until it finds a T.
 //
 // Returns the zero value and false if there is no T or if the type is nested
 // too deeply.
-func as[T Type](typ Type) (T, bool) {
+func As[T Type](typ Type) (T, bool) {
+	// NB: We can't make this function return (*T) since then
+	// we can't assert that a type matches an interface which
+	// embeds Type: as[composite](T).
 	for depth := 0; depth <= maxResolveDepth; depth++ {
 		switch v := (typ).(type) {
 		case T:
diff --git a/internal/cmd/gentypes/main.go b/internal/cmd/gentypes/main.go
index 60b0dbc..5aebdf9 100644
--- a/internal/cmd/gentypes/main.go
+++ b/internal/cmd/gentypes/main.go
@@ -683,7 +683,7 @@ func (fpe *failedPatchError) Error() string {
 }
 
 func outputPatchedStruct(gf *btf.GoFormatter, w *bytes.Buffer, id string, s *btf.Struct, patches []patch) error {
-	s = btf.Copy(s, nil).(*btf.Struct)
+	s = btf.Copy(s).(*btf.Struct)
 
 	for i, p := range patches {
 		if err := p(s); err != nil {


Test patch:
diff --git a/btf/core_test.go b/btf/core_test.go
index f56d3d4..b5ce7ff 100644
--- a/btf/core_test.go
+++ b/btf/core_test.go
@@ -582,55 +582,6 @@ func TestCORERelocation(t *testing.T) {
 	})
 }
 
-func TestCORECopyWithoutQualifiers(t *testing.T) {
-	qualifiers := []struct {
-		name string
-		fn   func(Type) Type
-	}{
-		{"const", func(t Type) Type { return &Const{Type: t} }},
-		{"volatile", func(t Type) Type { return &Volatile{Type: t} }},
-		{"restrict", func(t Type) Type { return &Restrict{Type: t} }},
-		{"typedef", func(t Type) Type { return &Typedef{Type: t} }},
-	}
-
-	for _, test := range qualifiers {
-		t.Run(test.name+" cycle", func(t *testing.T) {
-			root := &Volatile{}
-			root.Type = test.fn(root)
-
-			cycle, ok := Copy(root, UnderlyingType).(*cycle)
-			qt.Assert(t, qt.IsTrue(ok))
-			qt.Assert(t, qt.Equals[Type](cycle.root, root))
-		})
-	}
-
-	for _, a := range qualifiers {
-		for _, b := range qualifiers {
-			t.Run(a.name+" "+b.name, func(t *testing.T) {
-				v := a.fn(&Pointer{Target: b.fn(&Int{Name: "z"})})
-				want := &Pointer{Target: &Int{Name: "z"}}
-
-				got := Copy(v, UnderlyingType)
-				qt.Assert(t, qt.DeepEquals[Type](got, want))
-			})
-		}
-	}
-
-	t.Run("long chain", func(t *testing.T) {
-		rng := testutils.Rand(t)
-		root := &Int{Name: "abc"}
-		v := Type(root)
-		for i := 0; i < maxResolveDepth; i++ {
-			q := qualifiers[rng.Intn(len(qualifiers))]
-			v = q.fn(v)
-			t.Log(q.name)
-		}
-
-		got := Copy(v, UnderlyingType)
-		qt.Assert(t, qt.DeepEquals[Type](got, root))
-	})
-}
-
 func TestCOREReloFieldSigned(t *testing.T) {
 	for _, typ := range []Type{&Int{}, &Enum{}} {
 		t.Run(fmt.Sprintf("%T with invalid target", typ), func(t *testing.T) {
diff --git a/btf/marshal_test.go b/btf/marshal_test.go
index 9cbfc0b..0d3dbe6 100644
--- a/btf/marshal_test.go
+++ b/btf/marshal_test.go
@@ -79,16 +79,11 @@ func TestRoundtripVMlinux(t *testing.T) {
 		types[i+1], types[j+1] = types[j+1], types[i+1]
 	})
 
-	seen := make(map[Type]bool)
+	visited := make(map[Type]struct{})
 limitTypes:
 	for i, typ := range types {
-		iter := postorderTraversal(typ, func(t Type) (skip bool) {
-			return seen[t]
-		})
-		for iter.Next() {
-			seen[iter.Type] = true
-		}
-		if len(seen) >= math.MaxInt16 {
+		visitInPostorder(typ, visited, func(t Type) bool { return true })
+		if len(visited) >= math.MaxInt16 {
 			// IDs exceeding math.MaxUint16 can trigger a bug when loading BTF.
 			// This can be removed once the patch lands.
 			// See https://lore.kernel.org/bpf/20220909092107.3035-1-oss@lmb.io/
diff --git a/btf/traversal_test.go b/btf/traversal_test.go
index 2b87cca..3e7bfca 100644
--- a/btf/traversal_test.go
+++ b/btf/traversal_test.go
@@ -14,22 +14,22 @@ func TestPostorderTraversal(t *testing.T) {
 
 	t.Logf("%3v", ptr)
 	pending := []Type{str, cst, ptr}
-	iter := postorderTraversal(ptr, nil)
-	for iter.Next() {
-		qt.Assert(t, qt.Equals(iter.Type, pending[0]))
+	visitInPostorder(ptr, nil, func(typ Type) bool {
+		qt.Assert(t, qt.Equals(typ, pending[0]))
 		pending = pending[1:]
-	}
+		return true
+	})
 	qt.Assert(t, qt.HasLen(pending, 0))
 
 	i := &Int{Name: "foo"}
 	// i appears twice at the same nesting depth.
 	arr := &Array{Index: i, Type: i}
 	seen := make(map[Type]bool)
-	iter = postorderTraversal(arr, nil)
-	for iter.Next() {
-		qt.Assert(t, qt.IsFalse(seen[iter.Type]))
-		seen[iter.Type] = true
-	}
+	visitInPostorder(arr, nil, func(typ Type) bool {
+		qt.Assert(t, qt.IsFalse(seen[typ]))
+		seen[typ] = true
+		return true
+	})
 	qt.Assert(t, qt.IsTrue(seen[arr]))
 	qt.Assert(t, qt.IsTrue(seen[i]))
 }
@@ -46,60 +46,26 @@ func TestPostorderTraversalVmlinux(t *testing.T) {
 		t.Run(fmt.Sprintf("%s", typ), func(t *testing.T) {
 			seen := make(map[Type]bool)
 			var last Type
-			iter := postorderTraversal(typ, nil)
-			for iter.Next() {
-				if seen[iter.Type] {
-					t.Fatalf("%s visited twice", iter.Type)
+			visitInPostorder(typ, nil, func(typ Type) bool {
+				if seen[typ] {
+					t.Fatalf("%s visited twice", typ)
 				}
-				seen[iter.Type] = true
-				last = iter.Type
-			}
+				seen[typ] = true
+				last = typ
+				return true
+			})
 			if last != typ {
 				t.Fatalf("Expected %s got %s as last type", typ, last)
 			}
 
-			walkType(typ, func(child *Type) {
+			children(typ, func(child *Type) bool {
 				qt.Check(t, qt.IsTrue(seen[*child]), qt.Commentf("missing child %s", *child))
+				return true
 			})
 		})
 	}
 }
 
-func TestModifyGraph(t *testing.T) {
-	a := &Int{}
-	b := &Int{}
-	skipped := &Int{}
-	c := &Pointer{skipped}
-	root := &Struct{
-		Members: []Member{
-			{Type: a},
-			{Type: a},
-			{Type: b},
-			{Type: c},
-		},
-	}
-
-	counts := make(map[Type]int)
-	modifyGraphPreorder(root, func(node Type) (Type, bool) {
-		counts[node]++
-		if node == c {
-			return nil, false
-		}
-		return node, true
-	})
-
-	qt.Assert(t, qt.Equals(counts[root], 1))
-	qt.Assert(t, qt.Equals(counts[a], 1))
-	qt.Assert(t, qt.Equals(counts[b], 1))
-	qt.Assert(t, qt.Equals(counts[c], 1))
-	qt.Assert(t, qt.Equals(counts[skipped], 0))
-
-	qt.Assert(t, qt.Equals[Type](root.Members[0].Type, a))
-	qt.Assert(t, qt.Equals[Type](root.Members[1].Type, a))
-	qt.Assert(t, qt.Equals[Type](root.Members[2].Type, b))
-	qt.Assert(t, qt.IsNil(root.Members[3].Type))
-}
-
 func BenchmarkPostorderTraversal(b *testing.B) {
 	spec := vmlinuxTestdataSpec(b)
 
@@ -123,38 +89,7 @@ func BenchmarkPostorderTraversal(b *testing.B) {
 		b.Run(test.name, func(b *testing.B) {
 			b.ReportAllocs()
 			for i := 0; i < b.N; i++ {
-				iter := postorderTraversal(test.typ, nil)
-				for iter.Next() {
-				}
-			}
-		})
-	}
-}
-
-func BenchmarkPreorderTraversal(b *testing.B) {
-	spec := vmlinuxTestdataSpec(b)
-
-	var fn *Func
-	err := spec.TypeByName("gov_update_cpu_data", &fn)
-	if err != nil {
-		b.Fatal(err)
-	}
-
-	for _, test := range []struct {
-		name string
-		typ  Type
-	}{
-		{"single type", &Int{}},
-		{"cycle(1)", newCyclicalType(1)},
-		{"cycle(10)", newCyclicalType(10)},
-		{"gov_update_cpu_data", fn},
-	} {
-		b.Logf("%10v", test.typ)
-
-		b.Run(test.name, func(b *testing.B) {
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				modifyGraphPreorder(test.typ, func(t Type) (Type, bool) { return t, true })
+				visitInPostorder(test.typ, nil, func(t Type) bool { return true })
 			}
 		})
 	}
diff --git a/btf/types_test.go b/btf/types_test.go
index 21dab90..69063d9 100644
--- a/btf/types_test.go
+++ b/btf/types_test.go
@@ -62,10 +62,10 @@ func TestPow(t *testing.T) {
 }
 
 func TestCopy(t *testing.T) {
-	_ = Copy((*Void)(nil), nil)
+	_ = Copy((*Void)(nil))
 
 	in := &Int{Size: 4}
-	out := Copy(in, nil)
+	out := Copy(in)
 
 	in.Size = 8
 	if size := out.(*Int).Size; size != 4 {
@@ -73,7 +73,7 @@ func TestCopy(t *testing.T) {
 	}
 
 	t.Run("cyclical", func(t *testing.T) {
-		_ = Copy(newCyclicalType(2), nil)
+		_ = Copy(newCyclicalType(2))
 	})
 
 	t.Run("identity", func(t *testing.T) {
@@ -84,7 +84,7 @@ func TestCopy(t *testing.T) {
 				{Name: "a", Type: u16},
 				{Name: "b", Type: u16},
 			},
-		}, nil)
+		})
 
 		outStruct := out.(*Struct)
 		qt.Assert(t, qt.Equals(outStruct.Members[0].Type, outStruct.Members[1].Type))
@@ -99,20 +99,20 @@ func TestAs(t *testing.T) {
 	vol := &Volatile{cst}
 
 	// It's possible to retrieve qualifiers and Typedefs.
-	haveVol, ok := as[*Volatile](vol)
+	haveVol, ok := As[*Volatile](vol)
 	qt.Assert(t, qt.IsTrue(ok))
 	qt.Assert(t, qt.Equals(haveVol, vol))
 
-	haveTd, ok := as[*Typedef](vol)
+	haveTd, ok := As[*Typedef](vol)
 	qt.Assert(t, qt.IsTrue(ok))
 	qt.Assert(t, qt.Equals(haveTd, td))
 
-	haveCst, ok := as[*Const](vol)
+	haveCst, ok := As[*Const](vol)
 	qt.Assert(t, qt.IsTrue(ok))
 	qt.Assert(t, qt.Equals(haveCst, cst))
 
 	// Make sure we don't skip Pointer.
-	haveI, ok := as[*Int](vol)
+	haveI, ok := As[*Int](vol)
 	qt.Assert(t, qt.IsFalse(ok))
 	qt.Assert(t, qt.IsNil(haveI))
 
@@ -120,7 +120,7 @@ func TestAs(t *testing.T) {
 	for _, typ := range []Type{
 		td, cst, vol, ptr,
 	} {
-		have, ok := as[*Pointer](typ)
+		have, ok := As[*Pointer](typ)
 		qt.Assert(t, qt.IsTrue(ok))
 		qt.Assert(t, qt.Equals(have, ptr))
 	}
@@ -133,7 +133,7 @@ func BenchmarkCopy(b *testing.B) {
 	b.ResetTimer()
 
 	for i := 0; i < b.N; i++ {
-		Copy(typ, nil)
+		Copy(typ)
 	}
 }
 
@@ -216,7 +216,7 @@ func TestType(t *testing.T) {
 			}
 
 			var a []*Type
-			walkType(typ, func(t *Type) { a = append(a, t) })
+			children(typ, func(t *Type) bool { a = append(a, t); return true })
 
 			if _, ok := typ.(*cycle); !ok {
 				if n := countChildren(t, reflect.TypeOf(typ)); len(a) < n {
@@ -225,7 +225,7 @@ func TestType(t *testing.T) {
 			}
 
 			var b []*Type
-			walkType(typ, func(t *Type) { b = append(b, t) })
+			children(typ, func(t *Type) bool { b = append(b, t); return true })
 
 			if diff := cmp.Diff(a, b, compareTypes); diff != "" {
 				t.Errorf("Walk mismatch (-want +got):\n%s", diff)
@@ -491,7 +491,10 @@ func BenchmarkWalk(b *testing.B) {
 
 			for i := 0; i < b.N; i++ {
 				var dq typeDeque
-				walkType(typ, dq.Push)
+				children(typ, func(child *Type) bool {
+					dq.Push(child)
+					return true
+				})
 			}
 		})
 	}
@@ -519,11 +522,9 @@ func BenchmarkUnderlyingType(b *testing.B) {
 	})
 }
 
-// Copy can be used with UnderlyingType to strip qualifiers from a type graph.
-func ExampleCopy_stripQualifiers() {
+// As can be used to strip qualifiers from a Type.
+func ExampleAs() {
 	a := &Volatile{Type: &Pointer{Target: &Typedef{Name: "foo", Type: &Int{Size: 2}}}}
-	b := Copy(a, UnderlyingType)
-	// b has Volatile and Typedef removed.
-	fmt.Printf("%3v\n", b)
-	// Output: Pointer[target=Int[unsigned size=2]]
+	fmt.Println(As[*Pointer](a))
+	// Output: Pointer[target=Typedef:"foo"] true
 }