From caaa8c0f83e4d294c96630ef30cf807f048fdd80 Mon Sep 17 00:00:00 2001
From: Erik Grinaker <grinaker@cockroachlabs.com>
Date: Tue, 15 Mar 2022 12:49:36 +0000
Subject: [PATCH] roachpb: add `Key.Prevish()` to find a previous key

This patch adds `Key.Prevish()`, which returns a previous key in
lexicographical sort order. This is needed to expand a latch span
leftwards to peek at any left-adjacent range keys.

It is impossible to find the exact immediate predecessor of a key,
because it can have an infinite number of `0xff` bytes at the end, so
this returns the nearest previous key right-padded with `0xff` up to the
given length. It is still possible for an infinite number of keys to
exist between `Key` and `Key.Prevish()`, as keys have unbounded length.

Release note: None
---
 pkg/roachpb/data.go           | 21 +++++++++++++++++++++
 pkg/roachpb/data_test.go      | 22 ++++++++++++++++++++++
 pkg/util/encoding/encoding.go | 26 ++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/pkg/roachpb/data.go b/pkg/roachpb/data.go
index 35f1c56de3d2..842f3b3a4476 100644
--- a/pkg/roachpb/data.go
+++ b/pkg/roachpb/data.go
@@ -50,6 +50,12 @@ const (
 	localPrefixByte = '\x01'
 	// LocalMaxByte is the end of the local key range.
 	LocalMaxByte = '\x02'
+	// PrevishKeyLength is a reasonable key length to use for Key.Prevish(),
+	// typically when peeking to the left of a known key. We want this to be as
+	// tight as possible, since it can e.g. be used for latch spans. However, the
+	// exact previous key has infinite length, so we assume that most keys are
+	// less than 1024 bytes, or have a fairly unique 1024-byte prefix.
+	PrevishKeyLength = 1024
 )
 
 var (
@@ -157,6 +163,21 @@ func (k Key) Next() Key {
 	return Key(encoding.BytesNext(k))
 }
 
+// Prevish returns a previous key in lexicographic sort order. It is impossible
+// in general to find the exact immediate predecessor key, because it has an
+// infinite number of 0xff bytes at the end, so this returns the nearest
+// previous key right-padded with 0xff up to length bytes. An infinite number of
+// keys may exist between Key and Key.Prevish(), as keys have unbounded length.
+// This also implies that k.Prevish().IsPrev(k) will often be false.
+//
+// PrevishKeyLength can be used as a reasonable length in most situations.
+//
+// The method may only take a shallow copy of the Key, so both the receiver and
+// the return value should be treated as immutable after.
+func (k Key) Prevish(length int) Key {
+	return Key(encoding.BytesPrevish(k, length))
+}
+
 // IsPrev is a more efficient version of k.Next().Equal(m).
 func (k Key) IsPrev(m Key) bool {
 	l := len(m) - 1
diff --git a/pkg/roachpb/data_test.go b/pkg/roachpb/data_test.go
index 2a3f51f10418..ab3a960d7f9f 100644
--- a/pkg/roachpb/data_test.go
+++ b/pkg/roachpb/data_test.go
@@ -12,6 +12,7 @@ package roachpb
 
 import (
 	"bytes"
+	"encoding/hex"
 	"math"
 	"math/rand"
 	"reflect"
@@ -221,6 +222,27 @@ func TestNextKey(t *testing.T) {
 	}
 }
 
+func TestPrevish(t *testing.T) {
+	const length = 4
+	testcases := []struct {
+		key    Key
+		expect Key
+	}{
+		{nil, nil},
+		{[]byte{}, []byte{}},
+		{[]byte{0x00}, []byte{}},
+		{[]byte{0x01, 0x00}, []byte{0x01}},
+		{[]byte{0x01}, []byte{0x00, 0xff, 0xff, 0xff}},
+		{[]byte{0x01, 0x01}, []byte{0x01, 0x00, 0xff, 0xff}},
+		{[]byte{0xff, 0xff, 0xff, 0xff}, []byte{0xff, 0xff, 0xff, 0xfe}},
+	}
+	for _, tc := range testcases {
+		t.Run(hex.EncodeToString(tc.key), func(t *testing.T) {
+			require.Equal(t, tc.expect, tc.key.Prevish(length))
+		})
+	}
+}
+
 func TestIsPrev(t *testing.T) {
 	for i, tc := range []struct {
 		k, m Key
diff --git a/pkg/util/encoding/encoding.go b/pkg/util/encoding/encoding.go
index 4a24429dc4dc..f3db0756247e 100644
--- a/pkg/util/encoding/encoding.go
+++ b/pkg/util/encoding/encoding.go
@@ -3215,3 +3215,29 @@ func BytesNext(b []byte) []byte {
 	bn[len(bn)-1] = 0
 	return bn
 }
+
+// BytesPrevish returns a previous byte slice in lexicographical ordering. It is
+// impossible in general to find the exact previous byte slice, because it has
+// an infinite number of 0xff bytes at the end, so this returns the nearest
+// previous slice right-padded with 0xff up to length bytes. It may reuse the
+// given slice when possible.
+func BytesPrevish(b []byte, length int) []byte {
+	bLen := len(b)
+	// An empty slice has no previous slice.
+	if bLen == 0 {
+		return b
+	}
+	// If the last byte is 0, just remove it.
+	if b[bLen-1] == 0 {
+		return b[:bLen-1]
+	}
+	// Otherwise, decrement the last byte and right-pad with 0xff.
+	if bLen > length {
+		length = bLen
+	}
+	buf := make([]byte, length)
+	copy(buf, b)
+	buf[bLen-1]--
+	copy(buf[bLen:], bytes.Repeat([]byte{0xff}, length-bLen))
+	return buf
+}