Repository: kythe/kythe
Pull request #5434 (2022-11-08 22:23:16)

perf: update patcher to use binary search


Patch:
diff --git a/kythe/go/util/span/BUILD b/kythe/go/util/span/BUILD
index 536169ddd..8b609137e 100644
--- a/kythe/go/util/span/BUILD
+++ b/kythe/go/util/span/BUILD
@@ -23,6 +23,7 @@ go_test(
     deps = [
         "//kythe/go/test/testutil",
         "//kythe/go/util/compare",
+        "@com_github_google_go_cmp//cmp/cmpopts:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",
     ],
 )
diff --git a/kythe/go/util/span/span.go b/kythe/go/util/span/span.go
index 6351836c7..e7cd0ae3a 100644
--- a/kythe/go/util/span/span.go
+++ b/kythe/go/util/span/span.go
@@ -126,10 +126,22 @@ func Unmarshal(rec []byte) (*Patcher, error) {
 		default:
 			return nil, fmt.Errorf("unknown diff type: %s", db.SpanType[i])
 		}
+		if i != 0 {
+			updatePrefix(&spans[i-1], &spans[i])
+		}
 	}
 	return &Patcher{spans}, nil
 }
 
+func updatePrefix(prev, d *diff) {
+	d.oldPrefix = prev.oldPrefix
+	d.newPrefix = prev.newPrefix
+	d.oldPrefix.Type = del
+	d.newPrefix.Type = ins
+	d.oldPrefix.Update(*prev)
+	d.newPrefix.Update(*prev)
+}
+
 type diff struct {
 	Length int32
 	Type   diffmatchpatch.Operation
@@ -137,6 +149,8 @@ type diff struct {
 	Newlines     int32
 	FirstNewline int32
 	LastNewline  int32
+
+	oldPrefix, newPrefix offsetTracker
 }
 
 const (
@@ -168,6 +182,9 @@ func mapToOffsets(ds []diffmatchpatch.Diff) []diff {
 			FirstNewline: int32(first),
 			LastNewline:  int32(last),
 		}
+		if i != 0 {
+			updatePrefix(&res[i-1], &res[i])
+		}
 	}
 	return res
 }
@@ -206,41 +223,37 @@ func (p *Patcher) PatchSpan(s *cpb.Span) (span *cpb.Span, exists bool) {
 		return s, true
 	}
 
-	oldT := offsetTracker{Type: del}
-	newT := offsetTracker{Type: ins}
-
-	for _, d := range p.spans {
-		if oldT.Offset > spanStart {
-			return nil, false
-		}
-
-		if d.Type == eq && oldT.Offset <= spanStart && spanEnd <= oldT.Offset+d.Length {
-			// The span is located within the equal region.  It exists.
-			lineDiff := newT.Lines - oldT.Lines
-			colDiff := newT.ColumnOffset - oldT.ColumnOffset
-			if d.FirstNewline != -1 && spanStart-oldT.Offset >= d.FirstNewline {
-				// The given span is past the first newline so it has no column diff.
-				colDiff = 0
-			}
-			return &cpb.Span{
-				Start: &cpb.Point{
-					ByteOffset:   newT.Offset + (spanStart - oldT.Offset),
-					ColumnOffset: s.GetStart().GetColumnOffset() + colDiff,
-					LineNumber:   s.GetStart().GetLineNumber() + lineDiff,
-				},
-				End: &cpb.Point{
-					ByteOffset:   newT.Offset + (spanEnd - oldT.Offset),
-					ColumnOffset: s.GetEnd().GetColumnOffset() + colDiff,
-					LineNumber:   s.GetEnd().GetLineNumber() + lineDiff,
-				},
-			}, true
-		}
+	// Find the diff span that contains the starting offset.
+	idx := sort.Search(len(p.spans), func(i int) bool {
+		return spanStart < p.spans[i].oldPrefix.Offset
+	}) - 1
+	if idx < 0 {
+		return nil, false
+	}
 
-		oldT.Update(d)
-		newT.Update(d)
+	d := p.spans[idx]
+	if d.Type != eq || spanEnd > d.oldPrefix.Offset+d.Length {
+		return nil, false
 	}
 
-	return nil, false
+	lineDiff := d.newPrefix.Lines - d.oldPrefix.Lines
+	colDiff := d.newPrefix.ColumnOffset - d.oldPrefix.ColumnOffset
+	if d.FirstNewline != -1 && spanStart-d.oldPrefix.Offset >= d.FirstNewline {
+		// The given span is past the first newline so it has no column diff.
+		colDiff = 0
+	}
+	return &cpb.Span{
+		Start: &cpb.Point{
+			ByteOffset:   d.newPrefix.Offset + (spanStart - d.oldPrefix.Offset),
+			ColumnOffset: s.GetStart().GetColumnOffset() + colDiff,
+			LineNumber:   s.GetStart().GetLineNumber() + lineDiff,
+		},
+		End: &cpb.Point{
+			ByteOffset:   d.newPrefix.Offset + (spanEnd - d.oldPrefix.Offset),
+			ColumnOffset: s.GetEnd().GetColumnOffset() + colDiff,
+			LineNumber:   s.GetEnd().GetLineNumber() + lineDiff,
+		},
+	}, true
 }
 
 // ByteOffsets returns the starting and ending byte offsets of the Span.


