Ensure deterministic order for filtered public imports (#3734)

Author: emcfarlane

private/bufpkg/bufimage/bufimageutil/bufimageutil.go

@@ -266,9 +266,9 @@ func StripSourceRetentionOptions(image bufimage.Image) (bufimage.Image, error) {
 type transitiveClosure struct {
 	// The elements included in the transitive closure.
 	elements map[namedDescriptor]closureInclusionMode
-	// The ordered set of imports for each file. This allows for re-writing imports
+	// The set of imports for each file. This allows for re-writing imports
 	// for files whose contents have been pruned.
-	imports map[string]*orderedImports
+	imports map[string]map[string]struct{}
 }
 
 type closureInclusionMode int
@@ -295,7 +295,7 @@ const (
 func newTransitiveClosure() *transitiveClosure {
 	return &transitiveClosure{
 		elements: map[namedDescriptor]closureInclusionMode{},
-		imports:  map[string]*orderedImports{},
+		imports:  map[string]map[string]struct{}{},
 	}
 }
 
@@ -414,10 +414,10 @@ func (t *transitiveClosure) addImport(fromPath, toPath string) {
 	}
 	imps := t.imports[fromPath]
 	if imps == nil {
-		imps = newOrderedImports()
+		imps = make(map[string]struct{}, 2)
 		t.imports[fromPath] = imps
 	}
-	imps.add(toPath)
+	imps[toPath] = struct{}{}
 }
 
 func (t *transitiveClosure) addElement(
@@ -1066,60 +1066,3 @@ func stripSourceRetentionOptionsFromFile(imageFile bufimage.ImageFile) (bufimage
 		imageFile.UnusedDependencyIndexes(),
 	)
 }
-
-// orderedImports is a structure to maintain an ordered set of imports. This is needed
-// because we want to be able to iterate through imports in a deterministic way when filtering
-// the image.
-type orderedImports struct {
-	pathToIndex map[string]int
-	paths       []string
-}
-
-// newOrderedImports creates a new orderedImports structure.
-func newOrderedImports() *orderedImports {
-	return &orderedImports{
-		pathToIndex: map[string]int{},
-	}
-}
-
-// index returns the index for a given path. If the path does not exist in index map, -1
-// is returned and should be considered deleted.
-func (o *orderedImports) index(path string) int {
-	if index, ok := o.pathToIndex[path]; ok {
-		return index
-	}
-	return -1
-}
-
-// add appends a path to the paths list and the index in the map. If a key already exists,
-// then this is a no-op.
-func (o *orderedImports) add(path string) {
-	if _, ok := o.pathToIndex[path]; !ok {
-		o.pathToIndex[path] = len(o.paths)
-		o.paths = append(o.paths, path)
-	}
-}
-
-// delete removes a key from the index map of ordered imports. If a non-existent path is
-// set for deletion, then this is a no-op.
-// Note that the path is not removed from the paths list. If you want to iterate through
-// the paths, use keys() to get all non-deleted keys.
-func (o *orderedImports) delete(path string) {
-	delete(o.pathToIndex, path)
-}
-
-// keys provides all non-deleted keys from the ordered imports.
-func (o *orderedImports) keys() []string {
-	keys := make([]string, 0, len(o.pathToIndex))
-	for _, path := range o.paths {
-		if _, ok := o.pathToIndex[path]; ok {
-			keys = append(keys, path)
-		}
-	}
-	return keys
-}
-
-// len returns the number of paths in the ordered imports.
-func (o *orderedImports) len() int {
-	return len(o.pathToIndex)
-}

private/bufpkg/bufimage/bufimageutil/bufimageutil_test.go

@@ -494,6 +494,10 @@ func TestDependencies(t *testing.T) {
 		t.Parallel()
 		runDiffTest(t, "testdata/deps", "test.EnumA.txtar", WithIncludeTypes("test.EnumA"))
 	})
+	t.Run("PublicOrder", func(t *testing.T) {
+		t.Parallel()
+		runDiffTest(t, "testdata/deps", "test.PublicOrder.txtar", WithIncludeTypes("test.PublicOrder"))
+	})
 }
 
 func getImage(ctx context.Context, logger *slog.Logger, testdataDir string, options ...bufimage.BuildImageOption) (storage.ReadWriteBucket, bufimage.Image, error) {

private/bufpkg/bufimage/bufimageutil/image_filter.go

@@ -17,6 +17,7 @@ package bufimageutil
 import (
 	"fmt"
 	"slices"
+	"sort"
 	"strings"
 
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
@@ -240,12 +241,9 @@ func (b *sourcePathsBuilder) remapDependencies(
 
 	// Check if the imports need to be remapped.
 	importsRequired := b.closure.imports[fileDescriptor.GetName()]
-	importsCount := 0
-	if importsRequired != nil {
-		importsCount = importsRequired.len()
-	}
+	importsCount := len(importsRequired)
 	for _, importPath := range dependencies {
-		if importsRequired != nil && importsRequired.index(importPath) != -1 {
+		if _, ok := importsRequired[importPath]; ok {
 			importsCount--
 		} else {
 			importsCount = -1
@@ -266,26 +264,31 @@ func (b *sourcePathsBuilder) remapDependencies(
 	dependencyChanges := make([]int32, len(dependencies))
 	for _, importPath := range dependencies {
 		path := append(dependencyPath, indexFrom)
-		if importsRequired != nil && importsRequired.index(importPath) != -1 {
+		if _, ok := importsRequired[importPath]; ok {
 			dependencyChanges[indexFrom] = indexTo
 			if indexTo != indexFrom {
 				sourcePathsRemap.markMoved(path, indexTo)
 			}
 			newDependencies = append(newDependencies, importPath)
 			indexTo++
 			// delete them as we go, so we know which ones weren't in the list
-			importsRequired.delete(importPath)
+			delete(importsRequired, importPath)
 		} else {
 			sourcePathsRemap.markDeleted(path)
 			dependencyChanges[indexFrom] = -1
 		}
 		indexFrom++
 	}
-	if importsRequired != nil {
-		newDependencies = append(newDependencies, importsRequired.keys()...)
+	// Add imports picked up via a public import. The filtered files do not use public imports.
+	if publicImportCount := len(importsRequired); publicImportCount > 0 {
+		for importPath := range importsRequired {
+			newDependencies = append(newDependencies, importPath)
+		}
+		// Sort the public imports to ensure the output is deterministic.
+		sort.Strings(newDependencies[len(newDependencies)-publicImportCount:])
 	}
 
-	// Pulbic dependencies are always removed on remapping.
+	// Public dependencies are always removed on remapping.
 	publicDependencyPath := append(sourcePath, filePublicDependencyTag)
 	sourcePathsRemap.markDeleted(publicDependencyPath)
 

private/bufpkg/bufimage/bufimageutil/testdata/deps/old.proto

@@ -0,0 +1,8 @@
+syntax = "proto3";
+package old;
+
+import public "a.proto";
+import public "b.proto";
+import public "c.proto";
+
+message Old {}