Paralellize and cache the creation of FileDescriptorProto

experimental/fdp/descriptor.go

@@ -22,18 +22,29 @@ import (
 	"github.com/bufbuild/protocompile/experimental/ir"
 )
 
+// DescriptorProto generates a single [*descriptorpb.FileDescriptorProto] for the given
+// [*ir.File].
+func DescriptorProto(file *ir.File, options ...DescriptorOption) (*descriptorpb.FileDescriptorProto, error) {
+	var g generator
+	g.init(options...)
+
+	if g.exclude != nil && g.exclude.Exclude(file) {
+		return nil, nil
+	}
+
+	fdp := new(descriptorpb.FileDescriptorProto)
+	g.file(file, fdp)
+	return fdp, nil
+}
+
 // DescriptorSetBytes generates a FileDescriptorSet for the given files, and returns the
 // result as an encoded byte slice.
 //
 // The resulting FileDescriptorSet is always fully linked: it contains all dependencies except
 // the WKTs, and all names are fully-qualified.
 func DescriptorSetBytes(files []*ir.File, options ...DescriptorOption) ([]byte, error) {
 	var g generator
-	for _, opt := range options {
-		if opt != nil {
-			opt(&g)
-		}
-	}
+	g.init(options...)
 
 	fds := new(descriptorpb.FileDescriptorSet)
 	g.files(files, fds)
@@ -46,44 +57,88 @@ func DescriptorSetBytes(files []*ir.File, options ...DescriptorOption) ([]byte,
 // The resulting FileDescriptorProto is fully linked: all names are fully-qualified.
 func DescriptorProtoBytes(file *ir.File, options ...DescriptorOption) ([]byte, error) {
 	var g generator
-	for _, opt := range options {
-		if opt != nil {
-			opt(&g)
-		}
-	}
+	g.init(options...)
 
 	fdp := new(descriptorpb.FileDescriptorProto)
 	g.file(file, fdp)
 	return proto.Marshal(fdp)
 }
 
-// DescriptorOption is an option to pass to [DescriptorSetBytes] or [DescriptorProtoBytes].
-type DescriptorOption func(*generator)
+// DescriptorOption is an option to pass to [DescriptorSetBytes], [DescriptorProtoBytes],
+// or DescriptorProto.
+type DescriptorOption interface {
+	apply(*Options)
+}
+
+// This lets us use arbitrary closures as a DescriptorOption. We use this in
+// [IncludeSourceCodeInfo], [GenerateExtraOptionLocations], and [ExcludeFiles].
+type descriptorOption func(*Options)
 
-// IncludeSourceCodeInfo sets whether or not to include google.protobuf.SourceCodeInfo in
-// the output.
-func IncludeSourceCodeInfo(flag bool) DescriptorOption {
-	return func(g *generator) {
-		if flag {
-			g.debug = new(debug)
-		} else {
-			g.debug = nil
+// [DescriptorOption] instance for [descriptorOption].
+func (dopt descriptorOption) apply(o *Options) {
+	dopt(o)
+}
+
+// [DescriptorOption] instance for [Options].
+//
+// This allows us for example to pass a value of [Options] to [DescriptorProto].
+func (o *Options) apply(that *Options) {
+	*that = *o
+}
+
+// Apply applies the given [DescriptorOption] to this [*Options].
+//
+// Nil values are ignored; does nothing if opt is nil.
+func (o *Options) Apply(options ...DescriptorOption) *Options {
+	if o != nil {
+		for _, option := range options {
+			if option != nil {
+				option.apply(o)
+			}
 		}
 	}
+	return o
 }
 
-// ExcludeFiles excludes the given files from the output of [DescriptorSetBytes].
-func ExcludeFiles(exclude func(*ir.File) bool) DescriptorOption {
-	return func(g *generator) {
-		g.exclude = exclude
+// apply applies the given [DescriptorOption] to this [*generator] and initializes some of its fields.
+func (g *generator) init(options ...DescriptorOption) {
+	g.Apply(options...)
+	if g.includeSourceCodeInfo {
+		g.debug = new(debug)
+	} else {
+		g.debug = nil
 	}
 }
 
+// IncludeSourceCodeInfo sets whether or not to include google.protobuf.SourceCodeInfo in
+// the output.
+func IncludeSourceCodeInfo(flag bool) DescriptorOption {
+	return descriptorOption(func(o *Options) {
+		o.includeSourceCodeInfo = flag
+	})
+}
+
 // GenerateExtraOptionLocations set whether or not to generate additional locations for
 // elements inside of message literals in option values. This option is a no-op if
 // [IncludeSourceCodeInfo] is not set.
 func GenerateExtraOptionLocations(flag bool) DescriptorOption {
-	return func(g *generator) {
-		g.generateExtraOptionLocations = flag
-	}
+	return descriptorOption(func(o *Options) {
+		o.generateExtraOptionLocations = flag
+	})
+}
+
+// Excluder is used with [ExcludeFiles].
+//
+// This is an interface, rather than a function, so that implementations can be comparable for
+// use in queries.
+type Excluder interface {
+	Exclude(*ir.File) bool
+}
+
+// ExcludeFiles excludes the given files from the output of [DescriptorSetBytes] and
+// [DescriptorProto].
+func ExcludeFiles(exclude Excluder) DescriptorOption {
+	return descriptorOption(func(o *Options) {
+		o.exclude = exclude
+	})
 }

experimental/fdp/generator.go

@@ -39,11 +39,18 @@ import (
 	"github.com/bufbuild/protocompile/internal/tags"
 )
 
-type generator struct {
-	currentFile                  *ir.File
+// Options records a set of [DescriptorOptions].
+//
+// This type is intended for making options comparable, such as for use in queries.
+type Options struct {
+	includeSourceCodeInfo        bool
 	generateExtraOptionLocations bool
-	exclude                      func(*ir.File) bool
+	exclude                      Excluder
+}
 
+type generator struct {
+	currentFile *ir.File
+	Options
 	debug *debug
 }
 
@@ -52,7 +59,7 @@ func (g *generator) files(files []*ir.File, fds *descriptorpb.FileDescriptorSet)
 	// imports for each file because we want the result to be sorted
 	// topologically.
 	for file := range ir.TopoSort(files) {
-		if g.exclude != nil && g.exclude(file) {
+		if g.exclude != nil && g.exclude.Exclude(file) {
 			continue
 		}
 

experimental/incremental/queries/fdp.go

@@ -0,0 +1,47 @@
+// Copyright 2020-2025 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 queries
+
+import (
+	"google.golang.org/protobuf/types/descriptorpb"
+
+	"github.com/bufbuild/protocompile/experimental/fdp"
+	"github.com/bufbuild/protocompile/experimental/incremental"
+	"github.com/bufbuild/protocompile/experimental/ir"
+)
+
+// FDP is an [incremental.Query] that converts the lowered IR files [*ir.File]
+// to [*descriptorpb.FileDescriptorProto]
+//
+// FDP queries with different File and Options are considered distinct.
+// The File field must not be edited between different FDP queries!
+type FDP struct {
+	*ir.File // Must be comparable.
+	fdp.Options
+}
+
+var _ incremental.Query[*descriptorpb.FileDescriptorProto] = FDP{}
+
+// Key implements [incremental.Query].
+func (l FDP) Key() any {
+	return l
+}
+
+// Execute implements [incremental.Query].
+func (l FDP) Execute(t *incremental.Task) (*descriptorpb.FileDescriptorProto, error) {
+	t.Report().Options.Stage += stageFDP
+
+	return fdp.DescriptorProto(l.File, &l.Options)
+}

experimental/incremental/queries/fds.go

@@ -0,0 +1,91 @@
+// Copyright 2020-2025 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 queries
+
+import (
+	"slices"
+
+	"google.golang.org/protobuf/types/descriptorpb"
+
+	"github.com/bufbuild/protocompile/experimental/fdp"
+	"github.com/bufbuild/protocompile/experimental/incremental"
+	"github.com/bufbuild/protocompile/experimental/ir"
+	"github.com/bufbuild/protocompile/experimental/source"
+	"github.com/bufbuild/protocompile/internal/ext/slicesx"
+)
+
+// FDS is an [incremental.Query] that produces a [*descriptorpb.FileDescriptorSet].
+// It takes the IR files [ir.File] produced via [Link], topologically
+// sorts the resulting IR files, and converts each to a [*descriptorpb.FileDescriptorProto]
+// via [FDP]. It then bundles them all together into a [*descriptorpb.FileDescriptorSet].
+//
+// FDS queries with different Openers, options, and workspaces are considered distinct.
+type FDS struct {
+	source.Opener // Must be comparable.
+	*ir.Session
+	source.Workspace // Must be comparable
+	fdp.Options
+}
+
+var _ incremental.Query[*descriptorpb.FileDescriptorSet] = FDS{}
+
+// Key implements [incremental.Query].
+func (l FDS) Key() any {
+	return l
+}
+
+// Execute implements [incremental.Query].
+func (l FDS) Execute(t *incremental.Task) (*descriptorpb.FileDescriptorSet, error) {
+	t.Report().Options.Stage += stageFDS
+
+	linkQuery := Link{
+		Opener:    l.Opener,
+		Session:   l.Session,
+		Workspace: l.Workspace,
+	}
+
+	linkResult, err := incremental.Resolve(t, linkQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	irs := linkResult[0].Value
+	irs = slices.DeleteFunc(irs, func(f *ir.File) bool { return f == nil })
+
+	fdpQueries := slicesx.Transform(
+		slices.Collect(ir.TopoSort(irs)),
+		func(f *ir.File) incremental.Query[*descriptorpb.FileDescriptorProto] {
+			return FDP{
+				File:    f,
+				Options: l.Options,
+			}
+		},
+	)
+
+	fdpResults, err := incremental.Resolve(t, fdpQueries...)
+	if err != nil {
+		return nil, err
+	}
+	fdps, err := fdpResults.Slice()
+	if err != nil {
+		return nil, err
+	}
+	fdps = slices.DeleteFunc(
+		fdps,
+		func(fdp *descriptorpb.FileDescriptorProto) bool { return fdp == nil },
+	)
+
+	return &descriptorpb.FileDescriptorSet{File: fdps}, nil
+}

experimental/incremental/queries/queries.go

@@ -24,4 +24,6 @@ const (
 	stageAST
 	stageIR
 	stageLink
+	stageFDP
+	stageFDS
 )