btf: Add spec types iterator

btf/btf.go

@@ -655,6 +655,30 @@ func (s *Spec) TypeByName(name string, typ interface{}) error {
 	return nil
 }
 
+// TypesIterator iterates over types of a given spec.
+type TypesIterator struct {
+	spec  *Spec
+	index int
+	// The last visited type in the spec.
+	Type Type
+}
+
+// Iterate returns the types iterator.
+func (s *Spec) Iterate() *TypesIterator {
+	return &TypesIterator{spec: s, index: 0}
+}
+
+// Next returns true as long as there are any remaining types.
+func (iter *TypesIterator) Next() bool {
+	if len(iter.spec.types) <= iter.index {
+		return false
+	}
+
+	iter.Type = iter.spec.types[iter.index]
+	iter.index++
+	return true
+}
+
 // Handle is a reference to BTF loaded into the kernel.
 type Handle struct {
 	spec *Spec

btf/btf_test.go

@@ -342,3 +342,38 @@ func ExampleSpec_TypeByName() {
 	// We've found struct foo
 	fmt.Println(foo.Name)
 }
+
+func TestTypesIterator(t *testing.T) {
+	spec, err := LoadSpecFromReader(readVMLinux(t))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(spec.types) < 1 {
+		t.Fatal("Not enough types")
+	}
+
+	// Assertion that 'iphdr' type exists within the spec
+	_, err = spec.AnyTypeByName("iphdr")
+	if err != nil {
+		t.Fatalf("Failed to find 'iphdr' type by name: %s", err)
+	}
+
+	found := false
+	count := 0
+
+	iter := spec.Iterate()
+	for iter.Next() {
+		if !found && iter.Type.TypeName() == "iphdr" {
+			found = true
+		}
+		count += 1
+	}
+
+	if l := len(spec.types); l != count {
+		t.Fatalf("Failed to iterate over all types (%d vs %d)", l, count)
+	}
+	if !found {
+		t.Fatal("Cannot find 'iphdr' type")
+	}
+}