From 065757e7e385753eaa2e528cc3f424edb31a05fa Mon Sep 17 00:00:00 2001
From: Mostafa <mostafa.sedaghat@gmail.com>
Date: Tue, 12 Sep 2023 23:02:03 +0800
Subject: [PATCH] fix: restoring at the first block

---
 sandbox/sandbox.go      | 12 ++++--------
 sandbox/sandbox_test.go |  5 ++---
 state/state.go          | 11 ++++++-----
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/sandbox/sandbox.go b/sandbox/sandbox.go
index d5db33dcc..5ef61b8c8 100644
--- a/sandbox/sandbox.go
+++ b/sandbox/sandbox.go
@@ -26,6 +26,7 @@ type sandbox struct {
 	accounts        map[crypto.Address]*sandboxAccount
 	validators      map[crypto.Address]*sandboxValidator
 	params          param.Params
+	height          uint32
 	totalAccounts   int32
 	totalValidators int32
 	totalPower      int64
@@ -43,10 +44,11 @@ type sandboxAccount struct {
 	updated bool
 }
 
-func NewSandbox(store store.Reader, params param.Params,
+func NewSandbox(height uint32, store store.Reader, params param.Params,
 	committee committee.Reader, totalPower int64,
 ) Sandbox {
 	sb := &sandbox{
+		height:     height,
 		store:      store,
 		committee:  committee,
 		totalPower: totalPower,
@@ -218,13 +220,7 @@ func (sb *sandbox) CurrentHeight() uint32 {
 	sb.lk.RLock()
 	defer sb.lk.RUnlock()
 
-	return sb.currentHeight()
-}
-
-func (sb *sandbox) currentHeight() uint32 {
-	h, _ := sb.store.LastCertificate()
-
-	return h + 1
+	return sb.height + 1
 }
 
 func (sb *sandbox) IterateAccounts(
diff --git a/sandbox/sandbox_test.go b/sandbox/sandbox_test.go
index 58fa4c6ab..279dcafb4 100644
--- a/sandbox/sandbox_test.go
+++ b/sandbox/sandbox_test.go
@@ -48,15 +48,14 @@ func setup(t *testing.T) *testData {
 		totalPower += val.Power()
 	}
 
-	sandbox := NewSandbox(store, params, committee, totalPower).(*sandbox)
-
-	assert.Equal(t, sandbox.CurrentHeight(), uint32(1))
 	lastHeight := uint32(21)
 	for i := uint32(1); i < lastHeight; i++ {
 		b := ts.GenerateTestBlock(nil)
 		c := ts.GenerateTestCertificate()
 		store.SaveBlock(i, b, c)
 	}
+	sandbox := NewSandbox(store.LastHeight,
+		store, params, committee, totalPower).(*sandbox)
 	assert.Equal(t, sandbox.CurrentHeight(), lastHeight)
 	assert.Equal(t, sandbox.Params(), params)
 
diff --git a/state/state.go b/state/state.go
index a69a8c7c3..09114a49d 100644
--- a/state/state.go
+++ b/state/state.go
@@ -67,15 +67,15 @@ func LoadOrNewState(
 	st.logger = logger.NewSubLogger("_state", st)
 	st.store = store
 
-	// The first account is Treasury Account at the genesis time.
-	// So if we have more account, we are not in the genesis height anymore.
-	if store.TotalAccounts() > 1 {
+	// Check if the number of accounts is greater than the genesis time;
+	// this indicates we are not at the genesis height anymore.
+	if store.TotalAccounts() > int32(len(genDoc.Accounts())) {
 		err := st.tryLoadLastInfo()
 		if err != nil {
 			return nil, err
 		}
 	} else {
-		// We are at the genesis height
+		// We are at the genesis height.
 		err := st.makeGenesisState(genDoc)
 		if err != nil {
 			return nil, err
@@ -94,7 +94,8 @@ func LoadOrNewState(
 }
 
 func (st *state) concreteSandbox() sandbox.Sandbox {
-	return sandbox.NewSandbox(st.store, st.params, st.committee, st.totalPower)
+	return sandbox.NewSandbox(st.lastInfo.BlockHeight(),
+		st.store, st.params, st.committee, st.totalPower)
 }
 
 func (st *state) tryLoadLastInfo() error {