Skip to content

Commit

Permalink
btf: replace modifyGraphPreorder with recursion
Browse files Browse the repository at this point in the history
Rewrite Copy() to use recursion instead of a manually mantained
stack.

Signed-off-by: Lorenz Bauer <lmb@isovalent.com>
  • Loading branch information
lmb committed Mar 27, 2024
1 parent 879831d commit a99904d
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 139 deletions.
27 changes: 4 additions & 23 deletions btf/btf.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
}

Expand Down
37 changes: 0 additions & 37 deletions btf/traversal.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,43 +87,6 @@ func (po *postorderIterator) Next() bool {
return po.Type != nil
}

// 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.
//
// If cont is true, fn is invoked for every child of replacement. Otherwise
// traversal stops.
//
// 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 {
children(*node, walk)
}
}

if cont {
children(sub, walk)
}
return sub
}

// children calls fn on each child of typ.
//
// Iteration stops early if fn returns false.
Expand Down
64 changes: 0 additions & 64 deletions btf/traversal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,41 +65,6 @@ func TestPostorderTraversalVmlinux(t *testing.T) {
}
}

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)

Expand Down Expand Up @@ -130,32 +95,3 @@ func BenchmarkPostorderTraversal(b *testing.B) {
})
}
}

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 })
}
})
}
}
31 changes: 16 additions & 15 deletions btf/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -671,27 +671,28 @@ func pow(n int) bool {
//
// typ may form a cycle.
func Copy(typ Type) Type {
copies := make(copier)
return copies.copy(typ)
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) 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

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) {
*child = copyType(*child, ids, copies, copiedIDs)
})

return cpy
}

type typeDeque = internal.Deque[*Type]
Expand Down

0 comments on commit a99904d

Please sign in to comment.