feat(columnar): add definitions for related node xrefs (#3228)

kythe/docs/rfc/2909.md

@@ -158,13 +158,15 @@ message PagedCrossReferences {
 
 ### Key-value format:
 
-| key | value | kind |
-|---|---|---|
-| `"xr"-source` | `CrossReferencesIndex` |  |
-| `"xr"-source-00-kind-file-start-end` | `srvpb.ExpandedAnchor` | Regular cross-reference |
-| `"xr"-source-10-kind-ordinal-node` | `srvpb.Node` | Related nodes |
-| `"xr"-source-20-caller` | `srvpb.ExpandedAnchor`+`MarkedSource` | Definition for Caller |
-| `"xr"-source-20-caller-file-start-end` | `srvpb.ExpandedAnchor` | Callsite |
+| key                                    | value                                 | kind                     |
+|----------------------------------------|---------------------------------------|--------------------------|
+| `"xr"-source`                          | `CrossReferencesIndex`                |                          |
+| `"xr"-source-00-kind-file-start-end`   | `srvpb.ExpandedAnchor`                | Regular cross-reference  |
+| `"xr"-source-10-kind-ordinal-node`     |                                       | Non-ref relations        |
+| `"xr"-source-20-caller`                | `srvpb.ExpandedAnchor`+`MarkedSource` | Definition for Caller    |
+| `"xr"-source-20-caller-file-start-end` | `srvpb.ExpandedAnchor`                | Callsite                 |
+| `"xr"-source-30-node`                  | `srvpb.Node`                          | Related nodes            |
+| `"xr"-source-40-node`                  | `srvpb.ExpandedAnchor`                | Related node definitions |
 
 
 Notes:

kythe/go/serving/pipeline/beam.go

@@ -46,13 +46,15 @@ import (
 )
 
 func init() {
+	beam.RegisterFunction(bareRevEdge)
 	beam.RegisterFunction(callEdge)
 	beam.RegisterFunction(combineEdgesIndex)
 	beam.RegisterFunction(completeDocument)
 	beam.RegisterFunction(constructCaller)
 	beam.RegisterFunction(defToDecorPiece)
 	beam.RegisterFunction(edgeTargets)
 	beam.RegisterFunction(edgeToCrossRefRelation)
+	beam.RegisterFunction(emitRelatedDefs)
 	beam.RegisterFunction(fileToDecorPiece)
 	beam.RegisterFunction(filterAnchorNodes)
 	beam.RegisterFunction(groupCrossRefs)
@@ -145,18 +147,48 @@ func (k *KytheBeam) SplitCrossReferences() beam.PCollection {
 		beam.ParDo(s, splitEdge, filter.Distinct(s, beam.ParDo(s, callEdge, callsites))),
 	))
 
-	// TODO(schroederc): need to add definitions to related nodes
-	relations := beam.ParDo(s, edgeToCrossRefRelation, k.edgeRelations())
+	edges := k.edgeRelations()
+	relatedDefs := beam.ParDo(s, emitRelatedDefs, beam.CoGroupByKey(s,
+		k.directDefinitions(),
+		beam.ParDo(s, splitEdge, filter.Distinct(s, beam.ParDo(s, bareRevEdge, edges))),
+	))
+	relations := beam.ParDo(s, edgeToCrossRefRelation, edges)
 
 	return beam.ParDo(s, encodeCrossRef, beam.Flatten(s,
 		idx,
 		refs,
 		relations,
+		relatedDefs,
 		callers,
 		callsites,
 	))
 }
 
+func emitRelatedDefs(target *spb.VName, defStream func(**srvpb.ExpandedAnchor) bool, srcStream func(**spb.VName) bool, emit func(*xspb.CrossReferences)) {
+	var def *srvpb.ExpandedAnchor
+	if !defStream(&def) {
+		return // no related node definition found
+	}
+	nodeDef := &xspb.CrossReferences_NodeDefinition_{&xspb.CrossReferences_NodeDefinition{
+		Node:     target,
+		Location: def,
+	}}
+
+	var src *spb.VName
+	for srcStream(&src) {
+		emit(&xspb.CrossReferences{Source: src, Entry: nodeDef})
+	}
+}
+
+func bareRevEdge(eg *gspb.Edges, emit func(*scpb.Edge)) error {
+	switch e := eg.Entry.(type) {
+	case *gspb.Edges_Edge_:
+		edge := e.Edge
+		emit(&scpb.Edge{Target: eg.Source, Source: edge.Target})
+	}
+	return nil
+}
+
 func constructCaller(caller *spb.VName, defStream func(**srvpb.ExpandedAnchor) bool, msStream func(**cpb.MarkedSource) bool, calleeStream func(**spb.VName) bool, emit func(*xspb.CrossReferences)) {
 	var def *srvpb.ExpandedAnchor
 	if !defStream(&def) {

kythe/go/serving/pipeline/beam_integration_test.go

@@ -339,6 +339,9 @@ func TestServingSimpleCrossReferences(t *testing.T) {
 		Edge: []*scpb.Edge{{
 			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_DEFINES_BINDING},
 			Target: &spb.VName{Signature: "caller"},
+		}, {
+			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_DEFINES_BINDING},
+			Target: &spb.VName{Signature: "interface"},
 		}, {
 			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_REF},
 			Target: src,
@@ -493,6 +496,59 @@ func TestServingSimpleCrossReferences(t *testing.T) {
 		},
 	}))
 
