cosmos · tac0turtle · Nov 7, 2024 · Oct 4, 2024 · Oct 4, 2024 · Oct 4, 2024
@@ -5,6 +5,7 @@ go 1.23
 // server v2 integration
 replace (
 	cosmossdk.io/api => ../../api
+	cosmossdk.io/core/testing => ../../core/testing
 	cosmossdk.io/server/v2/appmanager => ../../server/v2/appmanager
 	cosmossdk.io/server/v2/stf => ../../server/v2/stf
 	cosmossdk.io/store/v2 => ../../store/v2
@@ -30,7 +31,7 @@ require (
 require (
 	buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.35.1-20240701160653-fedbb9acfd2f.1 // indirect
 	buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88ef6483f90f.1 // indirect
-	cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29 // indirect
+	cosmossdk.io/core/testing v0.0.0 // indirect
 	cosmossdk.io/errors/v2 v2.0.0-20240731132947-df72853b3ca5 // indirect
 	github.com/DataDog/zstd v1.5.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect

@@ -4,8 +4,6 @@ buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88e
 buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88ef6483f90f.1/go.mod h1:zqi/LZjZhyvjCMTEVIwAf5VRlkLduuCfqmZxgoormq0=
 cosmossdk.io/core v1.0.0-alpha.5 h1:McjYXAQ6XcT20v2uHyH7PhoWH8V+mebzfVFqT3GinsI=
 cosmossdk.io/core v1.0.0-alpha.5/go.mod h1:3u9cWq1FAVtiiCrDPpo4LhR+9V6k/ycSG4/Y/tREWCY=
-cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29 h1:NxxUo0GMJUbIuVg0R70e3cbn9eFTEuMr7ev1AFvypdY=
-cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29/go.mod h1:8s2tPeJtSiQuoyPmr2Ag7meikonISO4Fv4MoO8+ORrs=
 cosmossdk.io/depinject v1.1.0 h1:wLan7LG35VM7Yo6ov0jId3RHWCGRhe8E8bsuARorl5E=
 cosmossdk.io/depinject v1.1.0/go.mod h1:kkI5H9jCGHeKeYWXTqYdruogYrEeWvBQCw1Pj4/eCFI=
 cosmossdk.io/errors/v2 v2.0.0-20240731132947-df72853b3ca5 h1:IQNdY2kB+k+1OM2DvqFG1+UgeU1JzZrWtwuWzI3ZfwA=

@@ -4,6 +4,7 @@ go 1.23.1
 
 replace (
 	cosmossdk.io/api => ../../../api
+	cosmossdk.io/core/testing => ../../../core/testing
 	cosmossdk.io/server/v2 => ../
 	cosmossdk.io/server/v2/appmanager => ../appmanager
 	cosmossdk.io/server/v2/stf => ../stf
@@ -43,7 +44,7 @@ require (
 require (
 	buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.35.1-20240701160653-fedbb9acfd2f.1 // indirect
 	buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88ef6483f90f.1 // indirect
-	cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29 // indirect
+	cosmossdk.io/core/testing v0.0.0 // indirect
 	cosmossdk.io/depinject v1.1.0 // indirect
 	cosmossdk.io/errors v1.0.1 // indirect
 	cosmossdk.io/math v1.3.0 // indirect

@@ -8,8 +8,6 @@ cosmossdk.io/collections v0.4.0 h1:PFmwj2W8szgpD5nOd8GWH6AbYNi1f2J6akWXJ7P5t9s=
 cosmossdk.io/collections v0.4.0/go.mod h1:oa5lUING2dP+gdDquow+QjlF45eL1t4TJDypgGd+tv0=
 cosmossdk.io/core v1.0.0-alpha.5 h1:McjYXAQ6XcT20v2uHyH7PhoWH8V+mebzfVFqT3GinsI=
 cosmossdk.io/core v1.0.0-alpha.5/go.mod h1:3u9cWq1FAVtiiCrDPpo4LhR+9V6k/ycSG4/Y/tREWCY=
-cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29 h1:NxxUo0GMJUbIuVg0R70e3cbn9eFTEuMr7ev1AFvypdY=
-cosmossdk.io/core/testing v0.0.0-20240923163230-04da382a9f29/go.mod h1:8s2tPeJtSiQuoyPmr2Ag7meikonISO4Fv4MoO8+ORrs=
 cosmossdk.io/depinject v1.1.0 h1:wLan7LG35VM7Yo6ov0jId3RHWCGRhe8E8bsuARorl5E=
 cosmossdk.io/depinject v1.1.0/go.mod h1:kkI5H9jCGHeKeYWXTqYdruogYrEeWvBQCw1Pj4/eCFI=
 cosmossdk.io/errors v1.0.1 h1:bzu+Kcr0kS/1DuPBtUFdWjzLqyUuCiyHjyJB6srBV/0=

@@ -0,0 +1,83 @@
+package branch
+
+import (
+	"fmt"
+	"testing"
+
+	"cosmossdk.io/core/store"
+	coretesting "cosmossdk.io/core/testing"
+)
+
+var (
+	stackSizes   = []int{1, 10, 100}
+	elemsInStack = 10
+)
-var (
-	stackSizes   = []int{1, 10, 100}
-	elemsInStack = 10
-)
+var (
+	stackSizes   = []int{100, 1000, 10000}
+	elemsInStack = 10
+)
-var (
-	stackSizes   = []int{1, 10, 100}
-	elemsInStack = 10
-)
+var (
+	stackSizes   = []int{100, 1000, 10000}
+	elemsInStack = 10
+)
+
+func Benchmark_CacheStack_Set(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				_ = bs.Set([]byte{0}, []byte{0})
+			}
+		})
+	}
+}
-func Benchmark_CacheStack_Set(b *testing.B) {
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				_ = bs.Set([]byte{0}, []byte{0})
-			}
-		})
-	}
-}
+func Benchmark_CacheStack_Set(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			key := make([]byte, 32)   // typical key size
+			value := make([]byte, 64) // typical value size
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				binary.BigEndian.PutUint64(key, uint64(i))
+				if err := bs.Set(key, value); err != nil {
+					b.Fatal(err)
+				}
+				b.StopTimer()
+				bs.Discard()  // cleanup between iterations
+				b.StartTimer()
+			}
+		})
+	}
+}
-func Benchmark_CacheStack_Set(b *testing.B) {
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				_ = bs.Set([]byte{0}, []byte{0})
-			}
-		})
-	}
-}
+func Benchmark_CacheStack_Set(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			key := make([]byte, 32)   // typical key size
+			value := make([]byte, 64) // typical value size
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				binary.BigEndian.PutUint64(key, uint64(i))
+				if err := bs.Set(key, value); err != nil {
+					b.Fatal(err)
+				}
+				b.StopTimer()
+				bs.Discard()  // cleanup between iterations
+				b.StartTimer()
+			}
+		})
+	}
+}
+
+func Benchmark_Get(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				_, _ = bs.Get([]byte{0})
+			}
+		})
+	}
+}
-func Benchmark_Get(b *testing.B) {
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				_, _ = bs.Get([]byte{0})
-			}
-		})
-	}
-}
+func Benchmark_Get(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			key := []byte("test-key")
+			expectedValue := []byte("test-value")
+			if err := bs.Set(key, expectedValue); err != nil {
+				b.Fatal(err)
+			}
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				value, err := bs.Get(key)
+				if err != nil {
+					b.Fatal(err)
+				}
+				if !bytes.Equal(value, expectedValue) {
+					b.Fatalf("unexpected value: got %x, want %x", value, expectedValue)
+				}
+			}
+		})
+	}
+}
-func Benchmark_Get(b *testing.B) {
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				_, _ = bs.Get([]byte{0})
-			}
-		})
-	}
-}
+func Benchmark_Get(b *testing.B) {
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			key := []byte("test-key")
+			expectedValue := []byte("test-value")
+			if err := bs.Set(key, expectedValue); err != nil {
+				b.Fatal(err)
+			}
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				value, err := bs.Get(key)
+				if err != nil {
+					b.Fatal(err)
+				}
+				if !bytes.Equal(value, expectedValue) {
+					b.Fatalf("unexpected value: got %x, want %x", value, expectedValue)
+				}
+			}
+		})
+	}
+}
+
+func Benchmark_Iterate(b *testing.B) {
+	var keySink, valueSink any
+
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				iter, _ := bs.Iterator(nil, nil)
+				for iter.Valid() {
+					keySink = iter.Key()
+					valueSink = iter.Value()
+					iter.Next()
+				}
+				_ = iter.Close()
+			}
+		})
+	}
+
+	_ = keySink
+	_ = valueSink
+}
-func Benchmark_Iterate(b *testing.B) {
-	var keySink, valueSink any
-
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				iter, _ := bs.Iterator(nil, nil)
-				for iter.Valid() {
-					keySink = iter.Key()
-					valueSink = iter.Value()
-					iter.Next()
-				}
-				_ = iter.Close()
-			}
-		})
-	}
-
-	_ = keySink
-	_ = valueSink
-}
+func Benchmark_Iterate(b *testing.B) {
+	var (
+		expectedKeys   = make([][]byte, 0, elemsInStack)
+		expectedValues = make([][]byte, 0, elemsInStack)
+	)
+
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			// Prepare expected data
+			for j := 0; j < elemsInStack; j++ {
+				expectedKeys = append(expectedKeys, []byte{byte(stackSize - 1), byte(j)})
+				expectedValues = append(expectedValues, []byte{byte(j)})
+			}
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				iter, err := bs.Iterator([]byte{byte(stackSize - 1)}, nil)
+				if err != nil {
+					b.Fatal(err)
+				}
+				keyIdx := 0
+				for iter.Valid() {
+					if !bytes.Equal(iter.Key(), expectedKeys[keyIdx]) {
+						b.Fatalf("unexpected key at %d: got %x, want %x", keyIdx, iter.Key(), expectedKeys[keyIdx])
+					}
+					if !bytes.Equal(iter.Value(), expectedValues[keyIdx]) {
+						b.Fatalf("unexpected value at %d: got %x, want %x", keyIdx, iter.Value(), expectedValues[keyIdx])
+					}
+					keyIdx++
+					iter.Next()
+				}
+				if err := iter.Close(); err != nil {
+					b.Fatal(err)
+				}
+				if keyIdx != len(expectedKeys) {
+					b.Fatalf("unexpected number of keys: got %d, want %d", keyIdx, len(expectedKeys))
+				}
+			}
+		})
+	}
+}
-func Benchmark_Iterate(b *testing.B) {
-	var keySink, valueSink any
-
-	for _, stackSize := range stackSizes {
-		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
-			bs := makeBranchStack(b, stackSize)
-			b.ResetTimer()
-			b.ReportAllocs()
-			for i := 0; i < b.N; i++ {
-				iter, _ := bs.Iterator(nil, nil)
-				for iter.Valid() {
-					keySink = iter.Key()
-					valueSink = iter.Value()
-					iter.Next()
-				}
-				_ = iter.Close()
-			}
-		})
-	}
-
-	_ = keySink
-	_ = valueSink
-}
+func Benchmark_Iterate(b *testing.B) {
+	var (
+		expectedKeys   = make([][]byte, 0, elemsInStack)
+		expectedValues = make([][]byte, 0, elemsInStack)
+	)
+
+	for _, stackSize := range stackSizes {
+		b.Run(fmt.Sprintf("StackSize%d", stackSize), func(b *testing.B) {
+			bs := makeBranchStack(b, stackSize)
+			// Prepare expected data
+			for j := 0; j < elemsInStack; j++ {
+				expectedKeys = append(expectedKeys, []byte{byte(stackSize - 1), byte(j)})
+				expectedValues = append(expectedValues, []byte{byte(j)})
+			}
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				iter, err := bs.Iterator([]byte{byte(stackSize - 1)}, nil)
+				if err != nil {
+					b.Fatal(err)
+				}
+				keyIdx := 0
+				for iter.Valid() {
+					if !bytes.Equal(iter.Key(), expectedKeys[keyIdx]) {
+						b.Fatalf("unexpected key at %d: got %x, want %x", keyIdx, iter.Key(), expectedKeys[keyIdx])
+					}
+					if !bytes.Equal(iter.Value(), expectedValues[keyIdx]) {
+						b.Fatalf("unexpected value at %d: got %x, want %x", keyIdx, iter.Value(), expectedValues[keyIdx])
+					}
+					keyIdx++
+					iter.Next()
+				}
+				if err := iter.Close(); err != nil {
+					b.Fatal(err)
+				}
+				if keyIdx != len(expectedKeys) {
+					b.Fatalf("unexpected number of keys: got %d, want %d", keyIdx, len(expectedKeys))
+				}
+			}
+		})
+	}
+}
+
+// makeBranchStack creates a branch stack of the given size and initializes it with unique key-value pairs.
+func makeBranchStack(b *testing.B, stackSize int) Store[store.KVStore] {
+	parent := coretesting.NewMemKV()
+	branch := NewStore[store.KVStore](parent)
+	for i := 1; i < stackSize; i++ {
+		branch = NewStore[store.KVStore](branch)
+		for j := 0; j < elemsInStack; j++ {
+			// create unique keys by including the branch index.
+			key := []byte{byte(i), byte(j)}
+			value := []byte{byte(j)}
+			err := branch.Set(key, value)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	}
+	return branch
+}
@@ -15,7 +15,7 @@ const (
 
 var errKeyEmpty = errors.New("key cannot be empty")
 
-// changeSet implements the sorted cache for cachekv store,
+// changeSet implements the sorted tree for cachekv store,
 // we don't use MemDB here because cachekv is used extensively in sdk core path,
 // we need it to be as fast as possible, while `MemDB` is mainly used as a mocking db in unit tests.
 //