Repository: kythe/kythe
Pull request #5509 (2023-01-20 22:47:55)

test: add utility to resolve MarkedSources from entries


Patch:
diff --git a/kythe/go/util/markedsource/BUILD b/kythe/go/util/markedsource/BUILD
index d308b422e..2f0d6c6eb 100644
--- a/kythe/go/util/markedsource/BUILD
+++ b/kythe/go/util/markedsource/BUILD
@@ -4,20 +4,34 @@ package(default_visibility = ["//kythe:default_visibility"])
 
 go_library(
     name = "markedsource",
-    srcs = ["markedsource.go"],
-    deps = ["//kythe/proto:common_go_proto"],
+    srcs = [
+        "markedsource.go",
+        "resolve.go",
+    ],
+    deps = [
+        "//kythe/go/util/kytheuri",
+        "//kythe/go/util/schema/edges",
+        "//kythe/go/util/schema/facts",
+        "//kythe/proto:common_go_proto",
+        "//kythe/proto:storage_go_proto",
+        "@org_golang_google_protobuf//proto:go_default_library",
+    ],
 )
 
 go_test(
     name = "markedsource_test",
     size = "small",
-    srcs = ["markedsource_test.go"],
+    srcs = [
+        "markedsource_test.go",
+        "resolve_test.go",
+    ],
     data = ["//kythe/cxx/doc"],
     library = "markedsource",
     visibility = ["//visibility:private"],
     deps = [
+        "//kythe/go/test/testutil",
+        "//kythe/go/util/compare",
         "//kythe/proto:common_go_proto",
         "@org_golang_google_protobuf//encoding/prototext:go_default_library",
-        "@org_golang_google_protobuf//proto:go_default_library",
     ],
 )