Test patch:
diff --git a/kythe/go/util/span/span_test.go b/kythe/go/util/span/span_test.go
index c62108348..33dcb26e7 100644
--- a/kythe/go/util/span/span_test.go
+++ b/kythe/go/util/span/span_test.go
@@ -21,6 +21,7 @@ import (
 	"strconv"
 	"testing"
 
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"kythe.io/kythe/go/test/testutil"
 	"kythe.io/kythe/go/util/compare"
 
@@ -303,10 +304,10 @@ func TestPatchSpan(t *testing.T) {
 func TestMarshal(t *testing.T) {
 	expected := &Patcher{
 		[]diff{
-			{2, ins, 1, 2, 5},
-			{10, del, 5, 0, 9},
-			{3, eq, 0, -1, -1},
-			{5, ins, 1, 2, 2},
+			{Length: 2, Type: ins, Newlines: 1, FirstNewline: 2, LastNewline: 5},
+			{Length: 10, Type: del, Newlines: 5, FirstNewline: 0, LastNewline: 9},
+			{Length: 3, Type: eq, Newlines: 0, FirstNewline: -1, LastNewline: -1},
+			{Length: 5, Type: ins, Newlines: 1, FirstNewline: 2, LastNewline: 2},
 		},
 	}
 
@@ -316,7 +317,23 @@ func TestMarshal(t *testing.T) {
 	found, err := Unmarshal(rec)
 	testutil.Fatalf(t, "Unmarshal: %v", err)
 
-	if diff := compare.ProtoDiff(expected.spans, found.spans); diff != "" {
+	// Check that the unmarshalled Patcher has correct prefix sums
+	ignoreTypeField := cmpopts.IgnoreFields(offsetTracker{}, "Type")
+	oldT := offsetTracker{Type: del}
+	newT := offsetTracker{Type: ins}
+	for _, d := range found.spans {
+		if diff := compare.ProtoDiff(oldT, d.oldPrefix, ignoreTypeField); diff != "" {
+			t.Fatalf("Old prefix sum incorrect (-expected; +found)\n%s", diff)
+		}
+		if diff := compare.ProtoDiff(newT, d.newPrefix, ignoreTypeField); diff != "" {
+			t.Fatalf("New prefix sum incorrect (-expected; +found)\n%s", diff)
+		}
+
+		oldT.Update(d)
+		newT.Update(d)
+	}
+
+	if diff := compare.ProtoDiff(expected.spans, found.spans, cmpopts.IgnoreUnexported(diff{})); diff != "" {
 		t.Errorf("(-expected; +found)\n%s", diff)
 	}
 }