btf: avoid heap allocations when walking types

Type.walk() forces the typeDeque argument to always escape, since the escape analyzer has
to be conservative when dealing with interfaces. We can apply two tricks that reduce
allocations.

First, Type.walk() is replaced with

    children() []*Type

This removes the typeDeque parameter, and for types that don't have any children (like Void) we
never incur an allocation at all. For types that do have children we've unfortunately swapped
heap allocating typeDeque with allocating the []*Type slice instead.

Which brings us to the second trick: by using a type switch we allow the compiler to inline the
children() function for the most important Types, which in turn enables allocating most []*Type on
the stack instead of on the heap.

One notable exception is FuncProto, which allocates the returned slice as such:

    types := make([]*Type, len(fp.Params)+1)

It seems like this trips up escape analysis for some reason as the slice is heap allocated.

    name                                      old time/op    new time/op    delta
    Walk/Void-4                                 22.6ns ± 2%     2.8ns ± 2%   -87.40%  (p=0.029 n=4+4)
    Walk/Int[unsigned_size=0]-4                 22.7ns ± 1%     3.2ns ± 1%   -86.09%  (p=0.029 n=4+4)
    Walk/Pointer[target=<nil>]-4                61.3ns ± 1%    38.5ns ± 1%   -37.26%  (p=0.029 n=4+4)
    Walk/Array[index=<nil>_type=<nil>_n=0]-4    63.8ns ± 1%    42.2ns ± 0%   -33.91%  (p=0.029 n=4+4)
    Walk/Struct[fields=2]-4                     64.9ns ± 1%    59.5ns ± 0%    -8.29%  (p=0.029 n=4+4)
    Walk/Union[fields=2]-4                      66.6ns ± 3%    59.3ns ± 0%   -10.84%  (p=0.029 n=4+4)
    Walk/Enum[size=0_values=0]-4                22.5ns ± 4%     2.9ns ± 2%   -87.22%  (p=0.029 n=4+4)
    Walk/Fwd[struct]-4                          22.2ns ± 1%     2.8ns ± 1%   -87.23%  (p=0.029 n=4+4)
    Walk/Typedef[<nil>]-4                       60.8ns ± 1%    38.7ns ± 1%   -36.42%  (p=0.029 n=4+4)
    Walk/Volatile[<nil>]-4                      62.9ns ± 3%    38.4ns ± 2%   -38.89%  (p=0.029 n=4+4)
    Walk/Const[<nil>]-4                         61.9ns ± 1%    39.8ns ± 3%   -35.78%  (p=0.029 n=4+4)
    Walk/Restrict[<nil>]-4                      64.2ns ± 1%    39.2ns ± 1%   -38.91%  (p=0.029 n=4+4)
    Walk/Func[static_proto=<nil>]-4             63.3ns ± 4%    39.5ns ± 2%   -37.57%  (p=0.029 n=4+4)
    Walk/FuncProto[args=2_return=<nil>]-4       67.8ns ± 5%    77.3ns ± 2%   +14.03%  (p=0.029 n=4+4)
    Walk/Var[static]-4                          63.2ns ± 1%    38.5ns ± 0%   -39.03%  (p=0.029 n=4+4)
    Walk/Datasec-4                              67.5ns ± 2%    58.5ns ± 1%   -13.31%  (p=0.029 n=4+4)

    name                                      old alloc/op   new alloc/op   delta
    Walk/Void-4                                  48.0B ± 0%      0.0B       -100.00%  (p=0.029 n=4+4)
    Walk/Int[unsigned_size=0]-4                  48.0B ± 0%      0.0B       -100.00%  (p=0.029 n=4+4)
    Walk/Pointer[target=<nil>]-4                  112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Array[index=<nil>_type=<nil>_n=0]-4      112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Struct[fields=2]-4                       112B ± 0%       80B ± 0%   -28.57%  (p=0.029 n=4+4)
    Walk/Union[fields=2]-4                        112B ± 0%       80B ± 0%   -28.57%  (p=0.029 n=4+4)
    Walk/Enum[size=0_values=0]-4                 48.0B ± 0%      0.0B       -100.00%  (p=0.029 n=4+4)
    Walk/Fwd[struct]-4                           48.0B ± 0%      0.0B       -100.00%  (p=0.029 n=4+4)
    Walk/Typedef[<nil>]-4                         112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Volatile[<nil>]-4                        112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Const[<nil>]-4                           112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Restrict[<nil>]-4                        112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Func[static_proto=<nil>]-4               112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/FuncProto[args=2_return=<nil>]-4         112B ± 0%       88B ± 0%   -21.43%  (p=0.029 n=4+4)
    Walk/Var[static]-4                            112B ± 0%       64B ± 0%   -42.86%  (p=0.029 n=4+4)
    Walk/Datasec-4                                112B ± 0%       80B ± 0%   -28.57%  (p=0.029 n=4+4)

    name                                      old allocs/op  new allocs/op  delta
    Walk/Void-4                                   1.00 ± 0%      0.00       -100.00%  (p=0.029 n=4+4)
    Walk/Int[unsigned_size=0]-4                   1.00 ± 0%      0.00       -100.00%  (p=0.029 n=4+4)
    Walk/Pointer[target=<nil>]-4                  2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Array[index=<nil>_type=<nil>_n=0]-4      2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Struct[fields=2]-4                       2.00 ± 0%      2.00 ± 0%      ~     (all equal)
    Walk/Union[fields=2]-4                        2.00 ± 0%      2.00 ± 0%      ~     (all equal)
    Walk/Enum[size=0_values=0]-4                  1.00 ± 0%      0.00       -100.00%  (p=0.029 n=4+4)
    Walk/Fwd[struct]-4                            1.00 ± 0%      0.00       -100.00%  (p=0.029 n=4+4)
    Walk/Typedef[<nil>]-4                         2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Volatile[<nil>]-4                        2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Const[<nil>]-4                           2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Restrict[<nil>]-4                        2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Func[static_proto=<nil>]-4               2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/FuncProto[args=2_return=<nil>]-4         2.00 ± 0%      2.00 ± 0%      ~     (all equal)
    Walk/Var[static]-4                            2.00 ± 0%      1.00 ± 0%   -50.00%  (p=0.029 n=4+4)
    Walk/Datasec-4                                2.00 ± 0%      2.00 ± 0%      ~     (all equal)

