Fix auto-numbering of values on builder.EnumBuilder (#569)

jhump · Sep 6, 2023 · 63b3dc8 · 63b3dc8
1 parent 8baa453
commit 63b3dc8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/desc/builder/builder_test.go b/desc/builder/builder_test.go
@@ -2,7 +2,7 @@ package builder
 
 import (
 	"fmt"
-	"io/ioutil"
+	"io"
 	"os"
 	"sort"
 	"strings"
@@ -549,7 +549,7 @@ func loadProtoset(path string) (*desc.FileDescriptor, error) {
 		return nil, err
 	}
 	defer f.Close()
-	bb, err := ioutil.ReadAll(f)
+	bb, err := io.ReadAll(f)
 	if err != nil {
 		return nil, err
 	}
@@ -1577,6 +1577,37 @@ func TestPruneDependencies(t *testing.T) {
 	testutil.Eq(t, extDesc.GetFile().GetName(), newFileDesc.GetDependencies()[0].GetName())
 }
 
+func TestInterleavedEnumNumbers(t *testing.T) {
+	en := NewEnum("Options").
+		AddValue(NewEnumValue("OPTION_1").SetNumber(-1)).
+		AddValue(NewEnumValue("OPTION_2")).
+		AddValue(NewEnumValue("OPTION_3").SetNumber(2)).
+		AddValue(NewEnumValue("OPTION_4").SetNumber(1)).
+		AddValue(NewEnumValue("OPTION_5")).
+		AddValue(NewEnumValue("OPTION_6").SetNumber(100))
+
+	ed, err := en.Build()
+	testutil.Ok(t, err)
+
+	testutil.Require(t, ed.FindValueByName("OPTION_1") != nil)
+	testutil.Eq(t, int32(-1), ed.FindValueByName("OPTION_1").GetNumber())
+
+	testutil.Require(t, ed.FindValueByName("OPTION_2") != nil)
+	testutil.Eq(t, int32(0), ed.FindValueByName("OPTION_2").GetNumber())
+
+	testutil.Require(t, ed.FindValueByName("OPTION_3") != nil)
+	testutil.Eq(t, int32(2), ed.FindValueByName("OPTION_3").GetNumber())
+
+	testutil.Require(t, ed.FindValueByName("OPTION_4") != nil)
+	testutil.Eq(t, int32(1), ed.FindValueByName("OPTION_4").GetNumber())
+
+	testutil.Require(t, ed.FindValueByName("OPTION_5") != nil)
+	testutil.Eq(t, int32(3), ed.FindValueByName("OPTION_5").GetNumber())
+
+	testutil.Require(t, ed.FindValueByName("OPTION_6") != nil)
+	testutil.Eq(t, int32(100), ed.FindValueByName("OPTION_6").GetNumber())
+}
+
 func TestInvalid(t *testing.T) {
 	testCases := []struct {
 		name          string

diff --git a/desc/builder/enum.go b/desc/builder/enum.go
@@ -250,10 +250,11 @@ func (eb *EnumBuilder) buildProto(path []int32, sourceInfo *descriptorpb.SourceC
 
 	if len(needNumbersAssigned) > 0 {
 		tags := make([]int, len(values)-len(needNumbersAssigned))
-		for i, ev := range values {
-			tag := ev.GetNumber()
-			if tag != 0 {
-				tags[i] = int(tag)
+		tagsIndex := 0
+		for _, evb := range eb.values {
+			if evb.numberSet {
+				tags[tagsIndex] = int(evb.GetNumber())
+				tagsIndex++
 			}
 		}
 		sort.Ints(tags)