kv: optimize IsPoint (pingcap#7722)

kv/key.go

@@ -78,7 +78,36 @@ type KeyRange struct {
 
 // IsPoint checks if the key range represents a point.
 func (r *KeyRange) IsPoint() bool {
-	return bytes.Equal(r.StartKey.PrefixNext(), r.EndKey)
+	if len(r.StartKey) != len(r.EndKey) {
+		// Works like
+		//   return bytes.Equal(r.StartKey.Next(), r.EndKey)
+
+		startLen := len(r.StartKey)
+		return startLen+1 == len(r.EndKey) &&
+			r.EndKey[startLen] == 0 &&
+			bytes.Equal(r.StartKey, r.EndKey[:startLen])
+	}
+	// Works like
+	//   return bytes.Equal(r.StartKey.PrefixNext(), r.EndKey)
+
+	i := len(r.StartKey) - 1
+	for ; i >= 0; i-- {
+		if r.StartKey[i] != 255 {
+			break
+		}
+		if r.EndKey[i] != 0 {
+			return false
+		}
+	}
+	if i < 0 {
+		// In case all bytes in StartKey are 255.
+		return false
+	}
+	// The byte at diffIdx in StartKey should be one less than the byte at diffIdx in EndKey.
+	// And bytes in StartKey and EndKey before diffIdx should be equal.
+	diffOneIdx := i
+	return r.StartKey[diffOneIdx]+1 == r.EndKey[diffOneIdx] &&
+		bytes.Equal(r.StartKey[:diffOneIdx], r.EndKey[:diffOneIdx])
 }
 
 // EncodedKey represents encoded key in low-level storage engine.

kv/key_test.go

@@ -15,6 +15,7 @@ package kv
 
 import (
 	"bytes"
+	"testing"
 	"time"
 
 	. "github.com/pingcap/check"
@@ -55,18 +56,63 @@ func (s *testKeySuite) TestPartialNext(c *C) {
 }
 
 func (s *testKeySuite) TestIsPoint(c *C) {
-	kr := KeyRange{
-		StartKey: Key("rowkey1"),
-		EndKey:   Key("rowkey2"),
+	tests := []struct {
+		start   []byte
+		end     []byte
+		isPoint bool
+	}{
+		{
+			start:   Key("rowkey1"),
+			end:     Key("rowkey2"),
+			isPoint: true,
+		},
+		{
+			start:   Key("rowkey1"),
+			end:     Key("rowkey3"),
+			isPoint: false,
+		},
+		{
+			start:   Key(""),
+			end:     []byte{0},
+			isPoint: true,
+		},
+		{
+			start:   []byte{123, 123, 255, 255},
+			end:     []byte{123, 124, 0, 0},
+			isPoint: true,
+		},
+		{
+			start:   []byte{123, 123, 255, 255},
+			end:     []byte{123, 124, 0, 1},
+			isPoint: false,
+		},
+		{
+			start:   []byte{123, 123},
+			end:     []byte{123, 123, 0},
+			isPoint: true,
+		},
+		{
+			start:   []byte{255},
+			end:     []byte{0},
+			isPoint: false,
+		},
 	}
-	c.Check(kr.IsPoint(), IsTrue)
-
-	kr.EndKey = Key("rowkey3")
-	c.Check(kr.IsPoint(), IsFalse)
+	for _, tt := range tests {
+		kr := KeyRange{
+			StartKey: tt.start,
+			EndKey:   tt.end,
+		}
+		c.Check(kr.IsPoint(), Equals, tt.isPoint)
+	}
+}
 
-	kr = KeyRange{
-		StartKey: Key(""),
-		EndKey:   Key([]byte{0}),
+func BenchmarkIsPoint(b *testing.B) {
+	b.ReportAllocs()
+	kr := KeyRange{
+		StartKey: []byte("rowkey1"),
+		EndKey:   []byte("rowkey2"),
+	}
+	for i := 0; i < b.N; i++ {
+		kr.IsPoint()
 	}
-	c.Check(kr.IsPoint(), IsTrue)
 }