From f8d5f290781e4b69e7a35001e9b9387f83e61e5b Mon Sep 17 00:00:00 2001
From: MiniFrenchBread <103425574+MiniFrenchBread@users.noreply.github.com>
Date: Fri, 14 Jun 2024 18:29:01 +0800
Subject: [PATCH] refactor: epoch quorum storage

---
 x/dasigners/v1/genesis.go           | 14 ++++++++++---
 x/dasigners/v1/keeper/grpc_query.go | 32 ++++++++++-------------------
 x/dasigners/v1/keeper/keeper.go     | 29 +++++++++++++++++---------
 x/dasigners/v1/types/keys.go        |  6 ++++--
 4 files changed, 45 insertions(+), 36 deletions(-)

diff --git a/x/dasigners/v1/genesis.go b/x/dasigners/v1/genesis.go
index 71a0339c..deb04433 100644
--- a/x/dasigners/v1/genesis.go
+++ b/x/dasigners/v1/genesis.go
@@ -40,11 +40,19 @@ func ExportGenesis(ctx sdk.Context, keeper keeper.Keeper) *types.GenesisState {
 	})
 	epochQuorums := make([]*types.Quorums, 0)
 	for i := 0; i < int(epochNumber); i += 1 {
-		quorums, found := keeper.GetEpochQuorums(ctx, uint64(i))
-		if !found {
+		quorumCnt, err := keeper.GetQuorumCount(ctx, uint64(i))
+		if err != nil {
 			panic("historical quorums not found")
 		}
-		epochQuorums = append(epochQuorums, &quorums)
+		quorums := make([]*types.Quorum, quorumCnt)
+		for quorumId := uint64(0); quorumId < quorumCnt; quorumId += 1 {
+			quorum, err := keeper.GetEpochQuorum(ctx, uint64(i), quorumId)
+			if err != nil {
+				panic("failed to load historical quorum")
+			}
+			quorums[quorumId] = &quorum
+		}
+		epochQuorums = append(epochQuorums, &types.Quorums{Quorums: quorums})
 	}
 	return types.NewGenesisState(params, epochNumber, signers, epochQuorums)
 }
