bpf2go: generate short names for enum elements

Constant names emitted for enum elements end up quite unwieldy. They are
prefixed with enum name because in C struct, union and enum names are in
a separate namespace, allowing for overlaps with identifiers, e.g:

enum E { E };

While such overlaps are possible in *theory*, people usually don't do
it. If a typicall C naming convention is followed, we get

enum something {
    SOMETHING_FOO, SOMETHING_BAR
};

generating <STEM>SomethingSOMETHING_FOO and <STEM>SomethingSOMETHING_BAR.

In addition to "safe" long names, generate shorter ones if the
respective name is not taken. <STEM>SOMETHING_FOO and
<STEM>SOMETHING_BAR are much nicer to work with.

Signed-off-by: Nick Zavaritsky <[email protected]>

Author: mejedi

btf/format.go

@@ -25,6 +25,10 @@ type GoFormatter struct {
 	// EnumIdentifier is called for each element of an enum. By default the
 	// name of the enum type is concatenated with Identifier(element).
 	EnumIdentifier func(name, element string) string
+
+	// ShortEnumIdentifier is called for each element of an enum. An
+	// empty result causes short name to be omitted (default).
+	ShortEnumIdentifier func(name, element string) string
 }
 
 // TypeDeclaration generates a Go type declaration for a BTF type.
@@ -52,6 +56,14 @@ func (gf *GoFormatter) enumIdentifier(name, element string) string {
 	return name + gf.identifier(element)
 }
 
+func (gf *GoFormatter) shortEnumIdentifier(name, element string) string {
+	if gf.ShortEnumIdentifier != nil {
+		return gf.ShortEnumIdentifier(name, element)
+	}
+
+	return ""
+}
+
 // writeTypeDecl outputs a declaration of the given type.
 //
 // It encodes https://golang.org/ref/spec#Type_declarations:
@@ -76,13 +88,17 @@ func (gf *GoFormatter) writeTypeDecl(name string, typ Type) error {
 
 	gf.w.WriteString("; const ( ")
 	for _, ev := range e.Values {
-		id := gf.enumIdentifier(name, ev.Name)
 		var value any
 		if e.Signed {
 			value = int64(ev.Value)
 		} else {
 			value = ev.Value
 		}
+		if shortID := gf.shortEnumIdentifier(name, ev.Name); shortID != "" {
+			// output short identifier first (stringer prefers earlier decalarations)
+			fmt.Fprintf(&gf.w, "%s %s = %d; ", shortID, name, value)
+		}
+		id := gf.enumIdentifier(name, ev.Name)
 		fmt.Fprintf(&gf.w, "%s %s = %d; ", id, name, value)
 	}
 	gf.w.WriteString(")")

cmd/bpf2go/gen/output.go

@@ -173,6 +173,17 @@ func Generate(args GenerateArgs) error {
 	gf := &btf.GoFormatter{
 		Names:      nameByType,
 		Identifier: args.Identifier,
+		ShortEnumIdentifier: func(_, element string) string {
+			elementName := args.Stem + args.Identifier(element)
+			if _, nameTaken := typeByName[elementName]; nameTaken {
+				return ""
+			}
+			if _, nameReserved := reservedNames[elementName]; nameReserved {
+				return ""
+			}
+			reservedNames[elementName] = struct{}{}
+			return elementName
+		},
 	}
 
 	ctx := struct {
@@ -201,7 +212,7 @@ func Generate(args GenerateArgs) error {
 
 	var buf bytes.Buffer
 	if err := commonTemplate.Execute(&buf, &ctx); err != nil {
-		return fmt.Errorf("can't generate types: %s", err)
+		return fmt.Errorf("can't generate types: %v", err)
 	}
 
 	return internal.WriteFormatted(buf.Bytes(), args.Output)

cmd/bpf2go/gen/output_test.go

@@ -3,11 +3,13 @@ package gen
 import (
 	"bytes"
 	"fmt"
+	"regexp"
 	"strings"
 	"testing"
 
 	"github.com/go-quicktest/qt"
 
+	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/cmd/bpf2go/internal"
 )
 
@@ -63,3 +65,46 @@ func TestObjects(t *testing.T) {
 	qt.Assert(t, qt.StringContains(str, "Var1 *ebpf.Variable `ebpf:\"var_1\"`"))
 	qt.Assert(t, qt.StringContains(str, "ProgFoo1 *ebpf.Program `ebpf:\"prog_foo_1\"`"))
 }
+
+func TestEnums(t *testing.T) {
+	for _, conflict := range []bool{false, true} {
+		t.Run(fmt.Sprintf("conflict=%v", conflict), func(t *testing.T) {
+			var buf bytes.Buffer
+			args := GenerateArgs{
+				Package: "foo",
+				Stem:    "bar",
+				Types: []btf.Type{
+					&btf.Enum{Name: "EnumName", Size: 4, Values: []btf.EnumValue{
+						{Name: "V1", Value: 1}, {Name: "V2", Value: 2}, {Name: "conflict", Value: 0}}},
+				},
+				Output: &buf,
+			}
+			if conflict {
+				args.Types = append(args.Types, &btf.Struct{Name: "conflict", Size: 4})
+			}
+			err := Generate(args)
+			qt.Assert(t, qt.IsNil(err))
+
+			str := buf.String()
+
+			qt.Assert(t, qt.Matches(str, wsSeparated("barEnumNameV1", "barEnumName", "=", "1")))
+			qt.Assert(t, qt.Matches(str, wsSeparated("barEnumNameV2", "barEnumName", "=", "2")))
+			qt.Assert(t, qt.Matches(str, wsSeparated("barEnumNameConflict", "barEnumName", "=", "0")))
+
+			// short enum element names, only generated if they don't conflict with other decls
+			qt.Assert(t, qt.Matches(str, wsSeparated("barV1", "barEnumName", "=", "1")))
+			qt.Assert(t, qt.Matches(str, wsSeparated("barV2", "barEnumName", "=", "2")))
+
+			pred := qt.Matches(str, wsSeparated("barConflict", "barEnumName", "=", "0"))
+			if conflict {
+				qt.Assert(t, qt.Not(pred))
+			} else {
+				qt.Assert(t, pred)
+			}
+		})
+	}
+}
+
+func wsSeparated(terms ...string) *regexp.Regexp {
+	return regexp.MustCompile(strings.Join(terms, `\s+`))
+}