feat(post_processing): implement extends_overrides. (#4995)

kythe/go/serving/pipeline/beam.go

@@ -73,6 +73,8 @@ func init() {
 	beam.RegisterFunction(nodeToDocs)
 	beam.RegisterFunction(nodeToEdges)
 	beam.RegisterFunction(nodeToReverseEdges)
+	beam.RegisterFunction(overriddenToDecor)
+	beam.RegisterFunction(overridingToFile)
 	beam.RegisterFunction(parseMarkedSource)
 	beam.RegisterFunction(refToCallsite)
 	beam.RegisterFunction(refToCrossRef)
@@ -419,10 +421,57 @@ func (k *KytheBeam) decorationPieces(s beam.Scope) beam.PCollection {
 		beam.CoGroupByKey(s, beam.ParDo(s, moveSourceToKey, bareNodes), targets))
 	defs := beam.ParDo(s, defToDecorPiece,
 		beam.CoGroupByKey(s, k.directDefinitions(), targets))
-	// TODO(schroederc): overrides
+	overrides := k.overrides(targets)
 	decorDiagnostics := k.diagnostics()
 
-	return beam.Flatten(s, decor, files, targetNodes, defs, decorDiagnostics)
+	return beam.Flatten(s, decor, files, targetNodes, defs, decorDiagnostics, overrides)
+}
+
+func (k *KytheBeam) overrides(targets beam.PCollection) beam.PCollection {
+	s := k.s.Scope("Overrides")
+	overriddenToEdge := beam.Seq(s, k.Nodes(), &nodes.Filter{IncludeEdges: []string{edges.Overrides, edges.Extends, edges.OverridesTransitive, edges.Satisfies}}, nodeToEdges)
+	overridingToDecor := beam.ParDo(s, overriddenToDecor, beam.CoGroupByKey(s, k.directDefinitions(), overriddenToEdge))
+	return beam.ParDo(s, overridingToFile, beam.CoGroupByKey(s, targets, overridingToDecor))
+}
+
+func overriddenToDecor(overridden *spb.VName, overriddenAnchors func(**srvpb.ExpandedAnchor) bool, edgeStream func(**scpb.Edge) bool, emit func(*spb.VName, *ppb.DecorationPiece)) {
+	var overriddenAnchor *srvpb.ExpandedAnchor
+	var e *scpb.Edge
+	if !overriddenAnchors(&overriddenAnchor) {
+		return
+	}
+	for edgeStream(&e) {
+		var kind srvpb.FileDecorations_Override_Kind
+		edgeKindString := schema.EdgeKindString(e.GetKytheKind())
+		if edges.IsVariant(edgeKindString, edges.Overrides) {
+			kind = srvpb.FileDecorations_Override_OVERRIDES
+		} else if edges.IsVariant(edgeKindString, edges.Extends) || edges.IsVariant(edgeKindString, edges.Satisfies) {
+			kind = srvpb.FileDecorations_Override_EXTENDS
+		} else {
+			continue
+		}
+		emit(e.Source, &ppb.DecorationPiece{
+			Piece: &ppb.DecorationPiece_TargetOverride{
+				TargetOverride: &xspb.FileDecorations_TargetOverride{
+					Overriding:           e.Source,
+					Overridden:           e.Target,
+					Kind:                 kind,
+					OverridingDefinition: overriddenAnchor,
+				},
+			},
+		})
+	}
+}
+
+func overridingToFile(target *spb.VName, files func(**spb.VName) bool, overrides func(**ppb.DecorationPiece) bool, emit func(*spb.VName, *ppb.DecorationPiece)) {
+	var file *spb.VName
+	if !files(&file) {
+		return
+	}
+	var override *ppb.DecorationPiece
+	for overrides(&override) {
+		emit(file, override)
+	}
 }
 
 func (k *KytheBeam) diagnostics() beam.PCollection {
@@ -580,6 +629,13 @@ func (c *combineDecorPieces) AddInput(accum *srvpb.FileDecorations, p *ppb.Decor
 		})
 	case *ppb.DecorationPiece_Diagnostic:
 		accum.Diagnostic = append(accum.Diagnostic, p.Diagnostic)
+	case *ppb.DecorationPiece_TargetOverride:
+		accum.TargetOverride = append(accum.TargetOverride, &srvpb.FileDecorations_Override{
+			Overriding:           kytheuri.ToString(p.TargetOverride.Overriding),
+			Overridden:           kytheuri.ToString(p.TargetOverride.Overridden),
+			OverriddenDefinition: p.TargetOverride.OverridingDefinition.Ticket,
+			Kind:                 p.TargetOverride.Kind,
+		})
 	default:
 		panic(fmt.Errorf("unhandled DecorationPiece: %T", p))
 	}