+	t.Run("node_definitions", makeXRefTestCase(ctx, xs, &xpb.CrossReferencesRequest{
+		Ticket:          []string{ticket},
+		Filter:          []string{facts.NodeKind},
+		NodeDefinitions: true,
+	}, &xpb.CrossReferencesReply{
+		CrossReferences: map[string]*xpb.CrossReferencesReply_CrossReferenceSet{
+			ticket: {
+				Ticket:       ticket,
+				MarkedSource: ms,
+				RelatedNode: []*xpb.CrossReferencesReply_RelatedNode{{
+					Ticket:       "kythe:#interface",
+					RelationKind: edges.Extends,
+				}},
+			},
+		},
+		Nodes: map[string]*cpb.NodeInfo{
+			ticket: {Facts: map[string][]byte{facts.NodeKind: []byte(nodes.Record)}},
+			"kythe:#interface": {
+				Facts:      map[string][]byte{facts.NodeKind: []byte(nodes.Interface)},
+				Definition: "kythe:?path=path#anchor1",
+			},
+		},
+		DefinitionLocations: map[string]*xpb.Anchor{
+			"kythe:?path=path#anchor1": {
+				Ticket: "kythe:?path=path#anchor1",
+				Parent: "kythe:?path=path",
+				Span: &cpb.Span{
+					Start: &cpb.Point{
+						ByteOffset:   5,
+						ColumnOffset: 5,
+						LineNumber:   1,
+					},
+					End: &cpb.Point{
+						ByteOffset:   9,
+						ColumnOffset: 9,
+						LineNumber:   1,
+					},
+				},
+				Snippet: "blah blah",
+				SnippetSpan: &cpb.Span{
+					Start: &cpb.Point{
+						LineNumber: 1,
+					},
+					End: &cpb.Point{
+						ByteOffset:   9,
+						ColumnOffset: 9,
+						LineNumber:   1,
+					},
+				},
+			},
+		},
+	}))
+
 	t.Run("call_refs", makeXRefTestCase(ctx, xs, &xpb.CrossReferencesRequest{
 		Ticket:        []string{ticket},
 		ReferenceKind: xpb.CrossReferencesRequest_CALL_REFERENCES,

kythe/go/serving/xrefs/columnar.go

@@ -256,6 +256,9 @@ func (c *ColumnarTable) CrossReferences(ctx context.Context, req *xpb.CrossRefer
 	if len(patterns) > 0 {
 		reply.Nodes = make(map[string]*cpb.NodeInfo)
 	}
+	if req.NodeDefinitions {
+		reply.DefinitionLocations = make(map[string]*xpb.Anchor)
+	}
 	emitSnippets := req.Snippets != xpb.SnippetsKind_NONE
 
 	// TODO(schroederc): implement paging xrefs in large CrossReferencesReply messages
@@ -361,6 +364,18 @@ func (c *ColumnarTable) CrossReferences(ctx context.Context, req *xpb.CrossRefer
 				if relatedNodes.Contains(relatedNode) {
 					addXRefNode(reply, patterns, e.RelatedNode.Node)
 				}
+			case *xspb.CrossReferences_NodeDefinition_:
+				if !req.NodeDefinitions || len(reply.Nodes) == 0 {
+					continue
+				}
+
+				relatedNode := kytheuri.ToString(e.NodeDefinition.Node)
+				if node := reply.Nodes[relatedNode]; node != nil {
+					loc := e.NodeDefinition.Location
+					node.Definition = loc.Ticket
+					a := a2a(loc, emitSnippets).Anchor
+					reply.DefinitionLocations[loc.Ticket] = a
+				}
 			case *xspb.CrossReferences_Caller_:
 				if req.CallerKind == xpb.CrossReferencesRequest_NO_CALLERS {
 					continue

kythe/go/serving/xrefs/columnar/columnar_encoding.go

@@ -377,11 +377,12 @@ func decodeDecorDiagnostic(file *spb.VName, key string, val []byte) (*xspb.FileD
 // Columnar file decorations group numbers.
 // See: kythe/proto/xref_serving.proto
 const (
-	columnarXRefsIndexGroup       = -1 // no group number
-	columnarXRefsReferenceGroup   = 0
-	columnarXRefsRelationGroup    = 10
-	columnarXRefsCallerGroup      = 20
-	columnarXRefsRelatedNodeGroup = 30
+	columnarXRefsIndexGroup          = -1 // no group number
+	columnarXRefsReferenceGroup      = 0
+	columnarXRefsRelationGroup       = 10
+	columnarXRefsCallerGroup         = 20
+	columnarXRefsRelatedNodeGroup    = 30
+	columnarXRefsNodeDefinitionGroup = 40
 )
 
 // EncodeCrossReferencesEntry encodes a columnar CrossReferences entry.
@@ -399,6 +400,8 @@ func EncodeCrossReferencesEntry(keyPrefix []byte, xr *xspb.CrossReferences) (*KV
 		return encodeXRefCallsite(keyPrefix, xr.Source, e.Callsite)
 	case *xspb.CrossReferences_RelatedNode_:
 		return encodeXRefRelatedNode(keyPrefix, xr.Source, e.RelatedNode)
+	case *xspb.CrossReferences_NodeDefinition_:
+		return encodeXRefNodeDefinition(keyPrefix, xr.Source, e.NodeDefinition)
 	default:
 		return nil, fmt.Errorf("unknown CrossReferences entry: %T", e)
 	}
@@ -504,6 +507,18 @@ func encodeXRefRelatedNode(prefix []byte, src *spb.VName, n *xspb.CrossReference
 	return &KV{key, val}, nil
 }
 
+func encodeXRefNodeDefinition(prefix []byte, src *spb.VName, def *xspb.CrossReferences_NodeDefinition) (*KV, error) {
+	key, err := keys.Append(prefix, src, columnarXRefsNodeDefinitionGroup, def.Node)
+	if err != nil {
+		return nil, err
+	}
+	val, err := proto.Marshal(&xspb.CrossReferences_NodeDefinition{Location: def.Location})
+	if err != nil {
+		return nil, err
+	}
+	return &KV{key, val}, nil
+}
+
 // DecodeCrossReferencesEntry decodes a columnar CrossReferences entry.
 func DecodeCrossReferencesEntry(src *spb.VName, key string, val []byte) (*xspb.CrossReferences, error) {
 	kind := columnarXRefsIndexGroup
@@ -525,6 +540,8 @@ func DecodeCrossReferencesEntry(src *spb.VName, key string, val []byte) (*xspb.C
 		return decodeXRefCallerOrCallsite(src, key, val)
 	case columnarXRefsRelatedNodeGroup:
 		return decodeXRefRelatedNode(src, key, val)
+	case columnarXRefsNodeDefinitionGroup:
+		return decodeXRefNodeDefinition(src, key, val)
 	default:
 		return nil, fmt.Errorf("unknown group kind: %d", kind)
 	}
@@ -642,3 +659,23 @@ func decodeXRefRelatedNode(src *spb.VName, key string, val []byte) (*xspb.CrossR
 		Entry:  &xspb.CrossReferences_RelatedNode_{&rn},
 	}, nil
 }
+
+func decodeXRefNodeDefinition(src *spb.VName, key string, val []byte) (*xspb.CrossReferences, error) {
+	var node spb.VName
+	key, err := keys.Parse(key, &node)
+	if err != nil {
+		return nil, err
+	}
+	if len(key) != 0 {
+		return nil, fmt.Errorf("unexpected trailing key: %q", key)
+	}
+	var def xspb.CrossReferences_NodeDefinition
+	if err := proto.Unmarshal(val, &def); err != nil {
+		return nil, err
+	}
+	def.Node = &node
+	return &xspb.CrossReferences{
+		Source: src,
+		Entry:  &xspb.CrossReferences_NodeDefinition_{&def},
+	}, nil
+}

kythe/go/serving/xrefs/columnar/columnar_test.go

@@ -217,6 +217,18 @@ func TestCrossReferencesEncodingRoundtrip(t *testing.T) {
 				Source: &spb.VName{Signature: "relatedNode"},
 			},
 		}},
+	}, {
+		Source: src,
+		Entry: &xspb.CrossReferences_NodeDefinition_{&xspb.CrossReferences_NodeDefinition{
+			Node: &spb.VName{Signature: "relatedNode"},
+			Location: &srvpb.ExpandedAnchor{
+				Ticket: "kythe:#relatedNodeDef",
+				Span: &cpb.Span{
+					Start: &cpb.Point{ByteOffset: 32},
+					End:   &cpb.Point{ByteOffset: 256},
+				},
+			},
+		}},
 	}}
 
 	for _, test := range tests {

kythe/go/serving/xrefs/columnar_test.go

@@ -26,6 +26,7 @@ import (
 	"kythe.io/kythe/go/storage/keyvalue"
 	"kythe.io/kythe/go/util/compare"
 	"kythe.io/kythe/go/util/kytheuri"
+	"kythe.io/kythe/go/util/schema/facts"
 
 	cpb "kythe.io/kythe/proto/common_go_proto"
 	scpb "kythe.io/kythe/proto/schema_go_proto"
@@ -439,6 +440,15 @@ func TestServingCrossReferences(t *testing.T) {
 				Kind:   &scpb.Node_KytheKind{scpb.NodeKind_FUNCTION},
 			},
 		}},
+	}, {
+		Source: src,
+		Entry: &xspb.CrossReferences_NodeDefinition_{&xspb.CrossReferences_NodeDefinition{
+			Node: &spb.VName{Signature: "relatedNode"},
+			Location: &srvpb.ExpandedAnchor{
+				Ticket: "kythe:#relatedNodeDef",
+				Span:   span,
+			},
+		}},
 	}}
 	for _, xr := range xrefs {
 		mustWriteXRef(t, w, xr)
@@ -601,6 +611,37 @@ func TestServingCrossReferences(t *testing.T) {
 		},
 	}))
 
