spanset: Add Overlaps() and PartiallyOverlaps helper functions

Analogous to how we have a spanset helper contains() that
understands the special span representation: [x-eps,x).
This commit adds a couple helpers to detect overlapping spans.

Moreover, this commit adds tests to verify the behaviour of
these helper functions.

Author: iskettaneh

pkg/kv/kvserver/spanset/spanset.go

@@ -76,6 +76,11 @@ type Span struct {
 	Timestamp hlc.Timestamp
 }
 
+// TrickySpan represents a span that supports a special encoding where a nil
+// start key with a non-nil end key represents the point span:
+// [EndKey.Prev(), EndKey).
+type TrickySpan roachpb.Span
+
 // SpanSet tracks the set of key spans touched by a command, broken into MVCC
 // and non-MVCC accesses. The set is divided into subsets for access type
 // (read-only or read/write) and key scope (local or global; used to facilitate
@@ -88,7 +93,7 @@ type SpanSet struct {
 	// shouldn't be accessed (forbidden). This allows for complex pattern matching
 	// like forbidding specific keys across all range IDs without enumerating them
 	// explicitly.
-	forbiddenSpansMatchers []func(roachpb.Span) error
+	forbiddenSpansMatchers []func(TrickySpan) error
 	allowUndeclared        bool
 	allowForbidden         bool
 }