btf/core.go

@@ -869,8 +869,8 @@ func coreAreTypesCompatible(localType Type, targetType Type) error {
 
 		case *Pointer, *Array:
 			depth++
-			localType.walk(&localTs)
-			targetType.walk(&targetTs)
+			walkType(localType, localTs.Push)
+			walkType(targetType, targetTs.Push)
 
 		case *FuncProto:
 			tv := targetType.(*FuncProto)
@@ -879,8 +879,8 @@ func coreAreTypesCompatible(localType Type, targetType Type) error {
 			}
 
 			depth++
-			localType.walk(&localTs)
-			targetType.walk(&targetTs)
+			walkType(localType, localTs.Push)
+			walkType(targetType, targetTs.Push)
 
 		default:
 			return fmt.Errorf("unsupported type %T", localType)

btf/traversal.go

@@ -0,0 +1,52 @@
+package btf
+
+// walkType calls fn on each child of typ.
+//
+// It's faster than directly invoking typ.children for some common types.
+func walkType(typ Type, fn func(*Type)) {
+	walk := func(children []*Type) {
+		for _, c := range children {
+			fn(c)
+		}
+	}
+
+	// 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.
+	switch v := typ.(type) {
+	case *Void:
+		walk(v.children())
+	case *Int:
+		walk(v.children())
+	case *Pointer:
+		walk(v.children())
+	case *Array:
+		walk(v.children())
+	case *Struct:
+		walk(v.children())
+	case *Union:
+		walk(v.children())
+	case *Enum:
+		walk(v.children())
+	case *Fwd:
+		walk(v.children())
+	case *Typedef:
+		walk(v.children())
+	case *Volatile:
+		walk(v.children())
+	case *Const:
+		walk(v.children())
+	case *Restrict:
+		walk(v.children())
+	case *Func:
+		walk(v.children())
+	case *FuncProto:
+		walk(v.children())
+	case *Var:
+		walk(v.children())
+	case *Datasec:
+		walk(v.children())
+	default:
+		walk(v.children())
+	}
+}

btf/types.go

@@ -35,9 +35,9 @@ type Type interface {
 	// Make a copy of the type, without copying Type members.
 	copy() Type
 
-	// Enumerate all nested Types. Repeated calls must visit nested
+	// Enumerate all nested Types. Repeated calls must return nested
 	// types in the same order.
-	walk(*typeDeque)
+	children() []*Type
 }
 
 var (
@@ -72,7 +72,7 @@ func (v *Void) Format(fs fmt.State, verb rune) { formatType(fs, verb, v) }
 func (v *Void) TypeName() string               { return "" }
 func (v *Void) size() uint32                   { return 0 }
 func (v *Void) copy() Type                     { return (*Void)(nil) }
-func (v *Void) walk(*typeDeque)                {}
+func (v *Void) children() []*Type              { return nil }
 
 type IntEncoding byte
 
@@ -124,9 +124,9 @@ func (i *Int) Format(fs fmt.State, verb rune) {
 	formatType(fs, verb, i, i.Encoding, "size=", i.Size*8)
 }
 
-func (i *Int) TypeName() string { return i.Name }
-func (i *Int) size() uint32     { return i.Size }
-func (i *Int) walk(*typeDeque)  {}
+func (i *Int) TypeName() string  { return i.Name }
+func (i *Int) size() uint32      { return i.Size }
+func (i *Int) children() []*Type { return nil }
 func (i *Int) copy() Type {
 	cpy := *i
 	return &cpy
@@ -141,9 +141,9 @@ func (p *Pointer) Format(fs fmt.State, verb rune) {
 	formatType(fs, verb, p, "target=", p.Target)
 }
 
-func (p *Pointer) TypeName() string    { return "" }
-func (p *Pointer) size() uint32        { return 8 }
-func (p *Pointer) walk(tdq *typeDeque) { tdq.push(&p.Target) }
+func (p *Pointer) TypeName() string  { return "" }
+func (p *Pointer) size() uint32      { return 8 }
+func (p *Pointer) children() []*Type { return []*Type{&p.Target} }
 func (p *Pointer) copy() Type {
 	cpy := *p
 	return &cpy
@@ -162,10 +162,7 @@ func (arr *Array) Format(fs fmt.State, verb rune) {
 
 func (arr *Array) TypeName() string { return "" }
 
-func (arr *Array) walk(tdq *typeDeque) {
-	tdq.push(&arr.Index)
-	tdq.push(&arr.Type)
-}
+func (arr *Array) children() []*Type { return []*Type{&arr.Index, &arr.Type} }
 
 func (arr *Array) copy() Type {
 	cpy := *arr
@@ -188,10 +185,12 @@ func (s *Struct) TypeName() string { return s.Name }
 
 func (s *Struct) size() uint32 { return s.Size }
 
-func (s *Struct) walk(tdq *typeDeque) {
+func (s *Struct) children() []*Type {
+	types := make([]*Type, len(s.Members))
 	for i := range s.Members {
-		tdq.push(&s.Members[i].Type)
+		types[i] = &s.Members[i].Type
 	}
+	return types
 }
 
 func (s *Struct) copy() Type {
@@ -220,10 +219,12 @@ func (u *Union) TypeName() string { return u.Name }
 
 func (u *Union) size() uint32 { return u.Size }
 
-func (u *Union) walk(tdq *typeDeque) {
+func (u *Union) children() []*Type {
+	types := make([]*Type, 0, len(u.Members))
 	for i := range u.Members {
-		tdq.push(&u.Members[i].Type)
+		types = append(types, &u.Members[i].Type)
 	}
+	return types
 }
 
 func (u *Union) copy() Type {
@@ -293,8 +294,8 @@ type EnumValue struct {
 	Value uint64
 }
 
-func (e *Enum) size() uint32    { return e.Size }
-func (e *Enum) walk(*typeDeque) {}
+func (e *Enum) size() uint32      { return e.Size }
+func (e *Enum) children() []*Type { return nil }
 func (e *Enum) copy() Type {
 	cpy := *e
 	cpy.Values = make([]EnumValue, len(e.Values))
@@ -334,7 +335,7 @@ func (f *Fwd) Format(fs fmt.State, verb rune) {
 
 func (f *Fwd) TypeName() string { return f.Name }
 
-func (f *Fwd) walk(*typeDeque) {}
+func (f *Fwd) children() []*Type { return nil }
 func (f *Fwd) copy() Type {
 	cpy := *f
 	return &cpy
@@ -352,7 +353,7 @@ func (td *Typedef) Format(fs fmt.State, verb rune) {
 
 func (td *Typedef) TypeName() string { return td.Name }
 
-func (td *Typedef) walk(tdq *typeDeque) { tdq.push(&td.Type) }
+func (td *Typedef) children() []*Type { return []*Type{&td.Type} }
 func (td *Typedef) copy() Type {
 	cpy := *td
 	return &cpy
@@ -369,8 +370,8 @@ func (v *Volatile) Format(fs fmt.State, verb rune) {
 
 func (v *Volatile) TypeName() string { return "" }
 
-func (v *Volatile) qualify() Type       { return v.Type }
-func (v *Volatile) walk(tdq *typeDeque) { tdq.push(&v.Type) }
+func (v *Volatile) qualify() Type     { return v.Type }
+func (v *Volatile) children() []*Type { return []*Type{&v.Type} }
 func (v *Volatile) copy() Type {
 	cpy := *v
 	return &cpy
@@ -387,8 +388,8 @@ func (c *Const) Format(fs fmt.State, verb rune) {
 
 func (c *Const) TypeName() string { return "" }
 
-func (c *Const) qualify() Type       { return c.Type }
-func (c *Const) walk(tdq *typeDeque) { tdq.push(&c.Type) }
+func (c *Const) qualify() Type     { return c.Type }
+func (c *Const) children() []*Type { return []*Type{&c.Type} }
 func (c *Const) copy() Type {
 	cpy := *c
 	return &cpy
@@ -405,8 +406,8 @@ func (r *Restrict) Format(fs fmt.State, verb rune) {
 
 func (r *Restrict) TypeName() string { return "" }
 
-func (r *Restrict) qualify() Type       { return r.Type }
-func (r *Restrict) walk(tdq *typeDeque) { tdq.push(&r.Type) }
+func (r *Restrict) qualify() Type     { return r.Type }
+func (r *Restrict) children() []*Type { return []*Type{&r.Type} }
 func (r *Restrict) copy() Type {
 	cpy := *r
 	return &cpy
@@ -430,7 +431,7 @@ func (f *Func) Format(fs fmt.State, verb rune) {
 
 func (f *Func) TypeName() string { return f.Name }
 
-func (f *Func) walk(tdq *typeDeque) { tdq.push(&f.Type) }
+func (f *Func) children() []*Type { return []*Type{&f.Type} }
 func (f *Func) copy() Type {
 	cpy := *f
 	return &cpy
@@ -448,11 +449,13 @@ func (fp *FuncProto) Format(fs fmt.State, verb rune) {
 
 func (fp *FuncProto) TypeName() string { return "" }
 
-func (fp *FuncProto) walk(tdq *typeDeque) {
-	tdq.push(&fp.Return)
+func (fp *FuncProto) children() []*Type {
+	types := make([]*Type, len(fp.Params)+1)
+	types[0] = &fp.Return
 	for i := range fp.Params {
-		tdq.push(&fp.Params[i].Type)
+		types[i+1] = &fp.Params[i].Type
 	}
+	return types
 }
 
 func (fp *FuncProto) copy() Type {
@@ -480,7 +483,7 @@ func (v *Var) Format(fs fmt.State, verb rune) {
 
 func (v *Var) TypeName() string { return v.Name }
 
-func (v *Var) walk(tdq *typeDeque) { tdq.push(&v.Type) }
+func (v *Var) children() []*Type { return []*Type{&v.Type} }
 func (v *Var) copy() Type {
 	cpy := *v
 	return &cpy
@@ -501,10 +504,12 @@ func (ds *Datasec) TypeName() string { return ds.Name }
 
 func (ds *Datasec) size() uint32 { return ds.Size }
 
-func (ds *Datasec) walk(tdq *typeDeque) {
+func (ds *Datasec) children() []*Type {
+	types := make([]*Type, len(ds.Vars))
 	for i := range ds.Vars {
-		tdq.push(&ds.Vars[i].Type)
+		types[i] = &ds.Vars[i].Type
 	}
+	return types
 }
 
 func (ds *Datasec) copy() Type {
@@ -535,9 +540,9 @@ func (f *Float) Format(fs fmt.State, verb rune) {
 	formatType(fs, verb, f, "size=", f.Size*8)
 }
 
-func (f *Float) TypeName() string { return f.Name }
-func (f *Float) size() uint32     { return f.Size }
-func (f *Float) walk(*typeDeque)  {}
+func (f *Float) TypeName() string  { return f.Name }
+func (f *Float) size() uint32      { return f.Size }
+func (f *Float) children() []*Type { return nil }
 func (f *Float) copy() Type {
 	cpy := *f
 	return &cpy
@@ -557,8 +562,8 @@ func (dt *declTag) Format(fs fmt.State, verb rune) {
 	formatType(fs, verb, dt, "type=", dt.Type, "value=", dt.Value, "index=", dt.Index)
 }
 
-func (dt *declTag) TypeName() string   { return "" }
-func (dt *declTag) walk(td *typeDeque) { td.push(&dt.Type) }
+func (dt *declTag) TypeName() string  { return "" }
+func (dt *declTag) children() []*Type { return []*Type{&dt.Type} }
 func (dt *declTag) copy() Type {
 	cpy := *dt
 	return &cpy
@@ -574,9 +579,9 @@ func (tt *typeTag) Format(fs fmt.State, verb rune) {
 	formatType(fs, verb, tt, "type=", tt.Type, "value=", tt.Value)
 }
 
-func (tt *typeTag) TypeName() string   { return "" }
-func (tt *typeTag) qualify() Type      { return tt.Type }
-func (tt *typeTag) walk(td *typeDeque) { td.push(&tt.Type) }
+func (tt *typeTag) TypeName() string  { return "" }
+func (tt *typeTag) qualify() Type     { return tt.Type }
+func (tt *typeTag) children() []*Type { return []*Type{&tt.Type} }
 func (tt *typeTag) copy() Type {
 	cpy := *tt
 	return &cpy
@@ -590,7 +595,7 @@ type cycle struct {
 func (c *cycle) ID() TypeID                     { return math.MaxUint32 }
 func (c *cycle) Format(fs fmt.State, verb rune) { formatType(fs, verb, c, "root=", c.root) }
 func (c *cycle) TypeName() string               { return "" }
-func (c *cycle) walk(*typeDeque)                {}
+func (c *cycle) children() []*Type              { return nil }
 func (c *cycle) copy() Type {
 	cpy := *c
 	return &cpy
@@ -739,7 +744,7 @@ func (c copier) copy(typ *Type, transform Transformer) {
 		*t = cpy
 
 		// Mark any nested types for copying.
-		cpy.walk(&work)
+		walkType(cpy, work.push)
 	}
 }
 

btf/types_test.go

@@ -153,11 +153,7 @@ func TestType(t *testing.T) {
 				t.Error("Copy doesn't copy")
 			}
 
-			var first, second typeDeque
-			typ.walk(&first)
-			typ.walk(&second)
-
-			if diff := cmp.Diff(first.all(), second.all(), compareTypes); diff != "" {
+			if diff := cmp.Diff(typ.children(), typ.children(), compareTypes); diff != "" {
 				t.Errorf("Walk mismatch (-want +got):\n%s", diff)
 			}
 		})
@@ -440,7 +436,8 @@ func BenchmarkWalk(b *testing.B) {
 			b.ReportAllocs()
 
 			for i := 0; i < b.N; i++ {
-				typ.walk(&typeDeque{})
+				var dq typeDeque
+				walkType(typ, dq.Push)
 			}
 		})
 	}