From b2773729afacb32b775460b63f18d25a99fa211d Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Sat, 10 Aug 2024 00:45:02 -0700
Subject: [PATCH 1/7] refactor iterate

---
 storage/pebble/operation/common.go      |  66 +++++++++++--
 storage/pebble/operation/common_test.go | 124 +++++++++++++++++++++++-
 2 files changed, 179 insertions(+), 11 deletions(-)

diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index 48e0861c418..34147dbb35a 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -157,6 +157,31 @@ func remove(key []byte) func(pebble.Writer) error {
 	}
 }
 
+func increment(end []byte) []byte {
+	if len(end) == 0 {
+		return nil
+	}
+
+	// Create a new slice with the same length as the original
+	newEnd := make([]byte, len(end))
+	copy(newEnd, end)
+
+	// Perform the increment operation
+	for i := len(newEnd) - 1; i >= 0; i-- {
+		newEnd[i]++
+		if newEnd[i] != 0 {
+			break
+		}
+	}
+
+	isOverflow := newEnd[0] == 0
+	if isOverflow {
+		return nil
+	}
+
+	return newEnd
+}
+
 // iterate iterates over a range of keys defined by a start and end key. The
 // start key may be higher than the end key, in which case we iterate in
 // reverse order.
@@ -175,17 +200,30 @@ func remove(key []byte) func(pebble.Writer) error {
 // provided handleFunc will be propagated back to the caller of iterate.
 func iterate(start []byte, end []byte, iteration iterationFunc, prefetchValues bool) func(pebble.Reader) error {
 	return func(tx pebble.Reader) error {
+
+		if len(start) == 0 {
+			return fmt.Errorf("start prefix is empty")
+		}
+
+		if len(end) == 0 {
+			return fmt.Errorf("end prefix is empty")
+		}
+
 		// Reverse iteration is not supported by pebble
 		if bytes.Compare(start, end) > 0 {
 			return fmt.Errorf("start key must be less than or equal to end key")
 		}
 
 		// initialize the default options and comparison modifier for iteration
-		maxEndWithPrefix := append(end, ffBytes...)
-
 		options := pebble.IterOptions{
 			LowerBound: start,
-			UpperBound: maxEndWithPrefix,
+			// LowerBound specifies the smallest key to iterate and it's inclusive.
+			// UpperBound specifies the largest key to iterate and it's exclusive (not inclusive)
+			// in order to match all keys prefixed with the `end` bytes, we increment the bytes of end by 1,
+			// for instance, to iterate keys between "hello" and "world",
+			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
+			// will all be included.
+			UpperBound: increment(end),
 		}
 
 		// In order to satisfy this function's prefix-wise inclusion semantics,
@@ -260,8 +298,6 @@ func iterate(start []byte, end []byte, iteration iterationFunc, prefetchValues b
 	}
 }
 
-var ffBytes = bytes.Repeat([]byte{0xFF}, 32)
-
 // traverse iterates over a range of keys defined by a prefix.
 //
 // The prefix must be shared by all keys in the iteration.
@@ -276,7 +312,13 @@ func traverse(prefix []byte, iteration iterationFunc) func(pebble.Reader) error
 
 		it, err := r.NewIter(&pebble.IterOptions{
 			LowerBound: prefix,
-			UpperBound: append(prefix, ffBytes...),
+			// LowerBound specifies the smallest key to iterate and it's inclusive.
+			// UpperBound specifies the largest key to iterate and it's exclusive (not inclusive)
+			// in order to match all keys prefixed with the `end` bytes, we increment the bytes of end by 1,
+			// for instance, to iterate keys between "hello" and "world",
+			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
+			// will all be included.
+			UpperBound: increment(prefix),
 		})
 
 		if err != nil {
@@ -373,9 +415,17 @@ func findHighestAtOrBelow(
 			return fmt.Errorf("prefix must not be empty")
 		}
 
-		key := append(prefix, b(height)...)
+		// why height+1? because:
+		// UpperBound specifies the largest key to iterate and it's exclusive (not inclusive)
+		// in order to match all keys indexed by height that is equal to or below the given height,
+		// we could increment the height by 1,
+		// for instance, to find highest key equal to or below 10, we first use 11 as the UpperBound, so that
+		// if there are 4 keys: [prefix-7, prefix-9, prefix-10, prefix-11], then all keys except
+		// prefix-11 will be included. And iterating them in the increasing order will find prefix-10
+		// as the highest key.
+		key := append(prefix, b(height+1)...)
 		it, err := r.NewIter(&pebble.IterOptions{
-			UpperBound: append(key, ffBytes...),
+			UpperBound: key,
 		})
 		if err != nil {
 			return fmt.Errorf("can not create iterator: %w", err)
diff --git a/storage/pebble/operation/common_test.go b/storage/pebble/operation/common_test.go
index 610c5788e8a..6a10996ba4b 100644
--- a/storage/pebble/operation/common_test.go
+++ b/storage/pebble/operation/common_test.go
@@ -452,10 +452,24 @@ func TestIterateBoundaries(t *testing.T) {
 		// before start -> not included in range
 		{0x09, 0xff},
 		// shares prefix with start -> included in range
+		{0x10},
 		{0x10, 0x00},
 		{0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
 		{0x10, 0xff},
-		{0x10, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
+		// prefix with a shared
+		{0x10,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+		},
+		{0x10,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff,
+		},
 		// prefix between start and end -> included in range
 		{0x11, 0x00},
 		{0x19, 0xff},
@@ -463,14 +477,118 @@ func TestIterateBoundaries(t *testing.T) {
 		{0x20, 0x00},
 		{0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
 		{0x20, 0xff},
-		{0x20, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
+		{0x20,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+		},
 		// after end -> not included in range
+		{0x21},
 		{0x21, 0x00},
 	}
 
 	// keys within the expected range
 	keysInRange := make([]string, 0)
-	for i := 1; i < 11; i++ {
+
+	firstNToExclude := 1
+	lastNToExclude := 2
+	for i := firstNToExclude; i < len(keys)-lastNToExclude; i++ {
+		key := keys[i]
+		keysInRange = append(keysInRange, hex.EncodeToString(key))
+	}
+
+	unittest.RunWithPebbleDB(t, func(db *pebble.DB) {
+
+		// insert the keys into the database
+		require.NoError(t, WithReaderBatchWriter(db, func(tx storage.PebbleReaderBatchWriter) error {
+			_, w := tx.ReaderWriter()
+			for _, key := range keys {
+				err := w.Set(key, []byte{0x00}, nil)
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		}))
+
+		// define iteration function that simply appends all traversed keys
+		found := make([]string, 0)
+
+		iteration := func() (checkFunc, createFunc, handleFunc) {
+			check := func(key []byte) bool {
+				found = append(found, hex.EncodeToString(key))
+				return false
+			}
+			create := func() interface{} {
+				return nil
+			}
+			handle := func() error {
+				return fmt.Errorf("shouldn't handle anything")
+			}
+			return check, create, handle
+		}
+
+		// iterate forward and check boundaries are included correctly
+		err := iterate(start, end, iteration, false)(db)
+		for i, f := range found {
+			t.Logf("forward %d: %x", i, f)
+		}
+		require.NoError(t, err, "should iterate forward without error")
+		assert.ElementsMatch(t, keysInRange, found, "forward iteration should go over correct keys")
+
+		// iterate backward and check boundaries are not supported
+	})
+}
+
+func TestIterateBoundariesOverflow(t *testing.T) {
+
+	// create range of keys covering all boundaries around our start/end values
+	start := []byte{0x10}
+	end := []byte{0xff}
+	keys := [][]byte{
+		// before start -> not included in range
+		{0x09, 0xff},
+		// shares prefix with start -> included in range
+		{0x10},
+		{0x10, 0x00},
+		{0x10, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+		{0x10, 0xff},
+		// prefix with a shared
+		{0x10,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+		},
+		{0x10,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff,
+		},
+		// prefix between start and end -> included in range
+		{0x11, 0x00},
+		{0x19, 0xff},
+		// shares prefix with end -> included in range
+		{0xff, 0x00},
+		{0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
+		{0xff, 0xff},
+		{0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+		},
+	}
+
+	// keys within the expected range
+	keysInRange := make([]string, 0)
+	firstNToExclude := 1
+	lastNToExclude := 0
+	for i := firstNToExclude; i < len(keys)-lastNToExclude; i++ {
 		key := keys[i]
 		keysInRange = append(keysInRange, hex.EncodeToString(key))
 	}

From d5d8b21adc69bb2789f6ad23b2c4b72272777435 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Wed, 14 Aug 2024 18:22:55 -0700
Subject: [PATCH 2/7] refactor using keyUpperBound

---
 storage/pebble/operation/common.go      | 35 +++++++++----------------
 storage/pebble/operation/common_test.go | 22 ++++++++++++++++
 2 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index 34147dbb35a..63aac3b17dc 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -157,29 +157,18 @@ func remove(key []byte) func(pebble.Writer) error {
 	}
 }
 
-func increment(end []byte) []byte {
-	if len(end) == 0 {
-		return nil
-	}
-
-	// Create a new slice with the same length as the original
-	newEnd := make([]byte, len(end))
-	copy(newEnd, end)
-
-	// Perform the increment operation
-	for i := len(newEnd) - 1; i >= 0; i-- {
-		newEnd[i]++
-		if newEnd[i] != 0 {
-			break
+// referred to https://pkg.go.dev/github.com/cockroachdb/pebble#example-Iterator-PrefixIteration
+func keyUpperBound(b []byte) []byte {
+	end := make([]byte, len(b))
+	copy(end, b)
+	for i := len(end) - 1; i >= 0; i-- {
+		// increment the bytes by 1
+		end[i] = end[i] + 1
+		if end[i] != 0 {
+			return end[:i+1]
 		}
 	}
-
-	isOverflow := newEnd[0] == 0
-	if isOverflow {
-		return nil
-	}
-
-	return newEnd
+	return nil // no upper-bound
 }
 
 // iterate iterates over a range of keys defined by a start and end key. The
@@ -223,7 +212,7 @@ func iterate(start []byte, end []byte, iteration iterationFunc, prefetchValues b
 			// for instance, to iterate keys between "hello" and "world",
 			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
 			// will all be included.
-			UpperBound: increment(end),
+			UpperBound: keyUpperBound(end),
 		}
 
 		// In order to satisfy this function's prefix-wise inclusion semantics,
@@ -318,7 +307,7 @@ func traverse(prefix []byte, iteration iterationFunc) func(pebble.Reader) error
 			// for instance, to iterate keys between "hello" and "world",
 			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
 			// will all be included.
-			UpperBound: increment(prefix),
+			UpperBound: keyUpperBound(prefix),
 		})
 
 		if err != nil {
diff --git a/storage/pebble/operation/common_test.go b/storage/pebble/operation/common_test.go
index 6a10996ba4b..bd5e8f81e84 100644
--- a/storage/pebble/operation/common_test.go
+++ b/storage/pebble/operation/common_test.go
@@ -463,6 +463,13 @@ func TestIterateBoundaries(t *testing.T) {
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 		},
+		{0x10,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0x00,
+		},
 		{0x10,
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
@@ -483,6 +490,20 @@ func TestIterateBoundaries(t *testing.T) {
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
 		},
+		{0x20,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0x00,
+		},
+		{0x20,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
+			0xff,
+		},
 		// after end -> not included in range
 		{0x21},
 		{0x21, 0x00},
@@ -541,6 +562,7 @@ func TestIterateBoundaries(t *testing.T) {
 	})
 }
 
+// to test the case where the end key is 0xffff..ff
 func TestIterateBoundariesOverflow(t *testing.T) {
 
 	// create range of keys covering all boundaries around our start/end values

From 26865bbb088e70574e36ab69cec39287a1ebafca Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 15 Aug 2024 16:36:47 -0700
Subject: [PATCH 3/7] address review comments

---
 storage/pebble/operation/cluster.go     | 2 +-
 storage/pebble/operation/common.go      | 8 ++------
 storage/pebble/operation/common_test.go | 6 +++---
 3 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/storage/pebble/operation/cluster.go b/storage/pebble/operation/cluster.go
index f99fc546770..dedb5bd8d8c 100644
--- a/storage/pebble/operation/cluster.go
+++ b/storage/pebble/operation/cluster.go
@@ -79,5 +79,5 @@ func LookupClusterBlocksByReferenceHeightRange(start, end uint64, clusterBlockID
 			return false
 		}
 		return check, nil, nil
-	}, false)
+	})
 }
diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index 63aac3b17dc..f72bbe966b9 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -187,7 +187,7 @@ func keyUpperBound(b []byte) []byte {
 // TODO: this function is unbounded – pass context.Context to this or calling functions to allow timing functions out.
 // No errors are expected during normal operation. Any errors returned by the
 // provided handleFunc will be propagated back to the caller of iterate.
-func iterate(start []byte, end []byte, iteration iterationFunc, prefetchValues bool) func(pebble.Reader) error {
+func iterate(start []byte, end []byte, iteration iterationFunc) func(pebble.Reader) error {
 	return func(tx pebble.Reader) error {
 
 		if len(start) == 0 {
@@ -256,11 +256,6 @@ func iterate(start []byte, end []byte, iteration iterationFunc, prefetchValues b
 				continue
 			}
 
-			// when prefetchValues is false, we skip loading the value
-			if !prefetchValues {
-				continue
-			}
-
 			binaryValue, err := it.ValueAndErr()
 			if err != nil {
 				return fmt.Errorf("failed to get value: %w", err)
@@ -414,6 +409,7 @@ func findHighestAtOrBelow(
 		// as the highest key.
 		key := append(prefix, b(height+1)...)
 		it, err := r.NewIter(&pebble.IterOptions{
+			LowerBound: prefix,
 			UpperBound: key,
 		})
 		if err != nil {
diff --git a/storage/pebble/operation/common_test.go b/storage/pebble/operation/common_test.go
index bd5e8f81e84..6ebf1d756d1 100644
--- a/storage/pebble/operation/common_test.go
+++ b/storage/pebble/operation/common_test.go
@@ -311,7 +311,7 @@ func TestIterate(t *testing.T) {
 			return check, create, handle
 		}
 
-		err := iterate(keys[0], keys[2], iterationFunc, true)(db)
+		err := iterate(keys[0], keys[2], iterationFunc)(db)
 		require.Nil(t, err)
 
 		assert.Equal(t, expected, actual)
@@ -551,7 +551,7 @@ func TestIterateBoundaries(t *testing.T) {
 		}
 
 		// iterate forward and check boundaries are included correctly
-		err := iterate(start, end, iteration, false)(db)
+		err := iterate(start, end, iteration)(db)
 		for i, f := range found {
 			t.Logf("forward %d: %x", i, f)
 		}
@@ -647,7 +647,7 @@ func TestIterateBoundariesOverflow(t *testing.T) {
 		}
 
 		// iterate forward and check boundaries are included correctly
-		err := iterate(start, end, iteration, false)(db)
+		err := iterate(start, end, iteration)(db)
 		for i, f := range found {
 			t.Logf("forward %d: %x", i, f)
 		}

From aad84b10d2d341a2a6a39fa26c292cefd281cddd Mon Sep 17 00:00:00 2001
From: Leo Zhang <zhangchiqing@gmail.com>
Date: Thu, 15 Aug 2024 16:37:37 -0700
Subject: [PATCH 4/7] Apply suggestions from code review

Co-authored-by: Jordan Schalm <jordan.schalm@flowfoundation.org>
---
 storage/pebble/operation/common.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index f72bbe966b9..521509cccb0 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -157,8 +157,12 @@ func remove(key []byte) func(pebble.Writer) error {
 	}
 }
 
+// prefixUpperBound returns a key K such that all possible keys beginning with the input prefix 
+// sort lower than K according to the byte-wise lexicographic key ordering used by Pebble.
+// This is used to define an upper bound for iteration, when we want to iterate over
+// all keys beginning with a given prefix. 
 // referred to https://pkg.go.dev/github.com/cockroachdb/pebble#example-Iterator-PrefixIteration
-func keyUpperBound(b []byte) []byte {
+func prefixUpperBound(prefix []byte) []byte {
 	end := make([]byte, len(b))
 	copy(end, b)
 	for i := len(end) - 1; i >= 0; i-- {
@@ -171,9 +175,9 @@ func keyUpperBound(b []byte) []byte {
 	return nil // no upper-bound
 }
 
-// iterate iterates over a range of keys defined by a start and end key. The
-// start key may be higher than the end key, in which case we iterate in
-// reverse order.
+// iterate iterates over a range of keys defined by a start and end key. 
+// The start key must be less than or equal to the end key by lexicographic ordering.
+// Both start and end keys must have non-zero length.
 //
 // The iteration range uses prefix-wise semantics. Specifically, all keys that
 // meet ANY of the following conditions are included in the iteration:

From c2ad1e9b893173ccf0a00cf95b4fdb9c543fea3e Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 15 Aug 2024 16:44:34 -0700
Subject: [PATCH 5/7] address review comments

---
 storage/pebble/operation/common.go | 37 ++++++------------------------
 1 file changed, 7 insertions(+), 30 deletions(-)

diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index 521509cccb0..a74592ec7f9 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -157,14 +157,14 @@ func remove(key []byte) func(pebble.Writer) error {
 	}
 }
 
-// prefixUpperBound returns a key K such that all possible keys beginning with the input prefix 
+// prefixUpperBound returns a key K such that all possible keys beginning with the input prefix
 // sort lower than K according to the byte-wise lexicographic key ordering used by Pebble.
 // This is used to define an upper bound for iteration, when we want to iterate over
-// all keys beginning with a given prefix. 
+// all keys beginning with a given prefix.
 // referred to https://pkg.go.dev/github.com/cockroachdb/pebble#example-Iterator-PrefixIteration
 func prefixUpperBound(prefix []byte) []byte {
-	end := make([]byte, len(b))
-	copy(end, b)
+	end := make([]byte, len(prefix))
+	copy(end, prefix)
 	for i := len(end) - 1; i >= 0; i-- {
 		// increment the bytes by 1
 		end[i] = end[i] + 1
@@ -175,7 +175,7 @@ func prefixUpperBound(prefix []byte) []byte {
 	return nil // no upper-bound
 }
 
-// iterate iterates over a range of keys defined by a start and end key. 
+// iterate iterates over a range of keys defined by a start and end key.
 // The start key must be less than or equal to the end key by lexicographic ordering.
 // Both start and end keys must have non-zero length.
 //
@@ -216,32 +216,9 @@ func iterate(start []byte, end []byte, iteration iterationFunc) func(pebble.Read
 			// for instance, to iterate keys between "hello" and "world",
 			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
 			// will all be included.
-			UpperBound: keyUpperBound(end),
+			UpperBound: prefixUpperBound(end),
 		}
 
-		// In order to satisfy this function's prefix-wise inclusion semantics,
-		// we append 0xff bytes to the largest of start and end.
-		// This ensures Badger will seek to the largest key with that prefix
-		// for reverse iteration, thus including all keys with a prefix matching
-		// the starting key. It also enables us to detect boundary conditions by
-		// simple lexicographic comparison (ie. bytes.Compare) rather than
-		// explicitly comparing prefixes.
-
-		// If start is bigger than end, we have a backwards iteration:
-		// 1) We set the reverse option on the iterator, so we step through all
-		//    the keys backwards. This modifies the behaviour of Seek to go to
-		//    the first key that is less than or equal to the start key (as
-		//    opposed to greater than or equal in a regular iteration).
-		// 2) In order to satisfy this function's prefix-wise inclusion semantics,
-		//    we append a 0xff-byte suffix to the start key so the seek will go
-		// to the right place.
-		// 3) For a regular iteration, we break the loop upon hitting the first
-		//    item that has a key higher than the end prefix. In order to reverse
-		//    this, we use a modifier for the comparison that reverses the check
-		//    and makes it stop upon the first item lower than the end prefix.
-		// for forward iteration, add the 0xff-bytes suffix to the end
-		// prefix, to ensure we include all keys with that prefix before
-		// finishing.
 		it, err := tx.NewIter(&options)
 		if err != nil {
 			return fmt.Errorf("can not create iterator: %w", err)
@@ -306,7 +283,7 @@ func traverse(prefix []byte, iteration iterationFunc) func(pebble.Reader) error
 			// for instance, to iterate keys between "hello" and "world",
 			// we use "hello" as LowerBound, "worle" as UpperBound, so that "world", "world1", "worldffff...ffff"
 			// will all be included.
-			UpperBound: keyUpperBound(prefix),
+			UpperBound: prefixUpperBound(prefix),
 		})
 
 		if err != nil {

From d2c101677d4f8df08d0a033f07b649d51ab8d53a Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 15 Aug 2024 16:58:29 -0700
Subject: [PATCH 6/7] refactor getStartEndKeys with prefixUpperBound

---
 storage/pebble/operation/common.go      | 24 +-----------------------
 storage/pebble/operation/common_test.go | 20 ++++++++------------
 2 files changed, 9 insertions(+), 35 deletions(-)

diff --git a/storage/pebble/operation/common.go b/storage/pebble/operation/common.go
index a74592ec7f9..8124f7431f6 100644
--- a/storage/pebble/operation/common.go
+++ b/storage/pebble/operation/common.go
@@ -336,32 +336,10 @@ func traverse(prefix []byte, iteration iterationFunc) func(pebble.Reader) error
 // No errors are expected during normal operation.
 func removeByPrefix(prefix []byte) func(pebble.Writer) error {
 	return func(tx pebble.Writer) error {
-		start, end := getStartEndKeys(prefix)
-		return tx.DeleteRange(start, end, nil)
+		return tx.DeleteRange(prefix, prefixUpperBound(prefix), nil)
 	}
 }
 
-// getStartEndKeys calculates the start and end keys for a given prefix.
-// note: DeleteRange takes [start, end) (inclusion on start, exclusive on end),
-func getStartEndKeys(prefix []byte) (start, end []byte) {
-	start = prefix
-
-	// End key is the prefix with the last byte incremented by 1
-	end = append([]byte{}, prefix...)
-	for i := len(end) - 1; i >= 0; i-- {
-		if end[i] < 0xff {
-			end[i]++
-			break
-		}
-		end[i] = 0
-		if i == 0 {
-			end = append(end, 0)
-		}
-	}
-
-	return start, end
-}
-
 func convertNotFoundError(err error) error {
 	if errors.Is(err, pebble.ErrNotFound) {
 		return storage.ErrNotFound
diff --git a/storage/pebble/operation/common_test.go b/storage/pebble/operation/common_test.go
index 6ebf1d756d1..07050dbcae6 100644
--- a/storage/pebble/operation/common_test.go
+++ b/storage/pebble/operation/common_test.go
@@ -18,24 +18,20 @@ import (
 var upsert = insert
 var update = insert
 
-func TestGetStartEndKeys(t *testing.T) {
+func TestPrefixUpperBound(t *testing.T) {
 	tests := []struct {
-		prefix        []byte
-		expectedStart []byte
-		expectedEnd   []byte
+		prefix      []byte
+		expectedEnd []byte
 	}{
-		{[]byte("a"), []byte("a"), []byte("b")},
-		{[]byte("abc"), []byte("abc"), []byte("abd")},
-		{[]byte("prefix"), []byte("prefix"), []byte("prefiy")},
+		{[]byte("a"), []byte("b")},
+		{[]byte("abc"), []byte("abd")},
+		{[]byte("prefix"), []byte("prefiy")},
 	}
 
 	for _, test := range tests {
-		start, end := getStartEndKeys(test.prefix)
-		if !bytes.Equal(start, test.expectedStart) {
-			t.Errorf("getStartEndKeys(%q) start = %q; want %q", test.prefix, start, test.expectedStart)
-		}
+		end := prefixUpperBound(test.prefix)
 		if !bytes.Equal(end, test.expectedEnd) {
-			t.Errorf("getStartEndKeys(%q) end = %q; want %q", test.prefix, end, test.expectedEnd)
+			t.Errorf("prefixUpperBound(%q) end = %q; want %q", test.prefix, end, test.expectedEnd)
 		}
 	}
 }

From 2ea56984c26d81b20ddf0ffbaf201795d63b36b6 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 16 Aug 2024 10:47:12 -0700
Subject: [PATCH 7/7] add test case to events

---
 storage/pebble/operation/events_test.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/storage/pebble/operation/events_test.go b/storage/pebble/operation/events_test.go
index 348ae9ec663..24d4fccd142 100644
--- a/storage/pebble/operation/events_test.go
+++ b/storage/pebble/operation/events_test.go
@@ -19,7 +19,9 @@ func TestRetrieveEventByBlockIDTxID(t *testing.T) {
 
 		// create block ids, transaction ids and event types slices
 		blockIDs := []flow.Identifier{flow.HashToID([]byte{0x01}), flow.HashToID([]byte{0x02})}
-		txIDs := []flow.Identifier{flow.HashToID([]byte{0x11}), flow.HashToID([]byte{0x12})}
+		txIDs := []flow.Identifier{flow.HashToID([]byte{0x11}), flow.HashToID([]byte{0x12}),
+			[32]byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF},
+		}
 		eTypes := []flow.EventType{flow.EventAccountCreated, flow.EventAccountUpdated}
 
 		// create map of block id to event, tx id to event and event type to event