[Supplier] refactor: supplier module keys (#264)

* refactor: `NewMinedRelay` to shared testutil * refactor: claim & proof protobuf types * refactor: rename supplier keeper `UpsertClaim` & `UpsertProof` * refactor: misc. claim-side improvements * chore: add TODOs * refactor: supplier module keys * chore: review feedback improvements Co-authored-by: harry <[email protected]> Co-authored-by: Daniel Olshansky <[email protected]> * chore: review feedback improvements Co-authored-by: Daniel Olshansky <[email protected]> Co-authored-by: harry <[email protected]> * chore: add TODOs * trigger CI * chore: add TODO * chore: review feedback improvements Co-authored-by: Daniel Olshansky <[email protected]> * fix: usage raw string literal --------- Co-authored-by: harry <[email protected]> Co-authored-by: Daniel Olshansky <[email protected]>
pokt-network · Jan 10, 2024 · 1cc9ea4 · 1cc9ea4
1 parent 6b67ff9
commit 1cc9ea4
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 33 deletions.
diff --git a/x/supplier/keeper/claim.go b/x/supplier/keeper/claim.go
@@ -34,7 +34,7 @@ func (k Keeper) UpsertClaim(ctx sdk.Context, claim types.Claim) {
 
 	// Update the session end height index: sessionEndHeight -> [ClaimPrimaryKey]
 	sessionHeightStoreIndex := prefix.NewStore(parentStore, types.KeyPrefix(types.ClaimSessionEndHeightPrefix))
-	sessionEndBlockHeight := uint64(claim.GetSessionHeader().GetSessionEndBlockHeight())
+	sessionEndBlockHeight := claim.GetSessionHeader().GetSessionEndBlockHeight()
 	heightKey := types.ClaimSupplierEndSessionHeightKey(sessionEndBlockHeight, primaryKey)
 	sessionHeightStoreIndex.Set(heightKey, primaryKey)
 
@@ -61,7 +61,7 @@ func (k Keeper) RemoveClaim(ctx sdk.Context, sessionId, supplierAddr string) {
 	sessionHeightStoreIndex := prefix.NewStore(parentStore, types.KeyPrefix(types.ClaimSessionEndHeightPrefix))
 
 	addressKey := types.ClaimSupplierAddressKey(claim.SupplierAddress, primaryKey)
-	sessionEndBlockHeight := uint64(claim.GetSessionHeader().GetSessionEndBlockHeight())
+	sessionEndBlockHeight := claim.GetSessionHeader().GetSessionEndBlockHeight()
 	heightKey := types.ClaimSupplierEndSessionHeightKey(sessionEndBlockHeight, primaryKey)
 
 	// Delete all the entries (primary store and secondary indices)

diff --git a/x/supplier/types/key_claim.go b/x/supplier/types/key_claim.go
@@ -18,43 +18,22 @@ const (
 	ClaimSessionEndHeightPrefix = "Claim/height/"
 )
 
-// ClaimPrimaryKey returns the primary store key to retrieve a Claim by creating a composite key of the sessionId and supplierAddr
+// ClaimPrimaryKey returns the primary store key used to retrieve a Claim by creating a composite key of the sessionId and supplierAddr.
 func ClaimPrimaryKey(sessionId, supplierAddr string) []byte {
-	var key []byte
-
 	// We are guaranteed uniqueness of the primary key if it's a composite of the (sessionId, supplierAddr)
 	// because every supplier can only have one claim per session.
-	key = append(key, []byte(sessionId)...)
-	key = append(key, []byte("/")...)
-	key = append(key, []byte(supplierAddr)...)
-	key = append(key, []byte("/")...)
-
-	return key
+	return KeyComposite([]byte(sessionId), []byte(supplierAddr))
 }
 
-// ClaimSupplierAddressKey returns the address key to iterate through claims given a supplier Address
+// ClaimSupplierAddressKey returns the key used to iterate through claims given a supplier Address.
 func ClaimSupplierAddressKey(supplierAddr string, primaryKey []byte) []byte {
-	var key []byte
-
-	key = append(key, []byte(supplierAddr)...)
-	key = append(key, []byte("/")...)
-	key = append(key, primaryKey...)
-	key = append(key, []byte("/")...)
-
-	return key
+	return KeyComposite([]byte(supplierAddr), primaryKey)
 }
 
-// ClaimSupplierAddressKey returns the address key to iterate through claims given a supplier Address
-func ClaimSupplierEndSessionHeightKey(sessionEndHeight uint64, primaryKey []byte) []byte {
-	var key []byte
-
+// ClaimSupplierEndSessionHeightKey returns the key used to iterate through claims given a session end height.
+func ClaimSupplierEndSessionHeightKey(sessionEndHeight int64, primaryKey []byte) []byte {
 	heightBz := make([]byte, 8)
-	binary.BigEndian.PutUint64(heightBz, sessionEndHeight)
-
-	key = append(key, []byte(heightBz)...)
-	key = append(key, []byte("/")...)
-	key = append(key, primaryKey...)
-	key = append(key, []byte("/")...)
+	binary.BigEndian.PutUint64(heightBz, uint64(sessionEndHeight))
 
-	return key
+	return KeyComposite(heightBz, primaryKey)
 }
diff --git a/x/supplier/types/keys.go b/x/supplier/types/keys.go
@@ -1,5 +1,9 @@
 package types
 
+import (
+	"bytes"
+)
+
 const (
 	// ModuleName defines the module name
 	ModuleName = "supplier"
@@ -14,6 +18,15 @@ const (
 	MemStoreKey = "mem_supplier"
 )
 
-func KeyPrefix(p string) []byte {
-	return []byte(p)
+// KeyDelimiter is the delimiter for composite keys.
+var KeyDelimiter = []byte("/")
+
+// KeyPrefix returns the given prefix as a byte slice for use with the KVStore.
+func KeyPrefix(prefix string) []byte {
+	return []byte(prefix)
+}
+
+// KeyComposite combines the given keys into a single key for use with KVStore.
+func KeyComposite(keys ...[]byte) []byte {
+	return bytes.Join(keys, KeyDelimiter)
 }