Optionally include reachable fieldpaths in prompt

bazel/BUILD.bazel

@@ -0,0 +1,6 @@
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+
+bzl_library(
+    name = "proto_source_info",
+    srcs = ["proto_source_info.bzl"],
+)

bazel/proto_source_info.bzl

@@ -0,0 +1,54 @@
+"""Build rule for preserving source information in proto descriptor sets."""
+
+load("@com_google_protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
+
+def _source_info_proto_descriptor_set(ctx):
+    """Returns a proto descriptor set with source information preserved."""
+    srcs = depset([s for dep in ctx.attr.proto_libs for s in dep[ProtoInfo].direct_sources])
+    deps = depset(transitive = [dep[ProtoInfo].transitive_descriptor_sets for dep in ctx.attr.proto_libs])
+
+    src_files = srcs.to_list()
+    dep_files = deps.to_list()
+
+    args = ctx.actions.args()
+    args.add("--descriptor_set_out=" + ctx.outputs.out.path)
+    args.add("--include_imports")
+    args.add("--include_source_info=true")
+    args.add("--proto_path=.")
+    args.add("--proto_path=" + ctx.configuration.genfiles_dir.path)
+    args.add("--descriptor_set_in=" + ":".join([d.path for d in dep_files]))
+    args.add_all(src_files)
+
+    ctx.actions.run(
+        executable = ctx.executable._protoc,
+        inputs = src_files + dep_files,
+        outputs = [ctx.outputs.out],
+        arguments = [args],
+        mnemonic = "SourceInfoProtoDescriptorSet",
+        progress_message = "Generating proto descriptor set with source information for %{label}",
+    )
+
+source_info_proto_descriptor_set = rule(
+    doc = """
+Rule for generating a proto descriptor set for the transitive dependencies of proto libraries
+with source information preserved.
+
+This can dramatically increase the size of the descriptor set, so only use it
+when necessary (e.g. for formatting documentation about a CEL environment).
+
+Source info is only preserved for input files for each proto_library label in
+protolibs. Transitive dependencies are included with source info stripped.
+""",
+    attrs = {
+        "proto_libs": attr.label_list(providers = [[ProtoInfo]]),
+        "_protoc": attr.label(
+            default = "@com_google_protobuf//:protoc",
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+    outputs = {
+        "out": "%{name}-transitive-descriptor-set-source-info.proto.bin",
+    },
+    implementation = _source_info_proto_descriptor_set,
+)

cel/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "cel.go",
         "decls.go",
         "env.go",
+        "fieldpaths.go",
         "folding.go",
         "inlining.go",
         "io.go",
@@ -43,6 +44,7 @@ go_library(
         "//interpreter:go_default_library",
         "//parser:go_default_library",
         "@dev_cel_expr//:expr",
+        "@dev_cel_expr//conformance/proto3:go_default_library",
         "@org_golang_google_genproto_googleapis_api//expr/v1alpha1:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
         "@org_golang_google_protobuf//reflect/protodesc:go_default_library",
@@ -63,6 +65,7 @@ go_test(
         "cel_test.go",
         "decls_test.go",
         "env_test.go",
+        "fieldpaths_test.go",
         "folding_test.go",
         "inlining_test.go",
         "io_test.go",
@@ -78,6 +81,7 @@ go_test(
     ],
     embedsrcs = [
         "//cel/testdata:prompts",
+        "//cel/testdata:test_fds_with_source_info",
     ],
     deps = [
         "//common/operators:go_default_library",
@@ -89,6 +93,7 @@ go_test(
         "//test:go_default_library",
         "//test/proto2pb:go_default_library",
         "//test/proto3pb:go_default_library",
+        "@com_github_google_go_cmp//cmp:go_default_library",
         "@org_golang_google_genproto_googleapis_api//expr/v1alpha1:go_default_library",
         "@org_golang_google_protobuf//encoding/prototext:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
@@ -100,4 +105,4 @@ go_test(
 exports_files(
     ["templates/authoring.tmpl"],
     visibility = ["//visibility:public"],
-)
+)

cel/fieldpaths.go

@@ -0,0 +1,163 @@
+package cel
+
+import (
+	"slices"
+	"strings"
+
+	"github.com/google/cel-go/common"
+	"github.com/google/cel-go/common/types"
+)
+
+// fieldPath represents a selection path to a field from a variable in a CEL environment.
+type fieldPath struct {
+	celType *Type
+	// path represents the selection path to the field.
+	path        string
+	description string
+	isLeaf      bool
+}
+
+// Documentation implements the Documentor interface.
+func (f *fieldPath) Documentation() *common.Doc {
+	return common.NewFieldDoc(f.path, f.celType.String(), f.description)
+}
+
+type documentationProvider interface {
+	// FindStructFieldDescription returns documentation for a field if available.
+	// Returns false if the field could not be found.
+	FindStructFieldDescription(typeName, fieldName string) (string, bool)
+}
+
+type backtrack struct {
+	// provider used to resolve types.
+	provider types.Provider
+	// paths of fields that have been visited along the path.
+	path []string
+	// types of fields that have been visited along the path. used to avoid cycles.
+	types []*Type
+}
+
+func (b *backtrack) push(pathStep string, celType *Type) {
+	b.path = append(b.path, pathStep)
+	b.types = append(b.types, celType)
+}
+
+func (b *backtrack) pop() {
+	b.path = b.path[:len(b.path)-1]
+	b.types = b.types[:len(b.types)-1]
+}
+
+func formatPath(path []string) string {
+	var buffer strings.Builder
+	for i, p := range path {
+		if i == 0 {
+			buffer.WriteString(p)
+			continue
+		}
+		if strings.HasPrefix(p, "[") {
+			buffer.WriteString(p)
+			continue
+		}
+		buffer.WriteString(".")
+		buffer.WriteString(p)
+	}
+	return buffer.String()
+}
+
+func (b *backtrack) expandFieldPaths(celType *Type, paths []*fieldPath) []*fieldPath {
+	if slices.ContainsFunc(b.types[:len(b.types)-1], func(t *Type) bool { return t.String() == celType.String() }) {
+		// Cycle detected, so stop expanding.
+		paths[len(paths)-1].isLeaf = false
+		return paths
+	}
+	switch celType.Kind() {
+	case types.StructKind:
+		fields, ok := b.provider.FindStructFieldNames(celType.String())
+		if !ok {
+			// Caller added this type to the path, so it must be a leaf.
+			paths[len(paths)-1].isLeaf = true
+			return paths
+		}
+		for _, field := range fields {
+			fieldType, ok := b.provider.FindStructFieldType(celType.String(), field)
+			if !ok {
+				// Field not found, either hidden or an error.
+				continue
+			}
+			b.push(field, celType)
+			description := ""
+			if docProvider, ok := b.provider.(documentationProvider); ok {
+				description, _ = docProvider.FindStructFieldDescription(celType.String(), field)
+			}
+			path := &fieldPath{
+				celType:     fieldType.Type,
+				path:        formatPath(b.path),
+				description: description,
+				isLeaf:      false,
+			}
+			paths = append(paths, path)
+			paths = b.expandFieldPaths(fieldType.Type, paths)
+			b.pop()
+		}
+		return paths
+	case types.MapKind:
+		if len(celType.Parameters()) != 2 {
+			// dynamic map, so treat as a leaf.
+			paths[len(paths)-1].isLeaf = true
+			return paths
+		}
+		mapKeyType := celType.Parameters()[0]
+		mapValueType := celType.Parameters()[1]
+		// Add a placeholder for the map key kind (the zero value).
+		keyIdentifier := ""
+		switch mapKeyType.Kind() {
+		case types.StringKind:
+			keyIdentifier = "[\"\"]"
+		case types.IntKind:
+			keyIdentifier = "[0]"
+		case types.UintKind:
+			keyIdentifier = "[0u]"
+		case types.BoolKind:
+			keyIdentifier = "[false]"
+		default:
+			// Caller added this type to the path, so it must be a leaf.
+			paths[len(paths)-1].isLeaf = true
+			return paths
+		}
+		b.push(keyIdentifier, mapValueType)
+		defer b.pop()
+		return b.expandFieldPaths(mapValueType, paths)
+	case types.ListKind:
+		if len(celType.Parameters()) != 1 {
+			// dynamic list, so treat as a leaf.
+			paths[len(paths)-1].isLeaf = true
+			return paths
+		}
+		listElemType := celType.Parameters()[0]
+		b.push("[0]", listElemType)
+		defer b.pop()
+		return b.expandFieldPaths(listElemType, paths)
+	default:
+		paths[len(paths)-1].isLeaf = true
+	}
+
+	return paths
+}
+
+// fieldPathsForType expands the reachable fields from the given root identifier.
+func fieldPathsForType(provider types.Provider, identifier string, celType *Type) []*fieldPath {
+	b := &backtrack{
+		provider: provider,
+		path:     []string{identifier},
+		types:    []*Type{celType},
+	}
+	paths := []*fieldPath{
+		{
+			celType: celType,
+			path:    identifier,
+			isLeaf:  false,
+		},
+	}
+
+	return b.expandFieldPaths(celType, paths)
+}