diff --git a/x/dasigners/v1/keeper/grpc_query.go b/x/dasigners/v1/keeper/grpc_query.go
index ddc64093..36208f44 100644
--- a/x/dasigners/v1/keeper/grpc_query.go
+++ b/x/dasigners/v1/keeper/grpc_query.go
@@ -57,26 +57,20 @@ func (k Keeper) QuorumCount(
 
 func (k Keeper) EpochQuorum(c context.Context, request *types.QueryEpochQuorumRequest) (*types.QueryEpochQuorumResponse, error) {
 	ctx := sdk.UnwrapSDKContext(c)
-	quorums, found := k.GetEpochQuorums(ctx, request.EpochNumber)
-	if !found {
-		return nil, types.ErrQuorumNotFound
+	quorum, err := k.GetEpochQuorum(ctx, request.EpochNumber, request.QuorumId)
+	if err != nil {
+		return nil, err
 	}
-	if len(quorums.Quorums) <= int(request.QuorumId) {
-		return nil, types.ErrQuorumIdOutOfBound
-	}
-	return &types.QueryEpochQuorumResponse{Quorum: quorums.Quorums[request.QuorumId]}, nil
+	return &types.QueryEpochQuorumResponse{Quorum: &quorum}, nil
 }
 
 func (k Keeper) EpochQuorumRow(c context.Context, request *types.QueryEpochQuorumRowRequest) (*types.QueryEpochQuorumRowResponse, error) {
 	ctx := sdk.UnwrapSDKContext(c)
-	quorums, found := k.GetEpochQuorums(ctx, request.EpochNumber)
-	if !found {
-		return nil, types.ErrQuorumNotFound
+	quorum, err := k.GetEpochQuorum(ctx, request.EpochNumber, request.QuorumId)
+	if err != nil {
+		return nil, err
 	}
-	if len(quorums.Quorums) <= int(request.QuorumId) {
-		return nil, types.ErrQuorumIdOutOfBound
-	}
-	signers := quorums.Quorums[request.QuorumId].Signers
+	signers := quorum.Signers
 	if len(signers) <= int(request.RowIndex) {
 		return nil, types.ErrRowIndexOutOfBound
 	}
@@ -85,14 +79,10 @@ func (k Keeper) EpochQuorumRow(c context.Context, request *types.QueryEpochQuoru
 
 func (k Keeper) AggregatePubkeyG1(c context.Context, request *types.QueryAggregatePubkeyG1Request) (*types.QueryAggregatePubkeyG1Response, error) {
 	ctx := sdk.UnwrapSDKContext(c)
-	quorums, found := k.GetEpochQuorums(ctx, request.EpochNumber)
-	if !found {
-		return nil, types.ErrQuorumNotFound
+	quorum, err := k.GetEpochQuorum(ctx, request.EpochNumber, request.QuorumId)
+	if err != nil {
+		return nil, err
 	}
-	if len(quorums.Quorums) <= int(request.QuorumId) {
-		return nil, types.ErrQuorumIdOutOfBound
-	}
-	quorum := quorums.Quorums[request.QuorumId]
 	if (len(quorum.Signers)+7)/8 != len(request.QuorumBitmap) {
 		return nil, types.ErrQuorumBitmapLengthMismatch
 	}
diff --git a/x/dasigners/v1/keeper/keeper.go b/x/dasigners/v1/keeper/keeper.go
index 530b84c9..7dfcfd3a 100644
--- a/x/dasigners/v1/keeper/keeper.go
+++ b/x/dasigners/v1/keeper/keeper.go
@@ -142,21 +142,30 @@ func (k Keeper) IterateSigners(ctx sdk.Context, fn func(index int64, signer type
 	}
 }
 
-func (k Keeper) GetEpochQuorums(ctx sdk.Context, epoch uint64) (types.Quorums, bool) {
-	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.EpochQuorumsKeyPrefix)
-	bz := store.Get(types.GetEpochQuorumsKeyFromEpoch(epoch))
-	if bz == nil {
-		return types.Quorums{}, false
+func (k Keeper) GetEpochQuorum(ctx sdk.Context, epoch uint64, quorumId uint64) (types.Quorum, error) {
+	quorumCount, err := k.GetQuorumCount(ctx, epoch)
+	if err != nil {
+		return types.Quorum{}, err
 	}
-	var quorums types.Quorums
-	k.cdc.MustUnmarshal(bz, &quorums)
-	return quorums, true
+	if quorumCount <= quorumId {
+		return types.Quorum{}, types.ErrQuorumIdOutOfBound
+	}
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.EpochQuorumsKeyPrefix)
+	bz := store.Get(types.GetEpochQuorumKey(epoch, quorumId))
+	if bz == nil {
+		return types.Quorum{}, types.ErrQuorumNotFound
+	}
+	var quorum types.Quorum
+	k.cdc.MustUnmarshal(bz, &quorum)
+	return quorum, nil
 }
 
 func (k Keeper) SetEpochQuorums(ctx sdk.Context, epoch uint64, quorums types.Quorums) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.EpochQuorumsKeyPrefix)
-	bz := k.cdc.MustMarshal(&quorums)
-	store.Set(types.GetEpochQuorumsKeyFromEpoch(epoch), bz)
+	for quorumId, quorum := range quorums.Quorums {
+		bz := k.cdc.MustMarshal(quorum)
+		store.Set(types.GetEpochQuorumKey(epoch, uint64(quorumId)), bz)
+	}
 	k.SetQuorumCount(ctx, epoch, uint64(len(quorums.Quorums)))
 }
 
diff --git a/x/dasigners/v1/types/keys.go b/x/dasigners/v1/types/keys.go
index 163cdd09..a821d5cb 100644
--- a/x/dasigners/v1/types/keys.go
+++ b/x/dasigners/v1/types/keys.go
@@ -33,8 +33,10 @@ func GetSignerKeyFromAccount(account string) ([]byte, error) {
 	return hex.DecodeString(account)
 }
 
-func GetEpochQuorumsKeyFromEpoch(epoch uint64) []byte {
-	return sdk.Uint64ToBigEndian(epoch)
+func GetEpochQuorumKey(epoch uint64, quorumId uint64) []byte {
+	b := sdk.Uint64ToBigEndian(epoch)
+	b = append(b, sdk.Uint64ToBigEndian(quorumId)...)
+	return b
 }
 
 func GetQuorumCountKey(epoch uint64) []byte {