diff --git a/kythe/go/util/markedsource/resolve.go b/kythe/go/util/markedsource/resolve.go
new file mode 100644
index 000000000..a16e1b3cd
--- /dev/null
+++ b/kythe/go/util/markedsource/resolve.go
@@ -0,0 +1,153 @@
+/*
+ * Copyright 2023 The Kythe Authors. All rights reserved.
+ *
+ * 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 markedsource
+
+import (
+	"fmt"
+
+	"kythe.io/kythe/go/util/kytheuri"
+	"kythe.io/kythe/go/util/schema/edges"
+	"kythe.io/kythe/go/util/schema/facts"
+
+	"google.golang.org/protobuf/proto"
+
+	cpb "kythe.io/kythe/proto/common_go_proto"
+	spb "kythe.io/kythe/proto/storage_go_proto"
+)
+
+// Resolver produces fully resolved MarkedSources from a set of entries.
+type Resolver struct {
+	nodes map[string]*cpb.MarkedSource
+
+	params  map[string][]string
+	tparams map[string][]string
+	typed   map[string]string
+}
+
+// NewResolver constructs a MarkedSource resolver from the given entries.
+func NewResolver(entries []*spb.Entry) (*Resolver, error) {
+	r := &Resolver{
+		nodes:   make(map[string]*cpb.MarkedSource),
+		params:  make(map[string][]string),
+		tparams: make(map[string][]string),
+		typed:   make(map[string]string),
+	}
+	for _, e := range entries {
+		if e.GetFactName() == facts.Code {
+			ticket := kytheuri.ToString(e.GetSource())
+			var ms cpb.MarkedSource
+			if err := proto.Unmarshal(e.GetFactValue(), &ms); err != nil {
+				return nil, fmt.Errorf("error unmarshalling code for %s: %v", ticket, err)
+			}
+			r.nodes[ticket] = &ms
+		} else if e.GetEdgeKind() != "" {
+			ticket := kytheuri.ToString(e.GetSource())
+			kind, ord, _ := edges.ParseOrdinal(e.GetEdgeKind())
+			if ord < 0 {
+				return nil, fmt.Errorf("invalid ordinal: %d", ord)
+			}
+			switch kind {
+			case edges.Typed:
+				r.typed[ticket] = kytheuri.ToString(e.GetTarget())
+			case edges.Param:
+				params := r.params[ticket]
+				if len(params)-1 < ord {
+					n := make([]string, ord+1)
+					copy(n, params)
+					params = n
+					r.params[ticket] = params
+				}
+				params[ord] = kytheuri.ToString(e.GetTarget())
+			case edges.TParam:
+				tparams := r.tparams[ticket]
+				if len(tparams)-1 < ord {
+					n := make([]string, ord+1)
+					copy(n, tparams)
+					tparams = n
+					r.tparams[ticket] = tparams
+				}
+				tparams[ord] = kytheuri.ToString(e.GetTarget())
+			}
+		}
+	}
+	return r, nil
+}
+
+// Resolve returns the fully resolved MarkedSource for the given source VName.
+// May return nil if no MarkedSource is found.
+func (r *Resolver) Resolve(src *spb.VName) *cpb.MarkedSource {
+	return r.ResolveTicket(kytheuri.ToString(src))
+}
+
+// Resolve returns the fully resolved MarkedSource for the given source ticket.
+// May return nil if no MarkedSource is found.
+func (r *Resolver) ResolveTicket(ticket string) *cpb.MarkedSource {
+	if ticket == "" {
+		return nil
+	}
+	return r.resolve(ticket, r.nodes[ticket])
+}
+
+func (r *Resolver) resolve(ticket string, ms *cpb.MarkedSource) *cpb.MarkedSource {
+	if ms == nil {
+		return ms
+	}
+	switch ms.GetKind() {
+	case cpb.MarkedSource_LOOKUP_BY_PARAM:
+		// TODO: determine what to do when a lookup isn't found
+		params := r.params[ticket]
+		if p := params[ms.GetLookupIndex()]; p != "" {
+			return r.ResolveTicket(p)
+		}
+	case cpb.MarkedSource_LOOKUP_BY_TPARAM:
+		tparams := r.tparams[ticket]
+		if p := tparams[ms.GetLookupIndex()]; p != "" {
+			return r.ResolveTicket(p)
+		}
+	case cpb.MarkedSource_LOOKUP_BY_TYPED:
+		return r.ResolveTicket(r.typed[ticket])
+	case cpb.MarkedSource_PARAMETER_LOOKUP_BY_PARAM_WITH_DEFAULTS, cpb.MarkedSource_PARAMETER_LOOKUP_BY_PARAM:
+		// TODO: handle param/default
+		params := r.params[ticket][ms.GetLookupIndex():]
+		return r.resolveParameters(ms, params)
+	case cpb.MarkedSource_PARAMETER_LOOKUP_BY_TPARAM:
+		tparams := r.tparams[ticket][ms.GetLookupIndex():]
+		return r.resolveParameters(ms, tparams)
+	}
+	return r.resolveChildren(ticket, ms)
+}
+
+func (r *Resolver) resolveParameters(base *cpb.MarkedSource, params []string) *cpb.MarkedSource {
+	n := proto.Clone(base).(*cpb.MarkedSource)
+	n.LookupIndex = 0
+	n.Kind = cpb.MarkedSource_PARAMETER
+	n.Child = make([]*cpb.MarkedSource, len(params))
+	for i, p := range params {
+		n.Child[i] = r.ResolveTicket(p)
+	}
+	return n
+}
+
+func (r *Resolver) resolveChildren(ticket string, ms *cpb.MarkedSource) *cpb.MarkedSource {
+	n := proto.Clone(ms).(*cpb.MarkedSource)
+	children := make([]*cpb.MarkedSource, len(n.GetChild()))
+	for i, ms := range n.GetChild() {
+		children[i] = r.resolve(ticket, ms)
+	}
+	n.Child = children
+	return n
+}
diff --git a/kythe/go/util/tools/markedsource/BUILD b/kythe/go/util/tools/markedsource/BUILD
new file mode 100644
index 000000000..47a2314e9
--- /dev/null
+++ b/kythe/go/util/tools/markedsource/BUILD
@@ -0,0 +1,14 @@
+load("//tools:build_rules/shims.bzl", "go_binary")
+
+package(default_visibility = ["//kythe:default_visibility"])
+
+go_binary(
+    name = "markedsource",
+    srcs = ["markedsource.go"],
+    deps = [
+        "//kythe/go/storage/stream",
+        "//kythe/go/util/markedsource",
+        "//kythe/proto:storage_go_proto",
+        "@org_golang_google_protobuf//encoding/protojson:go_default_library",
+    ],
+)
diff --git a/kythe/go/util/tools/markedsource/markedsource.go b/kythe/go/util/tools/markedsource/markedsource.go
new file mode 100644
index 000000000..651c5df08
--- /dev/null
+++ b/kythe/go/util/tools/markedsource/markedsource.go
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2023 The Kythe Authors. All rights reserved.
+ *
+ * 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.
+ */
+
+// Binary markedsource provides a utility for printing the resolved MarkedSource
+// for any ticket from a delimited stream of wire-encoded kythe.proto.Entry
+// messages.
+//
+// Example usage:
+//
+//	markedsource kythe://kythe#myTicket < entries
+package main
+
+import (
+	"bufio"
+	"flag"
+	"log"
+	"os"
+
+	"kythe.io/kythe/go/storage/stream"
+	"kythe.io/kythe/go/util/markedsource"
+
+	"google.golang.org/protobuf/encoding/protojson"
+
+	spb "kythe.io/kythe/proto/storage_go_proto"
+)
+
+func main() {
+	flag.Parse()
+
+	in := bufio.NewReaderSize(os.Stdin, 2*4096)
+	rd := stream.NewReader(in)
+
+	var entries []*spb.Entry
+	if err := rd(func(e *spb.Entry) error {
+		entries = append(entries, e)
+		return nil
+	}); err != nil {
+		log.Fatalf("Failed to read entrystream: %v", err)
+	}
+
+	r, err := markedsource.NewResolver(entries)
+	if err != nil {
+		log.Fatalf("Failed to construct MarkedSource Resolver: %v", err)
+	}
+
+	wr := bufio.NewWriter(os.Stdout)
+	for _, ticket := range flag.Args() {
+		rec, err := protojson.Marshal(r.ResolveTicket(ticket))
+		if err != nil {
+			log.Fatalf("Error encoding MarkedSource: %v", err)
+		}
+		if _, err := wr.Write(rec); err != nil {
+			log.Fatalf("Error writing MarkedSource: %v", err)
+		}
+	}
+	if err := wr.Flush(); err != nil {
+		log.Fatalf("Error flushing stdout: %v", err)
+	}
+}