+	t.Run("node_definitions", makeXRefTestCase(ctx, xs, &xpb.CrossReferencesRequest{
+		Ticket:          []string{ticket},
+		Filter:          []string{facts.NodeKind},
+		NodeDefinitions: true,
+	}, &xpb.CrossReferencesReply{
+		CrossReferences: map[string]*xpb.CrossReferencesReply_CrossReferenceSet{
+			ticket: {
+				Ticket:       ticket,
+				MarkedSource: ms,
+				RelatedNode: []*xpb.CrossReferencesReply_RelatedNode{{
+					Ticket:       "kythe:#relatedNode",
+					RelationKind: "%/kythe/edge/childof",
+				}},
+			},
+		},
+		Nodes: map[string]*cpb.NodeInfo{
+			ticket: {Facts: map[string][]byte{"/kythe/node/kind": []byte("record")}},
+			"kythe:#relatedNode": {
+				Facts:      map[string][]byte{"/kythe/node/kind": []byte("function")},
+				Definition: "kythe:#relatedNodeDef",
+			},
+		},
+		DefinitionLocations: map[string]*xpb.Anchor{
+			"kythe:#relatedNodeDef": {
+				Ticket: "kythe:#relatedNodeDef",
+				Parent: "kythe:",
+				Span:   span,
+			},
+		},
+	}))
+
 	t.Run("non_call_refs", makeXRefTestCase(ctx, xs, &xpb.CrossReferencesRequest{
 		Ticket:        []string{ticket},
 		ReferenceKind: xpb.CrossReferencesRequest_NON_CALL_REFERENCES,

kythe/proto/xref_serving.proto

@@ -151,6 +151,7 @@ message CrossReferences {
     Caller caller = 5;
     Callsite callsite = 6;
     RelatedNode related_node = 7;
+    NodeDefinition node_definition = 8;
   }
 
   // Index for columnar cross-references data.
@@ -227,4 +228,12 @@ message CrossReferences {
   message RelatedNode {
     kythe.proto.schema.Node node = 1;
   }
+
+  // Definition location for each RelatedNode
+  //
+  // Columnar key: 40-node
+  message NodeDefinition {
+    kythe.proto.VName node = 1;
+    kythe.proto.serving.ExpandedAnchor location = 2;
+  }
 }