Fix locked and vested account balances

docs/api/docs.go

@@ -1431,6 +1431,9 @@ const docTemplate = `{
                 },
                 "reserved": {
                     "type": "number"
+                },
+                "vested": {
+                    "type": "number"
                 }
             }
         },

docs/api/swagger.json

@@ -1420,6 +1420,9 @@
                 },
                 "reserved": {
                     "type": "number"
+                },
+                "vested": {
+                    "type": "number"
                 }
             }
         },
@@ -3240,4 +3243,4 @@
             }
         }
     }
-}
+}

docs/api/swagger.yaml

@@ -11,6 +11,8 @@ definitions:
         type: integer
       reserved:
         type: number
+      vested:
+        type: number
     type: object
   github_com_itering_subscan_plugins_balance_model.Transfer:
     properties:

plugins/balance/dao/account.go

@@ -7,7 +7,6 @@ import (
 	bModel "github.com/itering/subscan/plugins/balance/model"
 	"github.com/itering/subscan/util/address"
 	"github.com/itering/substrate-api-rpc/rpc"
-	"github.com/shopspring/decimal"
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 )
@@ -70,24 +69,49 @@ func RefreshAccount(ctx context.Context, s *Storage, accountId string) error {
 	if q.RowsAffected == 1 {
 		_, _ = s.Pool.HINCRBY(ctx, model.MetadataCacheKey(), "total_account", 1)
 	}
-	return AfterAccountCreate(ctx, db, &account)
+	currentBlock, _ := s.Dao.GetCurrentBlockNum(ctx)
+	return AfterAccountCreate(ctx, db, &account, currentBlock)
 }
 
-func AfterAccountCreate(ctx context.Context, db *gorm.DB, account *bModel.Account) error {
+func AfterAccountCreate(ctx context.Context, db *gorm.DB, account *bModel.Account, currentBlock uint64) error {
 	accountDataRaw, err := rpc.ReadStorage(nil, "system", "account", "", account.Address)
 	if err != nil {
 		return err
 	}
 	accountData := new(bModel.AccountData)
 	accountDataRaw.ToAny(accountData)
+	locks := ReadBalanceLocks(account.Address)
+	vesting := ReadVesting(account.Address)
+	lockSummary := bModel.AccountLockSummary(accountData, locks)
 	return db.WithContext(ctx).Model(account).Where("address = ?", account.Address).UpdateColumns(map[string]interface{}{
 		"nonce":    accountData.Nonce,
 		"balance":  accountData.Data.Free.Add(accountData.Data.Reserved),
-		"locked":   decimal.Max(accountData.Data.MiscFrozen, accountData.Data.FeeFrozen),
+		"locked":   lockSummary.Locked,
 		"reserved": accountData.Data.Reserved,
+		"vested":   bModel.SummarizeVesting(vesting, currentBlock),
 	}).Error
 }
 
+func ReadBalanceLocks(accountId string) []bModel.BalanceLock {
+	locksRaw, err := rpc.ReadStorage(nil, "balances", "locks", "", accountId)
+	if err != nil {
+		return nil
+	}
+	var locks []bModel.BalanceLock
+	locksRaw.ToAny(&locks)
+	return locks
+}
+
+func ReadVesting(accountId string) []bModel.VestingInfo {
+	vestingRaw, err := rpc.ReadStorage(nil, "vesting", "vesting", "", accountId)
+	if err != nil {
+		return nil
+	}
+	var vesting []bModel.VestingInfo
+	vestingRaw.ToAny(&vesting)
+	return vesting
+}
+
 func (s *Storage) AddOrUpdateItem(c context.Context, item interface{}, keys []string, updates ...string) *gorm.DB {
 	var keyFields []clause.Column
 	for _, key := range keys {

plugins/balance/dao/script.go

@@ -9,28 +9,33 @@ import (
 	"github.com/itering/subscan/util"
 	"github.com/itering/subscan/util/address"
 	"github.com/panjf2000/ants/v2"
-	"github.com/shopspring/decimal"
 	"gorm.io/gorm"
 	"log"
 	"sync"
 )
 
 func InitAccount(sg *Storage) {
 	ctx := context.Background()
+	locksByAddress := readAllBalanceLocks(ctx)
+	vestingByAddress := readAllVesting(ctx)
+	currentBlock, _ := sg.Dao.GetCurrentBlockNum(ctx)
 	wg := new(sync.WaitGroup)
 	bp, _ := ants.NewPoolWithFunc(10, func(i interface{}) {
 		wg.Add(1)
 		defer wg.Done()
 		params := i.([]interface{})
 		addr := params[0].(string)
 		info := params[1].(*bModel.AccountData)
+		lockSummary := bModel.AccountLockSummary(info, locksByAddress[addr])
+		vested := bModel.SummarizeVesting(vestingByAddress[addr], currentBlock)
 		sg.AddOrUpdateItem(ctx, &bModel.Account{
 			Address:  addr,
 			Nonce:    info.Nonce,
 			Balance:  info.Data.Free.Add(info.Data.Reserved),
-			Locked:   decimal.Max(info.Data.MiscFrozen, info.Data.FeeFrozen),
+			Locked:   lockSummary.Locked,
 			Reserved: info.Data.Reserved,
-		}, []string{"address"}, "nonce", "balance", "locked", "reserved")
+			Vested:   vested,
+		}, []string{"address"}, "nonce", "balance", "locked", "reserved", "vested")
 	})
 	defer bp.Release()
 
@@ -51,6 +56,50 @@ func InitAccount(sg *Storage) {
 	wg.Wait()
 }
 
+func readAllBalanceLocks(ctx context.Context) map[string][]bModel.BalanceLock {
+	result := map[string][]bModel.BalanceLock{}
+	_ = substrate.BatchReadKeysPaged(ctx, "Balances", "Locks", "", func(keys []string, scaleType string) error {
+		r, _ := substrate.BatchStorageByKey(ctx, keys, scaleType, "")
+		for key, v := range r {
+			val, err := substrate.ParseStorageKey(key)
+			if err != nil || len(val) == 0 {
+				continue
+			}
+			addr := address.Format(val[0].ToString())
+			if addr == "" {
+				continue
+			}
+			var locks []bModel.BalanceLock
+			v.ToAny(&locks)
+			result[addr] = locks
+		}
+		return nil
+	})
+	return result
+}
+
+func readAllVesting(ctx context.Context) map[string][]bModel.VestingInfo {
+	result := map[string][]bModel.VestingInfo{}
+	_ = substrate.BatchReadKeysPaged(ctx, "Vesting", "Vesting", "", func(keys []string, scaleType string) error {
+		r, _ := substrate.BatchStorageByKey(ctx, keys, scaleType, "")
+		for key, v := range r {
+			val, err := substrate.ParseStorageKey(key)
+			if err != nil || len(val) == 0 {
+				continue
+			}
+			addr := address.Format(val[0].ToString())
+			if addr == "" {
+				continue
+			}
+			var vesting []bModel.VestingInfo
+			v.ToAny(&vesting)
+			result[addr] = vesting
+		}
+		return nil
+	})
+	return result
+}
+
 func RefreshAllAccount(_ *Storage) {
 
 }

plugins/balance/model/model.go

@@ -1,8 +1,6 @@
 package model
 
-import (
-	"github.com/shopspring/decimal"
-)
+import "github.com/shopspring/decimal"
 
 type Account struct {
 	ID       uint            `gorm:"primary_key" json:"-"`
@@ -11,6 +9,7 @@ type Account struct {
 	Balance  decimal.Decimal `json:"balance" gorm:"type:decimal(65,0);index:balance;index:balance_address,priority:1"`
 	Locked   decimal.Decimal `json:"locked" gorm:"type:decimal(65,0);"`
 	Reserved decimal.Decimal `json:"reserved" gorm:"type:decimal(65,0);"`
+	Vested   decimal.Decimal `json:"vested" gorm:"type:decimal(65,0);"`
 }
 
 func (a *Account) TableName() string {
@@ -25,9 +24,69 @@ type AccountData struct {
 		Reserved   decimal.Decimal `json:"reserved"`
 		MiscFrozen decimal.Decimal `json:"miscFrozen"`
 		FeeFrozen  decimal.Decimal `json:"feeFrozen"`
+		Frozen     decimal.Decimal `json:"frozen"`
 	} `json:"data"`
 }
 
+type BalanceLock struct {
+	ID     string          `json:"id"`
+	Amount decimal.Decimal `json:"amount"`
+}
+
+type LockSummary struct {
+	Locked decimal.Decimal
+}
+
+func (a AccountData) LockedBalance() decimal.Decimal {
+	return decimal.Max(a.Data.Frozen, decimal.Max(a.Data.MiscFrozen, a.Data.FeeFrozen))
+}
+
+func SummarizeLocks(locks []BalanceLock) LockSummary {
+	var summary LockSummary
+	for _, lock := range locks {
+		if lock.Amount.GreaterThan(summary.Locked) {
+			summary.Locked = lock.Amount
+		}
+	}
+	return summary
+}
+
+func AccountLockSummary(accountData *AccountData, locks []BalanceLock) LockSummary {
+	summary := SummarizeLocks(locks)
+	if accountData == nil {
+		return summary
+	}
+	if dataLocked := accountData.LockedBalance(); dataLocked.GreaterThan(summary.Locked) {
+		summary.Locked = dataLocked
+	}
+	return summary
+}
+
+type VestingInfo struct {
+	Locked        decimal.Decimal `json:"locked"`
+	PerBlock      decimal.Decimal `json:"perBlock"`
+	StartingBlock uint64          `json:"startingBlock"`
+}
+
+func (v VestingInfo) VestedAt(blockNum uint64) decimal.Decimal {
+	if blockNum <= v.StartingBlock {
+		return decimal.Zero
+	}
+	vested := v.PerBlock.Mul(decimal.NewFromUint64(blockNum - v.StartingBlock))
+	if vested.GreaterThan(v.Locked) {
+		return v.Locked
+	}
+	return vested
+}
+
+func SummarizeVesting(vesting []VestingInfo, blockNum uint64) decimal.Decimal {
+	var vested decimal.Decimal
+	for _, schedule := range vesting {
+		vested = vested.Add(schedule.VestedAt(blockNum))
+	}
+	return vested
+}
+
 type Transfer struct {
 	Id             uint            `json:"id" gorm:"primary_key;autoIncrement:false"`
 	BlockNum       uint            `json:"blockNum" gorm:"size:32"`

plugins/balance/model/model_test.go

@@ -0,0 +1,57 @@
+package model
+
+import (
+	"testing"
+
+	"github.com/shopspring/decimal"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAccountDataLockedBalanceSupportsLegacyAndCurrentFields(t *testing.T) {
+	accountData := AccountData{}
+	accountData.Data.MiscFrozen = decimal.NewFromInt(3)
+	accountData.Data.FeeFrozen = decimal.NewFromInt(5)
+	accountData.Data.Frozen = decimal.NewFromInt(7)
+
+	assert.Equal(t, "7", accountData.LockedBalance().String())
+}
+
+func TestAccountLockSummaryUsesMaxLockedValue(t *testing.T) {
+	accountData := AccountData{}
+	accountData.Data.Frozen = decimal.NewFromInt(10)
+
+	summary := AccountLockSummary(&accountData, []BalanceLock{
+		{ID: "0x7374616b696e6720", Amount: decimal.NewFromInt(12)},
+	})
+
+	assert.Equal(t, "12", summary.Locked.String())
+}
+
+func TestSummarizeVestingCalculatesVestedAmount(t *testing.T) {
+	summary := SummarizeVesting([]VestingInfo{
+		{
+			Locked:        decimal.NewFromInt(100),
+			PerBlock:      decimal.NewFromInt(3),
+			StartingBlock: 10,
+		},
+		{
+			Locked:        decimal.NewFromInt(50),
+			PerBlock:      decimal.NewFromInt(10),
+			StartingBlock: 12,
+		},
+	}, 16)
+
+	assert.Equal(t, "58", summary.String())
+}
+
+func TestVestingInfoVestedAtCapsAtLocked(t *testing.T) {
+	schedule := VestingInfo{
+		Locked:        decimal.NewFromInt(100),
+		PerBlock:      decimal.NewFromInt(3),
+		StartingBlock: 10,
+	}
+
+	assert.Equal(t, "0", schedule.VestedAt(10).String())
+	assert.Equal(t, "15", schedule.VestedAt(15).String())
+	assert.Equal(t, "100", schedule.VestedAt(100).String())
+}

ui-react/src/pages/sub/account/[id].tsx

@@ -37,7 +37,7 @@ export default function Page() {
     } else {
       return BIG_ZERO
     }
-  }, [accountData?.reserved, accountData?.locked, accountData?.balance, metadata?.enabledNewTransferableFormulas])
+  }, [accountData, metadata?.enabledNewTransferableFormulas])
 
   return (
     <PageContent>
@@ -69,6 +69,11 @@ export default function Page() {
                       <div>{getBalanceAmount(new BigNumber(accountData.locked), token?.decimals).toFormat()}</div>
                     </div>
                     <Divider className="my-2.5" />
+                    <div className="flex items-center">
+                      <div className="w-48">Vested</div>
+                      <div>{getBalanceAmount(new BigNumber(accountData.vested || 0), token?.decimals).toFormat()}</div>
+                    </div>
+                    <Divider className="my-2.5" />
                     <div className="flex items-center">
                       <div className="w-48">Reserved</div>
                       <div>{getBalanceAmount(new BigNumber(accountData.reserved), token?.decimals).toFormat()}</div>

ui-react/src/utils/api.ts

@@ -256,6 +256,7 @@ export type accountType = {
     locked: string
     nonce: string
     reserved: string
+    vested?: string
 }
 
 type getAccountParams = {