From c05418dd2d8705f89c0c2772b2d46be0e4258a1f Mon Sep 17 00:00:00 2001
From: Bilal Akhtar <bilal@cockroachlabs.com>
Date: Wed, 14 Jun 2023 11:30:23 -0400
Subject: [PATCH] db: don't use internal types in ScanInternal

This change updates the ScanInternal function signature
to not use the internal keyspan.Key type. Instead we use
the new InternalRangeKey type alias that exports it.

Also update the pointCollapsingIter to just call SeekGE
in NextPrefix instead of panicking.

Necessary for cockroachdb/cockroach#103028.
---
 db.go                 |  2 +-
 scan_internal.go      | 10 +++++++---
 scan_internal_test.go |  4 ++--
 snapshot.go           |  4 +---
 4 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/db.go b/db.go
index 67d9b357eb..0605d4b8d0 100644
--- a/db.go
+++ b/db.go
@@ -1183,7 +1183,7 @@ func (d *DB) ScanInternal(
 	lower, upper []byte,
 	visitPointKey func(key *InternalKey, value LazyValue) error,
 	visitRangeDel func(start, end []byte, seqNum uint64) error,
-	visitRangeKey func(start, end []byte, keys []keyspan.Key) error,
+	visitRangeKey func(start, end []byte, keys []InternalRangeKey) error,
 	visitSharedFile func(sst *SharedSSTMeta) error,
 ) error {
 	iter := d.newInternalIter(nil /* snapshot */, &scanInternalOptions{
diff --git a/scan_internal.go b/scan_internal.go
index f7d5aeb946..429ff87665 100644
--- a/scan_internal.go
+++ b/scan_internal.go
@@ -591,10 +591,10 @@ func (p *pointCollapsingIterator) Next() (*base.InternalKey, base.LazyValue) {
 
 // NextPrefix implements the InternalIterator interface.
 func (p *pointCollapsingIterator) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
-	// TODO(bilal): Implement this. It'll be similar to SeekGE, except we'll call
+	// TODO(bilal): Implement this optimally. It'll be similar to SeekGE, except we'll call
 	// the child iterator's NextPrefix, and have some special logic in case pos
 	// is pcIterPosNext.
-	panic("unimplemented")
+	return p.SeekGE(succKey, base.SeekGEFlagsNone)
 }
 
 // Prev implements the InternalIterator interface.
@@ -846,13 +846,17 @@ func (d *DB) truncateSharedFile(
 	return sst, false, nil
 }
 
+// InternalRangeKey is a type alias that exports keyspan.Key, for use with
+// ScanInternal. Has a helper method to decode the Kind of the trailer.
+type InternalRangeKey = keyspan.Key
+
 func scanInternalImpl(
 	ctx context.Context,
 	lower, upper []byte,
 	iter *scanInternalIterator,
 	visitPointKey func(key *InternalKey, value LazyValue) error,
 	visitRangeDel func(start, end []byte, seqNum uint64) error,
-	visitRangeKey func(start, end []byte, keys []keyspan.Key) error,
+	visitRangeKey func(start, end []byte, keys []InternalRangeKey) error,
 	visitSharedFile func(sst *SharedSSTMeta) error,
 ) error {
 	if visitSharedFile != nil && (lower == nil || upper == nil) {
diff --git a/scan_internal_test.go b/scan_internal_test.go
index f442eab2ff..77a637c965 100644
--- a/scan_internal_test.go
+++ b/scan_internal_test.go
@@ -31,7 +31,7 @@ func TestScanInternal(t *testing.T) {
 			ctx context.Context,
 			lower, upper []byte, visitPointKey func(key *InternalKey, value LazyValue) error,
 			visitRangeDel func(start, end []byte, seqNum uint64) error,
-			visitRangeKey func(start, end []byte, keys []keyspan.Key) error,
+			visitRangeKey func(start, end []byte, keys []InternalRangeKey) error,
 			visitSharedFile func(sst *SharedSSTMeta) error) error
 	}
 	batches := map[string]*Batch{}
@@ -243,7 +243,7 @@ func TestScanInternal(t *testing.T) {
 			}, func(start, end []byte, seqNum uint64) error {
 				fmt.Fprintf(&b, "%s-%s#%d,RANGEDEL\n", start, end, seqNum)
 				return nil
-			}, func(start, end []byte, keys []keyspan.Key) error {
+			}, func(start, end []byte, keys []InternalRangeKey) error {
 				s := keyspan.Span{Start: start, End: end, Keys: keys}
 				fmt.Fprintf(&b, "%s\n", s.String())
 				return nil
diff --git a/snapshot.go b/snapshot.go
index d2b8564d71..4d018967b8 100644
--- a/snapshot.go
+++ b/snapshot.go
@@ -8,8 +8,6 @@ import (
 	"context"
 	"io"
 	"math"
-
-	"github.com/cockroachdb/pebble/internal/keyspan"
 )
 
 // Snapshot provides a read-only point-in-time view of the DB state.
@@ -68,7 +66,7 @@ func (s *Snapshot) ScanInternal(
 	lower, upper []byte,
 	visitPointKey func(key *InternalKey, value LazyValue) error,
 	visitRangeDel func(start, end []byte, seqNum uint64) error,
-	visitRangeKey func(start, end []byte, keys []keyspan.Key) error,
+	visitRangeKey func(start, end []byte, keys []InternalRangeKey) error,
 	visitSharedFile func(sst *SharedSSTMeta) error,
 ) error {
 	if s.db == nil {