kythe/go/serving/pipeline/beam_test.go

@@ -451,6 +451,119 @@ func TestDecorations_targetDefinition(t *testing.T) {
 	}
 }
 
+func TestDecorations_overrides(t *testing.T) {
+	testNodes := []*scpb.Node{{
+		Source: &spb.VName{Path: "path"},
+		Kind:   &scpb.Node_KytheKind{scpb.NodeKind_FILE},
+		Fact: []*scpb.Fact{{
+			Name:  &scpb.Fact_KytheName{scpb.FactName_TEXT},
+			Value: []byte("some text\n"),
+		}},
+	}, {
+		Source: &spb.VName{Path: "path2", Signature: "def1"},
+		Kind:   &scpb.Node_KytheKind{scpb.NodeKind_ANCHOR},
+		Fact: []*scpb.Fact{{
+			Name:  &scpb.Fact_KytheName{scpb.FactName_LOC_START},
+			Value: []byte("5"),
+		}, {
+			Name:  &scpb.Fact_KytheName{scpb.FactName_LOC_END},
+			Value: []byte("8"),
+		}},
+		Edge: []*scpb.Edge{{
+			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_DEFINES_BINDING},
+			Target: &spb.VName{Signature: "node1"},
+		}},
+	}, {
+		Source: &spb.VName{Path: "path2", Signature: "def2"},
+		Kind:   &scpb.Node_KytheKind{scpb.NodeKind_ANCHOR},
+		Fact: []*scpb.Fact{{
+			Name:  &scpb.Fact_KytheName{scpb.FactName_LOC_START},
+			Value: []byte("4"),
+		}, {
+			Name:  &scpb.Fact_KytheName{scpb.FactName_LOC_END},
+			Value: []byte("7"),
+		}},
+		Edge: []*scpb.Edge{{
+			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_DEFINES_BINDING},
+			Target: &spb.VName{Signature: "node2"},
+		}},
+	}, {
+		Source: &spb.VName{Signature: "node2"},
+		Kind:   &scpb.Node_KytheKind{scpb.NodeKind_FUNCTION},
+		Edge: []*scpb.Edge{{
+			Kind:   &scpb.Edge_KytheKind{scpb.EdgeKind_OVERRIDES},
+			Target: &spb.VName{Signature: "node1"},
+		}},
+	}, {
+		Source: &spb.VName{Path: "path2"},
+		Kind:   &scpb.Node_KytheKind{scpb.NodeKind_FILE},
+		Fact: []*scpb.Fact{{
+			Name:  &scpb.Fact_KytheName{scpb.FactName_TEXT},
+			Value: []byte("some def\n"),
+		}},
+	}}
+
+	expected := []*srvpb.FileDecorations{{
+		File: &srvpb.File{Text: []byte("some def\n")},
+		Decoration: []*srvpb.FileDecorations_Decoration{{
+			Anchor:           &srvpb.RawAnchor{StartOffset: 4, EndOffset: 7},
+			Kind:             "/kythe/edge/defines/binding",
+			Target:           "kythe:#node2",
+			TargetDefinition: "kythe:?path=path2#def2",
+		}, {
+			Anchor:           &srvpb.RawAnchor{StartOffset: 5, EndOffset: 8},
+			Kind:             "/kythe/edge/defines/binding",
+			Target:           "kythe:#node1",
+			TargetDefinition: "kythe:?path=path2#def1",
+		}},
+		Target: []*srvpb.Node{{
+			Ticket: "kythe:#node2",
+			Fact:   []*cpb.Fact{{Name: "/kythe/node/kind", Value: []byte("function")}},
+		}},
+		TargetDefinitions: []*srvpb.ExpandedAnchor{{
+			Ticket: "kythe:?path=path2#def1",
+			Text:   "def",
+			Span: &cpb.Span{
+				Start: &cpb.Point{ByteOffset: 5, LineNumber: 1, ColumnOffset: 5},
+				End:   &cpb.Point{ByteOffset: 8, LineNumber: 1, ColumnOffset: 8},
+			},
+			Snippet: "some def",
+			SnippetSpan: &cpb.Span{
+				Start: &cpb.Point{LineNumber: 1},
+				End:   &cpb.Point{ByteOffset: 8, LineNumber: 1, ColumnOffset: 8},
+			},
+		}, {
+			Ticket: "kythe:?path=path2#def2",
+			Text:   " de",
+			Span: &cpb.Span{
+				Start: &cpb.Point{ByteOffset: 4, LineNumber: 1, ColumnOffset: 4},
+				End:   &cpb.Point{ByteOffset: 7, LineNumber: 1, ColumnOffset: 7},
+			},
+			Snippet: "some def",
+			SnippetSpan: &cpb.Span{
+				Start: &cpb.Point{LineNumber: 1},
+				End:   &cpb.Point{ByteOffset: 8, LineNumber: 1, ColumnOffset: 8},
+			},
+		}},
+		TargetOverride: []*srvpb.FileDecorations_Override{{
+			Overriding:           "kythe:#node2",
+			Overridden:           "kythe:#node1",
+			OverriddenDefinition: "kythe:?path=path2#def1",
+		}},
+	}, {
+		File: &srvpb.File{Text: []byte("some text\n")},
+	}}
+
+	p, s, nodes := ptest.CreateList(testNodes)
+	decor := FromNodes(s, nodes).Decorations()
+	debug.Print(s, decor)
+	passert.Equals(s, beam.DropKey(s, decor), beam.CreateList(s, expected))
+
+	if err := ptest.Run(p); err != nil {
+		t.Fatalf("Pipeline error: %+v", err)
+	}
+}
+
 func TestCrossReferences(t *testing.T) {
 	testNodes := []*scpb.Node{{
 		Source: &spb.VName{Signature: "node1"},

kythe/go/serving/pipeline/encoding.go

@@ -96,12 +96,27 @@ func encodeDecorPiece(file *spb.VName, p *ppb.DecorationPiece, emit func([]byte,
 		return encodeDecorDef(file, p.Definition, emit)
 	case *ppb.DecorationPiece_Diagnostic:
 		return encodeDecorDiagnostic(file, p.Diagnostic, emit)
+	case *ppb.DecorationPiece_TargetOverride:
+		return encodeDecorTargetOverride(file, p.TargetOverride, emit)
 	default:
-		// TODO(schroederc): add overrides
 		return fmt.Errorf("unknown DecorationPiece: %T", p)
 	}
 }
 
+func encodeDecorTargetOverride(file *spb.VName, o *xspb.FileDecorations_TargetOverride, emit func([]byte, []byte)) error {
+	e, err := columnar.EncodeDecorationsEntry(columnar.DecorationsKeyPrefix, &xspb.FileDecorations{
+		File: file,
+		Entry: &xspb.FileDecorations_TargetOverride_{
+			TargetOverride: o,
+		},
+	})
+	if err != nil {
+		return err
+	}
+	emit(e.Key, e.Value)
+	return nil
+}
+
 func encodeDecorFile(file *spb.VName, f *srvpb.File, emit func([]byte, []byte)) error {
 	// Emit FileDecorations Index
 	e, err := columnar.EncodeDecorationsEntry(columnar.DecorationsKeyPrefix, &xspb.FileDecorations{

kythe/go/serving/xrefs/columnar.go

@@ -41,6 +41,7 @@ import (
 
 	cpb "kythe.io/kythe/proto/common_go_proto"
 	scpb "kythe.io/kythe/proto/schema_go_proto"
+	"kythe.io/kythe/proto/serving_go_proto"
 	xpb "kythe.io/kythe/proto/xref_go_proto"
 	xspb "kythe.io/kythe/proto/xref_serving_go_proto"
 )
@@ -126,6 +127,7 @@ func (c *ColumnarTable) Decorations(ctx context.Context, req *xpb.DecorationsReq
 	var norm *span.Normalizer                                          // span normalizer for references
 	refsByTarget := make(map[string][]*xpb.DecorationsReply_Reference) // target -> set<Reference>
 	defs := stringset.New()                                            // set<needed definition tickets>
+	reply.ExtendsOverrides = make(map[string]*xpb.DecorationsReply_Overrides)
 	buildConfigs := stringset.New(req.BuildConfig...)
 	patterns := xrefs.ConvertFilters(req.Filter)
 	emitSnippets := req.Snippets != xpb.SnippetsKind_NONE
@@ -216,7 +218,27 @@ func (c *ColumnarTable) Decorations(ctx context.Context, req *xpb.DecorationsReq
 			refsByTarget[ref.TargetTicket] = append(refsByTarget[ref.TargetTicket], ref)
 			reply.Reference = append(reply.Reference, ref)
 		case *xspb.FileDecorations_TargetOverride_:
-			// TODO(schroederc): handle
+			overridingTicket := kytheuri.ToString(e.TargetOverride.Overriding)
+			t, ok := reply.ExtendsOverrides[overridingTicket]
+			if !ok {
+				t = &xpb.DecorationsReply_Overrides{}
+				reply.ExtendsOverrides[overridingTicket] = t
+			}
+			var kind xpb.DecorationsReply_Override_Kind
+			switch e.TargetOverride.Kind {
+			case serving_go_proto.FileDecorations_Override_OVERRIDES:
+				kind = xpb.DecorationsReply_Override_OVERRIDES
+			case serving_go_proto.FileDecorations_Override_EXTENDS:
+				kind = xpb.DecorationsReply_Override_EXTENDS
+			}
+			t.Override = append(t.Override, &xpb.DecorationsReply_Override{
+				Target:           kytheuri.ToString(e.TargetOverride.Overridden),
+				Kind:             kind,
+				TargetDefinition: e.TargetOverride.OverridingDefinition.Ticket,
+			})
+			if req.TargetDefinitions {
+				reply.DefinitionLocations[e.TargetOverride.OverridingDefinition.Ticket] = a2a(e.TargetOverride.OverridingDefinition, nil, false).Anchor
+			}
 		case *xspb.FileDecorations_TargetNode_:
 			if len(patterns) == 0 {
 				// TODO(schroederc): seek to next group

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

@@ -136,7 +136,7 @@ func encodeDecorTargetOverride(prefix []byte, file *spb.VName, to *xspb.FileDeco
 	if err != nil {
 		return nil, err
 	}
-	val, err := proto.Marshal(&xspb.FileDecorations_TargetOverride{})
+	val, err := proto.Marshal(to)
 	if err != nil {
 		return nil, err
 	}
@@ -295,13 +295,13 @@ func decodeDecorTargetOverride(file *spb.VName, key string, val []byte) (*xspb.F
 	} else if key != "" {
 		return nil, fmt.Errorf("unexpected TargetOverride key suffix: %q", key)
 	}
+	var to xspb.FileDecorations_TargetOverride
+	if err = proto.Unmarshal(val, &to); err != nil {
+		return nil, err
+	}
 	return &xspb.FileDecorations{
-		File: file,
-		Entry: &xspb.FileDecorations_TargetOverride_{&xspb.FileDecorations_TargetOverride{
-			Kind:       xspb.FileDecorations_TargetOverride_Kind(kind),
-			Overridden: &overridden,
-			Overriding: &overriding,
-		}},
+		File:  file,
+		Entry: &xspb.FileDecorations_TargetOverride_{&to},
 	}, nil
 }
 

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

@@ -61,7 +61,7 @@ func TestDecorationsEncodingRoundtrip(t *testing.T) {
 	}, {
 		File: file,
 		Entry: &xspb.FileDecorations_TargetOverride_{&xspb.FileDecorations_TargetOverride{
-			Kind:       xspb.FileDecorations_TargetOverride_EXTENDS,
+			Kind:       srvpb.FileDecorations_Override_EXTENDS,
 			Overridden: &spb.VName{Signature: "overridden"},
 			Overriding: &spb.VName{Signature: "overriding"},
 		}},

kythe/go/util/schema/edges/edges.go

@@ -42,6 +42,7 @@ const (
 	Generates               = Prefix + "generates"
 	Named                   = Prefix + "named"
 	Overrides               = Prefix + "overrides"
+	OverridesTransitive     = Prefix + "overrides/transitive"
 	Param                   = Prefix + "param"
 	Satisfies               = Prefix + "satisfies"
 	Typed                   = Prefix + "typed"

kythe/proto/BUILD

@@ -640,6 +640,7 @@ proto_library(
         ":schema_proto",
         ":serving_proto",
         ":storage_proto",
+        ":xref_serving_proto",
     ],
 )
 
@@ -650,6 +651,7 @@ go_kythe_proto(
         ":schema_go_proto",
         ":serving_go_proto",
         ":storage_go_proto",
+        ":xref_serving_go_proto",
     ],
 )
 

kythe/proto/pipeline.proto

@@ -25,6 +25,7 @@ import "kythe/proto/common.proto";
 import "kythe/proto/schema.proto";
 import "kythe/proto/serving.proto";
 import "kythe/proto/storage.proto";
+import "kythe/proto/xref_serving.proto";
 
 message Reference {
   kythe.proto.VName source = 1;
@@ -54,6 +55,7 @@ message DecorationPiece {
     kythe.proto.schema.Node node = 4;
     Definition definition = 5;
     kythe.proto.common.Diagnostic diagnostic = 6;
-    // TODO(schroederc): overrides
+    kythe.proto.serving.xrefs.FileDecorations.TargetOverride target_override =
+        7;
   }
 }