Fix import dropping for options

private/bufpkg/bufimage/bufimageutil/bufimageutil.go

@@ -380,7 +380,7 @@ func (t *transitiveClosure) includeType(
 	return nil
 }
 
-func (t *transitiveClosure) addImport(fromPath, toPath string) {
+func (t *transitiveClosure) addImport(wat namedDescriptor, fromPath, toPath string) {
 	if fromPath == toPath {
 		return // no need for a file to import itself
 	}
@@ -401,9 +401,8 @@ func (t *transitiveClosure) addElement(
 ) error {
 	descriptorInfo := imageIndex.ByDescriptor[descriptor]
 	if referrerFile != "" {
-		t.addImport(referrerFile, descriptorInfo.file.Path())
+		t.addImport(descriptor, referrerFile, descriptorInfo.file.Path())
 	}
-
 	if existingMode, ok := t.elements[descriptor]; ok && existingMode != inclusionModeEnclosing {
 		if existingMode == inclusionModeImplicit && !impliedByCustomOption {
 			// upgrade from implied to explicitly part of closure
@@ -417,15 +416,6 @@ func (t *transitiveClosure) addElement(
 		t.elements[descriptor] = inclusionModeExplicit
 	}
 
-	// if this type is enclosed inside another, add enclosing types
-	if err := t.addEnclosing(descriptorInfo.parent, descriptorInfo.file.Path(), imageIndex, opts); err != nil {
-		return err
-	}
-	// add any custom options and their dependencies
-	if err := t.exploreCustomOptions(descriptor, descriptorInfo.file.Path(), imageIndex, opts); err != nil {
-		return err
-	}
-
 	switch typedDescriptor := descriptor.(type) {
 	case *descriptorpb.FileDescriptorProto:
 		typeNames, ok := imageIndex.FileTypes[typedDescriptor.GetName()]
@@ -536,6 +526,15 @@ func (t *transitiveClosure) addElement(
 		return errorUnsupportedFilterType(descriptor, descriptorInfo.fullName)
 	}
 
+	// if this type is enclosed inside another, add enclosing types
+	if err := t.addEnclosing(descriptorInfo.parent, descriptorInfo.file.Path(), imageIndex, opts); err != nil {
+		return err
+	}
+	// add any custom options and their dependencies
+	if err := t.exploreCustomOptions(descriptor, descriptorInfo.file.Path(), imageIndex, opts); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -805,9 +804,9 @@ func (t *transitiveClosure) exploreCustomOptions(
 	optionsName := options.Descriptor().FullName()
 	var err error
 	options.Range(func(fd protoreflect.FieldDescriptor, val protoreflect.Value) bool {
-		//if !t.hasOption(fd, opts) {
-		//	return true
-		//}
+		if !t.hasOption(fd, opts) {
+			return true
+		}
 		// If the value contains an Any message, we should add the message type
 		// therein to the closure.
 		if err = t.exploreOptionValueForAny(fd, val, referrerFile, imageIndex, opts); err != nil {

private/bufpkg/bufimage/bufimageutil/bufimageutil_test.go

@@ -35,7 +35,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/tools/txtar"
-	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protodesc"
 	"google.golang.org/protobuf/reflect/protoreflect"
@@ -210,21 +209,51 @@ func TestOptions(t *testing.T) {
 func TestOptionImports(t *testing.T) {
 	t.Parallel()
 
+	// This checks that when excluding options the imports are correctly dropped.
+	// For this case when both options are removed only a.proto should be left.
+	testdataDir := "testdata/imports"
+	bucket, err := storageos.NewProvider().NewReadWriteBucket(testdataDir)
+	require.NoError(t, err)
+	testModuleData := []bufmoduletesting.ModuleData{
+		{
+			Bucket: storage.FilterReadBucket(bucket, storage.MatchPathEqual("a.proto")),
+		},
+		{
+			Bucket:      storage.FilterReadBucket(bucket, storage.MatchPathEqual("options.proto")),
+			NotTargeted: true,
+		},
+	}
+	moduleSet, err := bufmoduletesting.NewModuleSet(testModuleData...)
+	require.NoError(t, err)
+
+	// Safe to filter the image concurrently as its not being modified.
+	image, err := bufimage.BuildImage(
+		context.Background(),
+		slogtestext.NewLogger(t),
+		bufmodule.ModuleSetToModuleReadBucketWithOnlyProtoFiles(moduleSet),
+		bufimage.WithExcludeSourceCodeInfo(),
+	)
+	require.NoError(t, err)
+
 	t.Run("none_excluded", func(t *testing.T) {
 		t.Parallel()
-		runDiffTest(t, "testdata/imports", "none_excluded.txtar", WithExcludeOptions("google.protobuf.FieldOptions.jstype"))
+		generated := runFilterImage(t, image, WithExcludeOptions("google.protobuf.FieldOptions.jstype"))
+		checkExpectation(t, context.Background(), generated, bucket, "none_excluded.txtar")
 	})
 	t.Run("exclude_foo", func(t *testing.T) {
 		t.Parallel()
-		runDiffTest(t, "testdata/imports", "foo.txtar", WithExcludeOptions("message_foo"))
+		generated := runFilterImage(t, image, WithExcludeOptions("message_foo"))
+		checkExpectation(t, context.Background(), generated, bucket, "foo.txtar")
 	})
 	t.Run("exclude_foo_bar", func(t *testing.T) {
 		t.Parallel()
-		runDiffTest(t, "testdata/imports", "foo_bar.txtar", WithExcludeOptions("message_foo", "message_bar"))
+		generated := runFilterImage(t, image, WithExcludeOptions("message_foo", "message_bar"))
+		checkExpectation(t, context.Background(), generated, bucket, "foo_bar.txtar")
 	})
 	t.Run("exclude_bar", func(t *testing.T) {
 		t.Parallel()
-		runDiffTest(t, "testdata/imports", "bar.txtar", WithIncludeOptions("message_foo"), WithExcludeOptions("message_bar"))
+		generated := runFilterImage(t, image, WithIncludeOptions("message_foo"), WithExcludeOptions("message_bar"))
+		checkExpectation(t, context.Background(), generated, bucket, "bar.txtar")
 	})
 }
 
@@ -383,31 +412,26 @@ func runDiffTest(t *testing.T, testdataDir string, expectedFile string, opts ...
 	ctx := context.Background()
 	bucket, image, err := getImage(ctx, slogtestext.NewLogger(t), testdataDir, bufimage.WithExcludeSourceCodeInfo())
 	require.NoError(t, err)
+	generated := runFilterImage(t, image, opts...)
+	checkExpectation(t, ctx, generated, bucket, expectedFile)
+}
 
-	_ = protojson.MarshalOptions{}
-	b, _ := protojson.MarshalOptions{Indent: "  "}.Marshal(bufimage.ImageToFileDescriptorSet(image))
-	t.Log(string(b))
-
+func runFilterImage(t *testing.T, image bufimage.Image, opts ...ImageFilterOption) []byte {
 	filteredImage, err := FilterImage(image, opts...)
 	require.NoError(t, err)
 	assert.NotNil(t, image)
 	assert.True(t, imageIsDependencyOrdered(filteredImage), "image files not in dependency order")
 
 	// We may have filtered out custom options from the set in the step above. However, the options messages
 	// still contain extension fields that refer to the custom options, as a result of building the image.
-	// So we serialize and then de-serialize, and use only the filtered results to parse extensions. That
-	// way, the result will omit custom options that aren't present in the filtered set (as they will be
+	// So we serialize and then de-serialize, and use only the filtered results to parse extensions. That way, the result will omit custom options that aren't present in the filtered set (as they will be
 	// considered unrecognized fields).
 	data, err := protoencoding.NewWireMarshaler().Marshal(bufimage.ImageToFileDescriptorSet(filteredImage))
 	require.NoError(t, err)
 	fileDescriptorSet := &descriptorpb.FileDescriptorSet{}
 	err = protoencoding.NewWireUnmarshaler(filteredImage.Resolver()).Unmarshal(data, fileDescriptorSet)
 	require.NoError(t, err)
 
-	t.Log("--------------------------------------------")
-	//b, _ = protojson.MarshalOptions{Indent: "  "}.Marshal(bufimage.ImageToFileDescriptorSet(image))
-	//t.Log(string(b))
-
 	files, err := protodesc.NewFiles(fileDescriptorSet)
 	require.NoError(t, err)
 
@@ -432,7 +456,7 @@ func runDiffTest(t *testing.T, testdataDir string, expectedFile string, opts ...
 		return archive.Files[i].Name < archive.Files[j].Name
 	})
 	generated := txtar.Format(archive)
-	checkExpectation(t, ctx, generated, bucket, expectedFile)
+	return generated
 }
 
 func checkExpectation(t *testing.T, ctx context.Context, actual []byte, bucket storage.ReadWriteBucket, expectedFile string) {

private/bufpkg/bufimage/bufimageutil/image_filter.go

@@ -17,9 +17,11 @@ package bufimageutil
 import (
 	"fmt"
 	"slices"
+	"sort"
 	"strings"
 
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/pkg/slicesext"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/descriptorpb"
@@ -50,6 +52,9 @@ func filterImage(image bufimage.Image, options *imageFilterOptions) (bufimage.Im
 	// below, when remapping the descriptor.
 	if len(options.includeTypes) == 0 {
 		for _, file := range image.Files() {
+			if file.IsImport() {
+				continue
+			}
 			if err := closure.addElement(file.FileDescriptorProto(), "", false, imageIndex, options); err != nil {
 				return nil, err
 			}

private/bufpkg/bufimage/bufimageutil/testdata/options/options.jstype.include.txtar

@@ -1,7 +1,6 @@
 -- a.proto --
 syntax = "proto2";
 package pkg;
-import "options.proto";
 message Empty {
 }
 message Foo {