lotus/tvx/drivers/state_wrapper.go

package drivers

import (
	"context"
	"fmt"

	"github.com/filecoin-project/go-address"

	"github.com/filecoin-project/lotus/chain/state"
	"github.com/filecoin-project/lotus/lib/blockstore"
	"github.com/filecoin-project/specs-actors/actors/runtime"
	"github.com/ipfs/go-datastore"
	cbor "github.com/ipfs/go-ipld-cbor"

	"github.com/filecoin-project/specs-actors/actors/abi"
	"github.com/filecoin-project/specs-actors/actors/abi/big"
	"github.com/ipfs/go-cid"

	"github.com/filecoin-project/lotus/chain/types"
)

type StateWrapper struct {
	// The blockstore underlying the state tree and storage.
	bs blockstore.Blockstore

	ds datastore.Batching
	// HAMT-CBOR store on top of the blockstore.
	cst cbor.IpldStore

	// CID of the root of the state tree.
	stateRoot cid.Cid
}

func NewStateWrapper() *StateWrapper {
	bs := blockstore.NewTemporary()
	cst := cbor.NewCborStore(bs)
	// Put EmptyObjectCid value in the store. When an actor is initially created its Head is set to this value.
	_, err := cst.Put(context.TODO(), map[string]string{})
	if err != nil {
		panic(err)
	}

	treeImpl, err := state.NewStateTree(cst)
	if err != nil {
		panic(err) // Never returns error, the error return should be removed.
	}
	root, err := treeImpl.Flush(context.TODO())
	if err != nil {
		panic(err)
	}
	return &StateWrapper{
		bs:        bs,
		ds:        datastore.NewMapDatastore(),
		cst:       cst,
		stateRoot: root,
	}
}

func (s *StateWrapper) Root() cid.Cid {
	return s.stateRoot
}

// StoreGet the value at key from vm store
func (s *StateWrapper) StoreGet(key cid.Cid, out runtime.CBORUnmarshaler) error {
	tree, err := state.LoadStateTree(s.cst, s.stateRoot)
	if err != nil {
		return err
	}
	return tree.Store.Get(context.Background(), key, out)
}

// StorePut `value` into vm store
func (s *StateWrapper) StorePut(value runtime.CBORMarshaler) (cid.Cid, error) {
	tree, err := state.LoadStateTree(s.cst, s.stateRoot)
	if err != nil {
		return cid.Undef, err
	}
	return tree.Store.Put(context.Background(), value)
}

func (s *StateWrapper) Actor(addr address.Address) (Actor, error) {
	tree, err := state.LoadStateTree(s.cst, s.stateRoot)
	if err != nil {
		return nil, err
	}
	fcActor, err := tree.GetActor(addr)
	if err != nil {
		return nil, err
	}
	return &actorWrapper{*fcActor}, nil
}

func (s *StateWrapper) SetActorState(addr address.Address, balance abi.TokenAmount, actorState runtime.CBORMarshaler) (Actor, error) {
	tree, err := state.LoadStateTree(s.cst, s.stateRoot)
	if err != nil {
		return nil, err
	}
	// actor should exist
	act, err := tree.GetActor(addr)
	if err != nil {
		return nil, err
	}
	// add the state to the store and get a new head cid
	actHead, err := tree.Store.Put(context.Background(), actorState)
	if err != nil {
		return nil, err
	}
	// update the actor object with new head and balance parameter
	actr := &actorWrapper{types.Actor{
		Code:  act.Code,
		Nonce: act.Nonce,
		// updates
		Head:    actHead,
		Balance: balance,
	}}
	if err := tree.SetActor(addr, &actr.Actor); err != nil {
		return nil, err
	}
	return actr, s.flush(tree)
}

func (s *StateWrapper) CreateActor(code cid.Cid, addr address.Address, balance abi.TokenAmount, actorState runtime.CBORMarshaler) (Actor, address.Address, error) {
	idAddr := addr
	tree, err := state.LoadStateTree(s.cst, s.stateRoot)
	if err != nil {
		return nil, address.Undef, err
	}
	if addr.Protocol() != address.ID {

		actHead, err := tree.Store.Put(context.Background(), actorState)
		if err != nil {
			return nil, address.Undef, err
		}
		actr := &actorWrapper{types.Actor{
			Code:    code,
			Head:    actHead,
			Balance: balance,
		}}

		idAddr, err = tree.RegisterNewAddress(addr)
		if err != nil {
			return nil, address.Undef, fmt.Errorf("register new address for actor: %w", err)
		}

		if err := tree.SetActor(addr, &actr.Actor); err != nil {
			return nil, address.Undef, fmt.Errorf("setting new actor for actor: %w", err)
		}
	}

	// store newState
	head, err := tree.Store.Put(context.Background(), actorState)
	if err != nil {
		return nil, address.Undef, err
	}

	// create and store actor object
	a := types.Actor{
		Code:    code,
		Head:    head,
		Balance: balance,
	}
	if err := tree.SetActor(idAddr, &a); err != nil {
		return nil, address.Undef, err
	}

	return &actorWrapper{a}, idAddr, s.flush(tree)
}

// Flushes a state tree to storage and sets this state's root to that tree's root CID.
func (s *StateWrapper) flush(tree *state.StateTree) (err error) {
	s.stateRoot, err = tree.Flush(context.TODO())
	return
}

//
// Actor Wrapper
//

type actorWrapper struct {
	types.Actor
}

func (a *actorWrapper) Code() cid.Cid {
	return a.Actor.Code
}

func (a *actorWrapper) Head() cid.Cid {
	return a.Actor.Head
}

func (a *actorWrapper) CallSeqNum() uint64 {
	return a.Actor.Nonce
}

func (a *actorWrapper) Balance() big.Int {
	return a.Actor.Balance

}