Add PrimitiveSet.EntriesInKeysetOrder in Go. #tinkApiChange

PiperOrigin-RevId: 533279854

go/core/primitiveset/primitiveset.go

@@ -70,16 +70,19 @@ type PrimitiveSet struct {
 	// primitives sharing the prefix). This allows quickly retrieving the
 	// primitives sharing some particular prefix.
 	Entries map[string][]*Entry
+	// Stores entries in the original keyset key order.
+	EntriesInKeysetOrder []*Entry
 
 	Annotations map[string]string
 }
 
 // New returns an empty instance of PrimitiveSet.
 func New() *PrimitiveSet {
 	return &PrimitiveSet{
-		Primary:     nil,
-		Entries:     make(map[string][]*Entry),
-		Annotations: nil,
+		Primary:              nil,
+		Entries:              make(map[string][]*Entry),
+		EntriesInKeysetOrder: make([]*Entry, 0),
+		Annotations:          nil,
 	}
 }
 
@@ -121,5 +124,6 @@ func (ps *PrimitiveSet) Add(p interface{}, key *tinkpb.Keyset_Key) (*Entry, erro
 		key.GetKeyData().GetTypeUrl(),
 	)
 	ps.Entries[prefix] = append(ps.Entries[prefix], e)
+	ps.EntriesInKeysetOrder = append(ps.EntriesInKeysetOrder, e)
 	return e, nil
 }

go/core/primitiveset/primitiveset_test.go

@@ -34,12 +34,12 @@ func makeTestKey(keyID int, status tinkpb.KeyStatusType, outputPrefixType tinkpb
 
 func TestPrimitvesetNew(t *testing.T) {
 	ps := primitiveset.New()
-	if ps.Primary != nil || ps.Entries == nil {
+	if ps.Primary != nil || ps.Entries == nil || ps.EntriesInKeysetOrder == nil {
 		t.Errorf("expect primary to be nil and primitives is initialized")
 	}
 }
 
-func TestPrimitivesetAddEntries(t *testing.T) {
+func TestPrimitivesetAddAndEntriesInKeysetOrder(t *testing.T) {
 	keys := []*tinkpb.Keyset_Key{
 		makeTestKey(1234543, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK, "type.url.1"),
 		makeTestKey(7213743, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_LEGACY, "type.url.2"),
@@ -87,6 +87,9 @@ func TestPrimitivesetAddEntries(t *testing.T) {
 	if !cmp.Equal(got, want) {
 		t.Errorf("got = %v, want = %v", got, want)
 	}
+	if !cmp.Equal(ps.EntriesInKeysetOrder, want) {
+		t.Errorf("EntriesInKeysetOrder = %v, want = %v", ps.EntriesInKeysetOrder, want)
+	}
 }
 
 func TestPrimitivesetRawEntries(t *testing.T) {