@@ -213,7 +218,7 @@ func (s *SpanSet) AddMVCC(access SpanAccess, span roachpb.Span, timestamp hlc.Ti
 
 // AddForbiddenMatcher adds a forbidden span matcher. The matcher is a function
 // that is called for each span access to check if it should be forbidden.
-func (s *SpanSet) AddForbiddenMatcher(matcher func(roachpb.Span) error) {
+func (s *SpanSet) AddForbiddenMatcher(matcher func(TrickySpan) error) {
 	s.forbiddenSpansMatchers = append(s.forbiddenSpansMatchers, matcher)
 }
 
@@ -353,7 +358,7 @@ func (s *SpanSet) checkAllowed(
 	if !s.allowForbidden {
 		// Check if the span is forbidden.
 		for _, matcher := range s.forbiddenSpansMatchers {
-			if err := matcher(span); err != nil {
+			if err := matcher(TrickySpan(span)); err != nil {
 				return errors.Errorf("cannot %s span %s: matches forbidden pattern",
 					access, span)
 			}
@@ -370,7 +375,7 @@ func (s *SpanSet) checkAllowed(
 
 		for ac := access; ac < NumSpanAccess; ac++ {
 			for _, cur := range s.spans[ac][scope] {
-				if contains(cur.Span, span) && check(ac, cur) {
+				if Contains(cur.Span, TrickySpan(span)) && check(ac, cur) {
 					return nil
 				}
 			}
@@ -382,24 +387,63 @@ func (s *SpanSet) checkAllowed(
 	return nil
 }
 
-// contains returns whether s1 contains s2. Unlike Span.Contains, this function
-// supports spans with a nil start key and a non-nil end key (e.g. "[nil, c)").
-// In this form, s2.Key (inclusive) is considered to be the previous key to
-// s2.EndKey (exclusive).
-func contains(s1, s2 roachpb.Span) bool {
-	if s2.Key != nil {
-		// The common case.
-		return s1.Contains(s2)
+// Contains returns whether s1 contains s2, where s2 can be a TrickySpan.
+func Contains(s1 roachpb.Span, s2 TrickySpan) bool {
+	s2Span := roachpb.Span(s2)
+
+	if s2Span.Key != nil {
+		// The common case: s2 is a regular span with a non-nil start key.
+		return s1.Contains(s2Span)
 	}
 
+	// s2 is a TrickySpan with nil Key and non-nil EndKey.
 	// The following is equivalent to:
 	//   s1.Contains(roachpb.Span{Key: s2.EndKey.Prev()})
 
 	if s1.EndKey == nil {
-		return s1.Key.IsPrev(s2.EndKey)
+		return s1.Key.IsPrev(s2Span.EndKey)
+	}
+
+	return s1.Key.Compare(s2Span.EndKey) < 0 && s1.EndKey.Compare(s2Span.EndKey) >= 0
+}
+
+// Overlaps returns whether s1 overlaps s2, where s2 can be a TrickySpan.
+func Overlaps(s1 roachpb.Span, s2 TrickySpan) bool {
+	s2Span := roachpb.Span(s2)
+
+	// The common case: both spans have non-nil start keys.
+	if s2Span.Key != nil {
+		return s1.Overlaps(s2Span)
+	}
+
+	// s2 is a TrickySpan with nil Key and non-nil EndKey.
+	// The following is equivalent to:
+	//   s1.Overlaps(roachpb.Span{Key: s2.EndKey.Prev()})
+
+	if s1.EndKey == nil {
+		// s1 is a point span, overlaps with s2 iff s1.Key is the prev of s2.EndKey
+		return s1.Key.IsPrev(s2Span.EndKey)
 	}
 
-	return s1.Key.Compare(s2.EndKey) < 0 && s1.EndKey.Compare(s2.EndKey) >= 0
+	// s1 is [s1.Key, s1.EndKey), s2 is [s2.EndKey.Prev(), s2.EndKey)
+	// They overlap iff s2.EndKey.Prev() is in [s1.Key, s1.EndKey).
+	return s1.Key.Compare(s2Span.EndKey) < 0 && s1.EndKey.Compare(s2Span.EndKey) >= 0
+}
+
+// PartiallyOverlaps returns whether s1 partially overlaps s2,
+// where s2 can be a TrickySpan. If true, it means that s1 and s2 overlap, but
+// neither s1 nor s2 contain one another.
+func PartiallyOverlaps(s1 roachpb.Span, s2 TrickySpan) bool {
+	s2Span := roachpb.Span(s2)
+
+	// The common case: both spans have non-nil start keys.
+	if s2Span.Key != nil {
+		return s1.Overlaps(s2Span) && !s1.Contains(s2Span) && !s2Span.Contains(s1)
+	}
+
+	// At this point, we know that s2 is a TrickySpan, and it represents a special
+	// point span. Point spans can never partially overlap with any span.
+	return false
 }
 
 // Validate returns an error if any spans that have been added to the set

pkg/kv/kvserver/spanset/spanset_test.go

@@ -7,6 +7,7 @@ package spanset
 
 import (
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -394,3 +395,162 @@ func TestSpanSetWriteImpliesRead(t *testing.T) {
 		t.Errorf("expected to be allowed to read rwSpan, error: %+v", err)
 	}
 }
+
+// makeSpanHelper accepts strings like: "a-d", and returns a span with
+// startKey = a, and endKey = d. It also accepts `X` which represents nil. For
+// example, "X-d" returns a span with a nil startKey, and endKey = d.
+func makeSpanHelper(t *testing.T, s string) roachpb.Span {
+	parts := strings.Split(s, "-")
+	require.Len(t, parts, 2)
+
+	var start roachpb.Key
+	var end roachpb.Key
+
+	if parts[0] != "X" {
+		start = roachpb.Key(parts[0])
+	}
+
+	if parts[1] != "X" {
+		end = roachpb.Key(parts[1])
+	}
+
+	return roachpb.Span{
+		Key:    start,
+		EndKey: end,
+	}
+}
+
+// Test that Contains correctly determines if s1 contains s2, including
+// support for spans with nil start/end keys.
+func TestContains(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		name     string
+		s1       roachpb.Span
+		s2       roachpb.Span
+		expected bool
+	}{
+		// s1 equals s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "a-c"), expected: true},
+		// s1 contains s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "b-c"), expected: true},
+		// s1 contains point s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "a-X"), expected: true},
+		// Point s1 contains point s2.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "a-X"), expected: true},
+		// s1 contains point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "X-d"), expected: true},
+		// Point s1 contains point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-X"), s2: roachpb.Span{EndKey: roachpb.Key("a").Next()}, expected: true},
+		// s1 does not contain s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "d-f"), expected: false},
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "b-d"), expected: false},
+		{s1: makeSpanHelper(t, "b-c"), s2: makeSpanHelper(t, "a-d"), expected: false},
+		// s1 does not contain s2 point.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "d-X"), expected: false},
+		// Point s1 does not contain s2.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "b-d"), expected: false},
+		// Point s1 does not contain point s2.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "b-X"), expected: false},
+		// s1 does not contain point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "X-e"), expected: false},
+		// Point s1 does not contain point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "X-b"), expected: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, Contains(tc.s1, TrickySpan(tc.s2)))
+		})
+	}
+}
+
+// Test that Overlaps correctly determines if s1 overlaps s2, including
+// support for spans with nil start/end keys.
+func TestOverlaps(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		name     string
+		s1       roachpb.Span
+		s2       roachpb.Span
+		expected bool
+	}{
+		// s1 equals s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "a-c"), expected: true},
+		// s1 overlaps s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "b-d"), expected: true},
+		// s1 contains s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "b-c"), expected: true},
+		// s2 contains s1.
+		{s1: makeSpanHelper(t, "b-c"), s2: makeSpanHelper(t, "a-d"), expected: true},
+		// s1 overlaps point s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "a-X"), expected: true},
+		// s1 overlaps point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "X-d"), expected: true},
+		// Point s1 overlaps point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-X"), s2: roachpb.Span{EndKey: roachpb.Key("a").Next()}, expected: true},
+		// Point s1 overlaps point s2.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "a-X"), expected: true},
+		// s1 doesn't overlap s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "d-f"), expected: false},
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "c-d"), expected: false},
+		// s1 doesn't overlap point s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "c-X"), expected: false},
+		// s1 doesn't overlap point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "X-a"), expected: false},
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "X-d"), expected: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			require.Equal(t, tc.expected, Overlaps(tc.s1, TrickySpan(tc.s2)))
+		})
+	}
+}
+
+// Test that PartiallyOverlaps correctly determines if s1 partially overlaps s2,
+// including support for spans with nil start/end keys.
+func TestPartiallyOverlaps(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	testCases := []struct {
+		name     string
+		s1       roachpb.Span
+		s2       roachpb.Span
+		expected bool
+	}{
+		// s1 equals s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "a-c"), expected: false},
+		// s1 partially overlaps s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "b-d"), expected: true},
+		{s1: makeSpanHelper(t, "b-d"), s2: makeSpanHelper(t, "a-c"), expected: true},
+		// s1 contains s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "b-c"), expected: false},
+		// s2 contains s1.
+		{s1: makeSpanHelper(t, "b-c"), s2: makeSpanHelper(t, "a-d"), expected: false},
+		// s1 contains point s2.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "a-X"), expected: false},
+		// s1 contains point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-d"), s2: makeSpanHelper(t, "X-d"), expected: false},
+		// Point s1 contains point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-X"), s2: roachpb.Span{EndKey: roachpb.Key("a").Next()}, expected: false},
+		// Point s1 contains point s2.
+		{s1: makeSpanHelper(t, "a-X"), s2: makeSpanHelper(t, "a-X"), expected: false},
+		// s1 doesn't overlap s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "d-f"), expected: false},
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "c-d"), expected: false},
+		// s1 doesn't overlap point s2.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "c-X"), expected: false},
+		// s1 doesn't overlap point s2 with nil startKey.
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "X-a"), expected: false},
+		{s1: makeSpanHelper(t, "a-c"), s2: makeSpanHelper(t, "X-d"), expected: false},
+	}
+
+	for _, tc := range testCases {
+		t.Run("", func(t *testing.T) {
+			require.Equal(t, tc.expected, PartiallyOverlaps(tc.s1, TrickySpan(tc.s2)))
+		})
+	}
+}