status-im/status-go

package balance

import (
    "context"
    "math/big"
    "reflect"
    "sync"

    "github.com/ethereum/go-ethereum/common"
    "github.com/ethereum/go-ethereum/core/types"
)

// Reader interface for reading balance at a specified address.
type Reader interface {
    BalanceAt(ctx context.Context, account common.Address, blockNumber *big.Int) (*big.Int, error)
    NonceAt(ctx context.Context, account common.Address, blockNumber *big.Int) (uint64, error)
    HeaderByNumber(ctx context.Context, number *big.Int) (*types.Header, error)
    CallBlockHashByTransaction(ctx context.Context, blockNumber *big.Int, index uint) (common.Hash, error)
    NetworkID() uint64
}

// Cacher interface for caching balance to BalanceCache. Requires BalanceReader to fetch balance.
type Cacher interface {
    BalanceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*big.Int, error)
    NonceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*int64, error)
    Clear()
    Cache() CacheIface
}

// Interface for cache of balances.
type CacheIface interface {
    GetBalance(account common.Address, chainID uint64, blockNumber *big.Int) *big.Int
    GetNonce(account common.Address, chainID uint64, blockNumber *big.Int) *int64
    AddBalance(account common.Address, chainID uint64, blockNumber *big.Int, balance *big.Int)
    AddNonce(account common.Address, chainID uint64, blockNumber *big.Int, nonce *int64)
    BalanceSize(account common.Address, chainID uint64) int
    NonceSize(account common.Address, chainID uint64) int
    Clear()
}

type addressChainMap[T any] map[common.Address]map[uint64]T // address->chainID

type cacheIface[K comparable, V any] interface {
    get(K) V
    set(K, V)
    len() int
    keys() []K
    clear()
    init()
}

// genericCache is a generic implementation of CacheIface
type genericCache[B cacheIface[uint64, *big.Int], N cacheIface[uint64, *int64], NR cacheIface[int64, nonceRange]] struct {
    nonceRangeCache[NR]

    // balances maps an address and chain to a cache of a block number and the balance of this particular address on the chain
    balances addressChainMap[B]
    nonces   addressChainMap[N]
    rw       sync.RWMutex
}

func (b *genericCache[_, _, _]) GetBalance(account common.Address, chainID uint64, blockNumber *big.Int) *big.Int {
    b.rw.RLock()
    defer b.rw.RUnlock()

    _, exists := b.balances[account]
    if !exists {
        return nil
    }

    _, exists = b.balances[account][chainID]
    if !exists {
        return nil
    }

    return b.balances[account][chainID].get(blockNumber.Uint64())
}

func (b *genericCache[B, _, _]) AddBalance(account common.Address, chainID uint64, blockNumber *big.Int, balance *big.Int) {
    b.rw.Lock()
    defer b.rw.Unlock()

    _, exists := b.balances[account]
    if !exists {
        b.balances[account] = make(map[uint64]B)
    }

    _, exists = b.balances[account][chainID]
    if !exists {
        b.balances[account][chainID] = reflect.New(reflect.TypeOf(b.balances[account][chainID]).Elem()).Interface().(B)
        b.balances[account][chainID].init()
    }

    b.balances[account][chainID].set(blockNumber.Uint64(), balance)
}

func (b *genericCache[_, _, _]) GetNonce(account common.Address, chainID uint64, blockNumber *big.Int) *int64 {
    b.rw.RLock()
    defer b.rw.RUnlock()

    _, exists := b.nonces[account]
    if !exists {
        return nil
    }

    _, exists = b.nonces[account][chainID]
    if !exists {
        return nil
    }

    nonce := b.nonces[account][chainID].get(blockNumber.Uint64())
    if nonce != nil {
        return nonce
    }

    return b.findNonceInRange(account, chainID, blockNumber)
}

func (b *genericCache[_, N, _]) AddNonce(account common.Address, chainID uint64, blockNumber *big.Int, nonce *int64) {
    b.rw.Lock()
    defer b.rw.Unlock()

    _, exists := b.nonces[account]
    if !exists {
        b.nonces[account] = make(map[uint64]N)
    }

    _, exists = b.nonces[account][chainID]
    if !exists {
        b.nonces[account][chainID] = reflect.New(reflect.TypeOf(b.nonces[account][chainID]).Elem()).Interface().(N)
        b.nonces[account][chainID].init()
    }

    b.nonces[account][chainID].set(blockNumber.Uint64(), nonce)
    b.updateNonceRange(account, chainID, blockNumber, nonce)
}

func (b *genericCache[_, _, _]) BalanceSize(account common.Address, chainID uint64) int {
    b.rw.RLock()
    defer b.rw.RUnlock()

    _, exists := b.balances[account]
    if !exists {
        return 0
    }

    _, exists = b.balances[account][chainID]
    if !exists {
        return 0
    }

    return b.balances[account][chainID].len()
}

func (b *genericCache[_, N, _]) NonceSize(account common.Address, chainID uint64) int {
    b.rw.RLock()
    defer b.rw.RUnlock()

    _, exists := b.nonces[account]
    if !exists {
        return 0
    }

    _, exists = b.nonces[account][chainID]
    if !exists {
        return 0
    }

    return b.nonces[account][chainID].len()
}

// implements Cacher interface that caches balance and nonce in memory.
type cacherImpl struct {
    cache CacheIface
}

func newCacherImpl(cache CacheIface) *cacherImpl {
    return &cacherImpl{
        cache: cache,
    }
}

func (b *cacherImpl) BalanceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*big.Int, error) {
    cachedBalance := b.cache.GetBalance(account, client.NetworkID(), blockNumber)
    if cachedBalance != nil {
        return cachedBalance, nil
    }

    balance, err := client.BalanceAt(ctx, account, blockNumber)
    if err != nil {
        return nil, err
    }
    b.cache.AddBalance(account, client.NetworkID(), blockNumber, balance)
    return balance, nil
}

func (b *cacherImpl) NonceAt(ctx context.Context, client Reader, account common.Address, blockNumber *big.Int) (*int64, error) {
    cachedNonce := b.cache.GetNonce(account, client.NetworkID(), blockNumber)
    if cachedNonce != nil {
        return cachedNonce, nil
    }

    nonce, err := client.NonceAt(ctx, account, blockNumber)
    if err != nil {
        return nil, err
    }
    int64Nonce := int64(nonce)
    b.cache.AddNonce(account, client.NetworkID(), blockNumber, &int64Nonce)

    return &int64Nonce, nil
}

func (b *cacherImpl) Clear() {
    b.cache.Clear()
}

func (b *cacherImpl) Cache() CacheIface {
    return b.cache
}