Test patch:
diff --git a/kythe/go/util/markedsource/resolve_test.go b/kythe/go/util/markedsource/resolve_test.go
new file mode 100644
index 000000000..c2033aa98
--- /dev/null
+++ b/kythe/go/util/markedsource/resolve_test.go
@@ -0,0 +1,138 @@
+/*
+ * Copyright 2023 The Kythe Authors. All rights reserved.
+ *
+ * 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 markedsource
+
+import (
+	"strconv"
+	"testing"
+
+	"kythe.io/kythe/go/test/testutil"
+	"kythe.io/kythe/go/util/compare"
+	"kythe.io/kythe/go/util/kytheuri"
+	"kythe.io/kythe/go/util/schema/edges"
+	"kythe.io/kythe/go/util/schema/facts"
+
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
+
+	cpb "kythe.io/kythe/proto/common_go_proto"
+	spb "kythe.io/kythe/proto/storage_go_proto"
+)
+
+func TestResolve(t *testing.T) {
+	tests := []struct {
+		Entries  []*spb.Entry
+		Ticket   string
+		Expected string
+	}{
+		{
+			[]*spb.Entry{code(t, "kythe:#test", `kind: IDENTIFIER pre_text: "simple"`)},
+			"kythe:#test",
+			`kind: IDENTIFIER pre_text: "simple"`,
+		},
+		{
+			[]*spb.Entry{
+				code(t, "kythe:#test", `kind: BOX child: { kind: LOOKUP_BY_PARAM lookup_index: 0 }`),
+				edge(t, "kythe:#test", edges.ParamIndex(0), "kythe:#param"),
+				code(t, "kythe:#param", `kind: IDENTIFIER pre_text: "param"`),
+			},
+			"kythe:#test",
+			`kind: BOX child: { kind: IDENTIFIER pre_text: "param" }`,
+		},
+		{
+			[]*spb.Entry{
+				code(t, "kythe:#test", `kind: BOX child: { kind: LOOKUP_BY_PARAM lookup_index: 0 }`),
+				edge(t, "kythe:#test", edges.ParamIndex(0), "kythe:#paramLevel0"),
+				code(t, "kythe:#paramLevel0", `kind: LOOKUP_BY_PARAM lookup_index: 1`),
+				edge(t, "kythe:#paramLevel0", edges.ParamIndex(1), "kythe:#paramLevel1"),
+				code(t, "kythe:#paramLevel1", `kind: IDENTIFIER pre_text: "deep"`),
+			},
+			"kythe:#test",
+			`kind: BOX child: { kind: IDENTIFIER pre_text: "deep" }`,
+		},
+		{
+			[]*spb.Entry{
+				code(t, "kythe:#test", `kind: PARAMETER_LOOKUP_BY_PARAM post_child_text: ", "`),
+				edge(t, "kythe:#test", edges.ParamIndex(0), "kythe:#param0"),
+				edge(t, "kythe:#test", edges.ParamIndex(1), "kythe:#param1"),
+				edge(t, "kythe:#test", edges.ParamIndex(2), "kythe:#param2"),
+				code(t, "kythe:#param0", `kind: IDENTIFIER pre_text: "0"`),
+				code(t, "kythe:#param1", `kind: IDENTIFIER pre_text: "1"`),
+				code(t, "kythe:#param2", `kind: IDENTIFIER pre_text: "2"`),
+			},
+			"kythe:#test",
+			`kind: PARAMETER post_child_text: ", " child: { kind: IDENTIFIER pre_text: "0" } child: { kind: IDENTIFIER pre_text: "1" } child: { kind: IDENTIFIER pre_text: "2" }`,
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			r, err := NewResolver(test.Entries)
+			testutil.Fatalf(t, "NewResolver: %v", err)
+			v := parseVName(t, test.Ticket)
+			found := r.Resolve(v)
+			expected := parseMarkedSource(t, test.Expected)
+
+			if diff := compare.ProtoDiff(expected, found); diff != "" {
+				t.Fatalf("(- expected; + found)\n%s", diff)
+			}
+		})
+	}
+}
+
+func code(t testing.TB, src, code string) *spb.Entry {
+	t.Helper()
+	v := parseVName(t, src)
+	ms := parseMarkedSource(t, code)
+	rec, err := proto.Marshal(ms)
+	testutil.Fatalf(t, "proto.Marshal: %v", err)
+	return &spb.Entry{
+		Source:    v,
+		FactName:  facts.Code,
+		FactValue: rec,
+	}
+}
+
+func edge(t testing.TB, src, kind, tgt string) *spb.Entry {
+	t.Helper()
+	return &spb.Entry{
+		Source:   parseVName(t, src),
+		EdgeKind: kind,
+		Target:   parseVName(t, tgt),
+	}
+}
+
+func parseVName(t testing.TB, s string) *spb.VName {
+	t.Helper()
+	v, err := kytheuri.ToVName(s)
+	testutil.Fatalf(t, "kytheuri.ParseVName: %v", err)
+	return v
+}
+
+func parseMarkedSource(t testing.TB, s string) *cpb.MarkedSource {
+	t.Helper()
+	var ms cpb.MarkedSource
+	testutil.Fatalf(t, "prototext.Unmarshal: %v", prototext.Unmarshal([]byte(s), &ms))
+	return &ms
+}
+
+func e(t testing.TB, s string) *spb.Entry {
+	t.Helper()
+	var e spb.Entry
+	testutil.Fatalf(t, "prototext.Unmarshal: %v", prototext.Unmarshal([]byte(s), &e))
+	return &e
+}