Split imports for directory strategy in buf generate

private/buf/bufgen/generator.go

@@ -24,6 +24,7 @@ import (
 	"github.com/bufbuild/buf/private/buf/bufprotopluginexec"
 	"github.com/bufbuild/buf/private/bufpkg/bufconfig"
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/bufpkg/bufimage/bufimagegenerate"
 	"github.com/bufbuild/buf/private/bufpkg/bufimage/bufimagemodify"
 	"github.com/bufbuild/buf/private/bufpkg/bufprotoplugin"
 	"github.com/bufbuild/buf/private/bufpkg/bufprotoplugin/bufprotopluginos"
@@ -324,7 +325,7 @@ func (g *generator) execLocalPlugin(
 	if err != nil {
 		return nil, err
 	}
-	requests, err := bufimage.ImagesToCodeGeneratorRequests(
+	requests, err := bufimagegenerate.ImagesToCodeGeneratorRequests(
 		pluginImages,
 		pluginConfig.Opt(),
 		nil,

private/buf/bufgen/image_provider.go

@@ -19,6 +19,7 @@ import (
 	"sync"
 
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/bufpkg/bufimage/bufimagegenerate"
 )
 
 // imageProvider is used to provide the images used
@@ -29,7 +30,7 @@ import (
 // strategy.
 type imageProvider struct {
 	image       bufimage.Image
-	imagesByDir []bufimage.Image
+	imagesByDir []bufimagegenerate.ImageForGeneration
 	lock        sync.Mutex
 }
 
@@ -39,16 +40,16 @@ func newImageProvider(image bufimage.Image) *imageProvider {
 	}
 }
 
-func (p *imageProvider) GetImages(strategy Strategy) ([]bufimage.Image, error) {
+func (p *imageProvider) GetImages(strategy Strategy) ([]bufimagegenerate.ImageForGeneration, error) {
 	switch strategy {
 	case StrategyAll:
-		return []bufimage.Image{p.image}, nil
+		return []bufimagegenerate.ImageForGeneration{bufimagegenerate.NewImageForGenerationFromImage(p.image)}, nil
 	case StrategyDirectory:
 		p.lock.Lock()
 		defer p.lock.Unlock()
 		if p.imagesByDir == nil {
 			var err error
-			p.imagesByDir, err = bufimage.ImageByDir(p.image)
+			p.imagesByDir, err = bufimagegenerate.ImageByDirSplitImports(p.image)
 			if err != nil {
 				return nil, err
 			}

private/buf/cmd/buf/command/alpha/protoc/plugin.go

@@ -21,8 +21,10 @@ import (
 
 	"github.com/bufbuild/buf/private/buf/bufprotopluginexec"
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/bufpkg/bufimage/bufimagegenerate"
 	"github.com/bufbuild/buf/private/pkg/app"
 	"github.com/bufbuild/buf/private/pkg/command"
+	"github.com/bufbuild/buf/private/pkg/slicesext"
 	"github.com/bufbuild/buf/private/pkg/storage/storageos"
 	"github.com/bufbuild/buf/private/pkg/tracing"
 	"go.uber.org/zap"
@@ -59,8 +61,8 @@ func executePlugin(
 		storageosProvider,
 		runner,
 	)
-	requests, err := bufimage.ImagesToCodeGeneratorRequests(
-		images,
+	requests, err := bufimagegenerate.ImagesToCodeGeneratorRequests(
+		slicesext.Map(images, bufimagegenerate.NewImageForGenerationFromImage),
 		strings.Join(pluginInfo.Opt, ","),
 		bufprotopluginexec.DefaultVersion,
 		false,

private/buf/cmd/protoc-gen-buf-lint/lint_test.go

@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/bufpkg/bufimage/bufimagegenerate"
 	"github.com/bufbuild/buf/private/pkg/app"
 	"github.com/bufbuild/buf/private/pkg/command"
 	"github.com/bufbuild/buf/private/pkg/normalpath"
@@ -365,7 +366,7 @@ func testBuildRequest(
 	}
 	image, err := bufimage.NewImage(imageFiles)
 	require.NoError(t, err)
-	codeGenReq, err := bufimage.ImageToCodeGeneratorRequest(image, parameter, nil, false, false)
+	codeGenReq, err := bufimagegenerate.ImageToCodeGeneratorRequest(bufimagegenerate.NewImageForGenerationFromImage(image), parameter, nil, false, false)
 	require.NoError(t, err)
 	request, err := protoplugin.NewRequest(codeGenReq)
 	require.NoError(t, err)

private/bufpkg/bufimage/bufimage.go

@@ -577,94 +577,6 @@ func ImageToFileDescriptorProtos(image Image) []*descriptorpb.FileDescriptorProt
 	return imageFilesToFileDescriptorProtos(image.Files())
 }
 
-// ImageToCodeGeneratorRequest returns a new CodeGeneratorRequest for the Image.
-//
-// All non-imports are added as files to generate.
-// If includeImports is set, all non-well-known-type imports are also added as files to generate.
-// If includeWellKnownTypes is set, well-known-type imports are also added as files to generate.
-// includeWellKnownTypes has no effect if includeImports is not set.
-func ImageToCodeGeneratorRequest(
-	image Image,
-	parameter string,
-	compilerVersion *pluginpb.Version,
-	includeImports bool,
-	includeWellKnownTypes bool,
-) (*pluginpb.CodeGeneratorRequest, error) {
-	return imageToCodeGeneratorRequest(
-		image,
-		parameter,
-		compilerVersion,
-		includeImports,
-		includeWellKnownTypes,
-		nil,
-		nil,
-	)
-}
-
-// ImagesToCodeGeneratorRequests converts the Images to CodeGeneratorRequests.
-//
-// All non-imports are added as files to generate.
-// If includeImports is set, all non-well-known-type imports are also added as files to generate.
-// If includeImports is set, only one CodeGeneratorRequest will contain any given file as a FileToGenerate.
-// If includeWellKnownTypes is set, well-known-type imports are also added as files to generate.
-// includeWellKnownTypes has no effect if includeImports is not set.
-func ImagesToCodeGeneratorRequests(
-	images []Image,
-	parameter string,
-	compilerVersion *pluginpb.Version,
-	includeImports bool,
-	includeWellKnownTypes bool,
-) ([]*pluginpb.CodeGeneratorRequest, error) {
-	requests := make([]*pluginpb.CodeGeneratorRequest, len(images))
-	// alreadyUsedPaths is a map of paths that have already been added to an image.
-	//
-	// We track this if includeImports is set, so that when we find an import, we can
-	// see if the import was already added to a CodeGeneratorRequest via another Image
-	// in the Image slice. If the import was already added, we do not add duplicates
-	// across CodeGeneratorRequests.
-	var alreadyUsedPaths map[string]struct{}
-	// nonImportPaths is a map of non-import paths.
-	//
-	// We track this if includeImports is set. If we find a non-import file in Image A
-	// and this file is an import in Image B, the file will have already been added to
-	// a CodeGeneratorRequest via Image A, so do not add the duplicate to any other
-	// CodeGeneratorRequest.
-	var nonImportPaths map[string]struct{}
-	if includeImports {
-		// We don't need to track these if includeImports is false, so we only populate
-		// the maps if includeImports is true. If includeImports is false, only non-imports
-		// will be added to each CodeGeneratorRequest, so figuring out whether or not
-		// we should add a given import to a given CodeGeneratorRequest is unnecessary.
-		//
-		// imageToCodeGeneratorRequest checks if these maps are nil before every access.
-		alreadyUsedPaths = make(map[string]struct{})
-		nonImportPaths = make(map[string]struct{})
-		for _, image := range images {
-			for _, imageFile := range image.Files() {
-				if !imageFile.IsImport() {
-					nonImportPaths[imageFile.Path()] = struct{}{}
-				}
-			}
-		}
-	}
-	for i, image := range images {
-		var err error
-		requests[i], err = imageToCodeGeneratorRequest(
-			image,
-			parameter,
-			compilerVersion,
-			includeImports,
-			includeWellKnownTypes,
-			alreadyUsedPaths,
-			nonImportPaths,
-		)
-		if err != nil {
-			return nil, err
-		}
-	}
-	return requests, nil
-}
-
 type newImageForProtoOptions struct {
 	noReparse            bool
 	computeUnusedImports bool

private/bufpkg/bufimage/bufimagegenerate/bufimagegenerate.go

@@ -0,0 +1,71 @@
+// Copyright 2020-2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package bufimagegenerate
+
+import (
+	"github.com/bufbuild/buf/private/bufpkg/bufimage"
+	"github.com/bufbuild/buf/private/pkg/slicesext"
+)
+
+// ImageForGeneration is a buf image to be generated.
+type ImageForGeneration interface {
+	// files are the files that comprise the image, but not all files should be generated.
+	files() []imageFileForGeneration
+}
+
+// NewImageForGenerationFromImage creates an ImageForGeneration that works the exact same way as before
+// this type is added, and this is the one of the two exported functions that construct an ImageForGeneration.
+// The other one is ImageByDirSplitImports.
+//
+// TODO: update this doc
+func NewImageForGenerationFromImage(image bufimage.Image) ImageForGeneration {
+	imageFilesForGeneration := slicesext.Map(image.Files(), func(imageFile bufimage.ImageFile) imageFileForGeneration {
+		return imageFileForGeneration{
+			ImageFile:  imageFile,
+			toGenerate: true,
+		}
+	})
+	return imageForGeneration(imageFilesForGeneration)
+}
+
+// *** PRIVATE ***
+
+type imageFileForGeneration struct {
+	bufimage.ImageFile
+
+	// toGenerate returns whether the file may be generated. This is not necessarily the same as IsImport(),
+	// especially when strategy is set to directory. It also does not guarantee the file's inclusion in file_to_generate.
+	//
+	//  - If it returns true, the file will be generated unless includeImports or includeWKT says otherwise.
+	//  - If it returns false, it will not be generated regardless of includeImports and includeWKT.
+	toGenerate bool
+}
+
+type imageForGeneration []imageFileForGeneration
+
+func (i imageForGeneration) files() []imageFileForGeneration {
+	return []imageFileForGeneration(i)
+}
+
+func newImageForGeneration(image bufimage.Image, filesToGenerate map[string]struct{}) ImageForGeneration {
+	imageFilesForGeneration := imageForGeneration(slicesext.Map(image.Files(), func(imageFile bufimage.ImageFile) imageFileForGeneration {
+		_, ok := filesToGenerate[imageFile.Path()]
+		return imageFileForGeneration{
+			ImageFile:  imageFile,
+			toGenerate: ok,
+		}
+	}))
+	return imageForGeneration(imageFilesForGeneration)
+}