feat!: LegacyActorTree vs ActorTree, use HAMT v4 with generics

Ref: filecoin-project/go-hamt-ipld#122

builtin/actor_tree.go

@@ -1,17 +1,20 @@
 package builtin
 
 import (
+	"crypto/sha256"
+
 	"github.com/filecoin-project/go-address"
+	hamt "github.com/filecoin-project/go-hamt-ipld/v4"
 	"github.com/filecoin-project/go-state-types/abi"
 	"github.com/filecoin-project/go-state-types/big"
-	"github.com/filecoin-project/go-state-types/builtin/v8/util/adt"
+	"github.com/filecoin-project/go-state-types/builtin/v14/util/adt"
 	"github.com/ipfs/go-cid"
 	"golang.org/x/xerrors"
 )
 
-type ActorTree interface {
+type LegacyActorTree interface {
 	GetStore() adt.Store
-
+	GetMap() *adt.Map
 	Flush() (cid.Cid, error)
 	GetActorV4(addr address.Address) (*ActorV4, bool, error)
 	GetActorV5(addr address.Address) (*ActorV5, bool, error)
@@ -22,7 +25,19 @@ type ActorTree interface {
 	ForEachKey(fn func(addr address.Address) error) error
 }
 
-var _ ActorTree = (*actorTree)(nil)
+type ActorTree interface {
+	GetStore() adt.Store
+	Flush() (cid.Cid, error)
+	GetActorV5(addr address.Address) (*ActorV5, bool, error)
+	SetActorV5(addr address.Address, actor *ActorV5) error
+	ForEachV5(fn func(addr address.Address, actor *ActorV5) error) error
+	ForEachKey(fn func(addr address.Address) error) error
+}
+
+var _ LegacyActorTree = (*legacyActorTree)(nil)
+var _ ActorTree = (*legacyActorTree)(nil)
+
+// var _ ActorTree = (*actorTree)(nil)
 
 // Value type of the top level of the state tree.
 // Represents the on-chain state of a single actor.
@@ -43,47 +58,173 @@ type ActorV5 struct {
 	DelegatedAddress *address.Address // Delegated (f4) actor address
 }
 
+func (a *ActorV5) Equals(o *ActorV5) bool {
+	if a == nil && o == nil {
+		return true
+	}
+	if a == nil || o == nil {
+		return false
+	}
+	return a.Code == o.Code &&
+		a.Head == o.Head &&
+		a.CallSeqNum == o.CallSeqNum &&
+		a.Balance.Equals(o.Balance) &&
+		((a.DelegatedAddress == nil && o.DelegatedAddress == nil) || (a.DelegatedAddress != nil && o.DelegatedAddress != nil && *a.DelegatedAddress == *o.DelegatedAddress))
+}
+
+func (a *ActorV5) New() *ActorV5 {
+	return new(ActorV5)
+}
+
 // A specialization of a map of ID-addresses to actor heads.
 type actorTree struct {
+	Map      *hamt.Node[*ActorV5]
+	LastRoot cid.Cid
+	Store    adt.Store
+}
+
+// Initializes a new, empty state tree backed by a store.
+func NewTree(store adt.Store) (ActorTree, error) {
+	nd, err := hamt.NewNode[*ActorV5](
+		store,
+		hamt.UseHashFunction(func(input []byte) []byte {
+			res := sha256.Sum256(input)
+			return res[:]
+		}),
+		hamt.UseTreeBitWidth(DefaultHamtBitwidth),
+	)
+	if err != nil {
+		return nil, err
+	}
+	return &actorTree{
+		Map:      nd,
+		LastRoot: cid.Undef,
+		Store:    store,
+	}, nil
+}
+
+// Loads a tree from a root CID and store.
+func LoadTree(store adt.Store, root cid.Cid) (ActorTree, error) {
+	nd, err := hamt.LoadNode[*ActorV5](
+		store.Context(),
+		store,
+		root,
+		hamt.UseHashFunction(func(input []byte) []byte {
+			res := sha256.Sum256(input)
+			return res[:]
+		}),
+		hamt.UseTreeBitWidth(DefaultHamtBitwidth),
+	)
+	if err != nil {
+		return nil, err
+	}
+	return &actorTree{
+		Map:      nd,
+		LastRoot: root,
+		Store:    store,
+	}, nil
+}
+
+func (t *actorTree) GetStore() adt.Store {
+	return t.Store
+}
+
+// Writes the tree root node to the store, and returns its CID.
+func (t *actorTree) Flush() (cid.Cid, error) {
+	if err := t.Map.Flush(t.Store.Context()); err != nil {
+		return cid.Undef, xerrors.Errorf("failed to flush map root: %w", err)
+	}
+	c, err := t.Store.Put(t.Store.Context(), t.Map)
+	if err != nil {
+		return cid.Undef, xerrors.Errorf("writing map root object: %w", err)
+	}
+	t.LastRoot = c
+	return c, nil
+}
+
+func (t *actorTree) GetActorV5(addr address.Address) (*ActorV5, bool, error) {
+	if addr.Protocol() != address.ID {
+		return nil, false, xerrors.Errorf("non-ID address %v invalid as actor key", addr)
+	}
+	if addr.Protocol() != address.ID {
+		return nil, false, xerrors.Errorf("non-ID address %v invalid as actor key", addr)
+	}
+	return t.Map.Find(t.Store.Context(), abi.AddrKey(addr).Key())
+}
+
+func (t *actorTree) SetActorV5(addr address.Address, actor *ActorV5) error {
+	if addr.Protocol() != address.ID {
+		return xerrors.Errorf("non-ID address %v invalid as actor key", addr)
+	}
+	return t.Map.Set(t.Store.Context(), abi.AddrKey(addr).Key(), actor)
+}
+
+func (t *actorTree) ForEachV5(fn func(addr address.Address, actor *ActorV5) error) error {
+	return t.Map.ForEach(t.Store.Context(), func(key string, val *ActorV5) error {
+		addr, err := address.NewFromBytes([]byte(key))
+		if err != nil {
+			return err
+		}
+		return fn(addr, val)
+	})
+}
+
+// Traverses all keys in the tree, without decoding the actor states.
+func (t *actorTree) ForEachKey(fn func(addr address.Address) error) error {
+	return t.Map.ForEach(t.Store.Context(), func(key string, _ *ActorV5) error {
+		addr, err := address.NewFromBytes([]byte(key))
+		if err != nil {
+			return err
+		}
+		return fn(addr)
+	})
+}
+
+// A specialization of a map of ID-addresses to actor heads.
+type legacyActorTree struct {
 	Map   *adt.Map
 	Store adt.Store
 }
 
 // Initializes a new, empty state tree backed by a store.
-func NewTree(store adt.Store) (ActorTree, error) {
+func NewLegacyTree(store adt.Store) (LegacyActorTree, error) {
 	emptyMap, err := adt.MakeEmptyMap(store, DefaultHamtBitwidth)
 	if err != nil {
 		return nil, err
 	}
-	return &actorTree{
+	return &legacyActorTree{
 		Map:   emptyMap,
 		Store: store,
 	}, nil
 }
 
 // Loads a tree from a root CID and store.
-func LoadTree(s adt.Store, r cid.Cid) (ActorTree, error) {
+func LoadLegacyTree(s adt.Store, r cid.Cid) (LegacyActorTree, error) {
 	m, err := adt.AsMap(s, r, DefaultHamtBitwidth)
 	if err != nil {
 		return nil, err
 	}
-	return &actorTree{
+	return &legacyActorTree{
 		Map:   m,
 		Store: s,
 	}, nil
 }
 
-func (t *actorTree) GetStore() adt.Store {
+func (t *legacyActorTree) GetStore() adt.Store {
 	return t.Store
 }
 
+func (t *legacyActorTree) GetMap() *adt.Map {
+	return t.Map
+}
+
 // Writes the tree root node to the store, and returns its CID.
-func (t *actorTree) Flush() (cid.Cid, error) {
+func (t *legacyActorTree) Flush() (cid.Cid, error) {
 	return t.Map.Root()
 }
 
 // Loads the state associated with an address.
-func (t *actorTree) GetActorV4(addr address.Address) (*ActorV4, bool, error) {
+func (t *legacyActorTree) GetActorV4(addr address.Address) (*ActorV4, bool, error) {
 	if addr.Protocol() != address.ID {
 		return nil, false, xerrors.Errorf("non-ID address %v invalid as actor key", addr)
 	}
@@ -92,7 +233,7 @@ func (t *actorTree) GetActorV4(addr address.Address) (*ActorV4, bool, error) {
 	return &actor, found, err
 }
 
-func (t *actorTree) GetActorV5(addr address.Address) (*ActorV5, bool, error) {
+func (t *legacyActorTree) GetActorV5(addr address.Address) (*ActorV5, bool, error) {
 	if addr.Protocol() != address.ID {
 		return nil, false, xerrors.Errorf("non-ID address %v invalid as actor key", addr)
 	}
@@ -102,22 +243,22 @@ func (t *actorTree) GetActorV5(addr address.Address) (*ActorV5, bool, error) {
 }
 
 // Sets the state associated with an address, overwriting if it already present.
-func (t *actorTree) SetActorV4(addr address.Address, actor *ActorV4) error {
+func (t *legacyActorTree) SetActorV4(addr address.Address, actor *ActorV4) error {
 	if addr.Protocol() != address.ID {
 		return xerrors.Errorf("non-ID address %v invalid as actor key", addr)
 	}
 	return t.Map.Put(abi.AddrKey(addr), actor)
 }
 
-func (t *actorTree) SetActorV5(addr address.Address, actor *ActorV5) error {
+func (t *legacyActorTree) SetActorV5(addr address.Address, actor *ActorV5) error {
 	if addr.Protocol() != address.ID {
 		return xerrors.Errorf("non-ID address %v invalid as actor key", addr)
 	}
 	return t.Map.Put(abi.AddrKey(addr), actor)
 }
 
 // Traverses all entries in the tree.
-func (t *actorTree) ForEachV4(fn func(addr address.Address, actor *ActorV4) error) error {
+func (t *legacyActorTree) ForEachV4(fn func(addr address.Address, actor *ActorV4) error) error {
 	var val ActorV4
 	return t.Map.ForEach(&val, func(key string) error {
 		addr, err := address.NewFromBytes([]byte(key))
@@ -128,7 +269,7 @@ func (t *actorTree) ForEachV4(fn func(addr address.Address, actor *ActorV4) erro
 	})
 }
 
-func (t *actorTree) ForEachV5(fn func(addr address.Address, actor *ActorV5) error) error {
+func (t *legacyActorTree) ForEachV5(fn func(addr address.Address, actor *ActorV5) error) error {
 	var val ActorV5
 	return t.Map.ForEach(&val, func(key string) error {
 		addr, err := address.NewFromBytes([]byte(key))
@@ -140,7 +281,7 @@ func (t *actorTree) ForEachV5(fn func(addr address.Address, actor *ActorV5) erro
 }
 
 // Traverses all keys in the tree, without decoding the actor states.
-func (t *actorTree) ForEachKey(fn func(addr address.Address) error) error {
+func (t *legacyActorTree) ForEachKey(fn func(addr address.Address) error) error {
 	return t.Map.ForEach(nil, func(key string) error {
 		addr, err := address.NewFromBytes([]byte(key))
 		if err != nil {

builtin/v10/migration/top.go

@@ -33,11 +33,11 @@ func MigrateStateTree(ctx context.Context, store cbor.IpldStore, newManifestCID
 	adtStore := adt9.WrapStore(ctx, store)
 
 	// Load input and output state trees
-	actorsIn, err := builtin.LoadTree(adtStore, actorsRootIn)
+	actorsIn, err := builtin.LoadLegacyTree(adtStore, actorsRootIn)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("loading state tree: %w", err)
 	}
-	actorsOut, err := builtin.NewTree(adtStore)
+	actorsOut, err := builtin.NewLegacyTree(adtStore)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("creating new state tree: %w", err)
 	}

builtin/v11/migration/top.go

@@ -26,7 +26,7 @@ func MigrateStateTree(ctx context.Context, store cbor.IpldStore, newManifestCID
 	adtStore := adt10.WrapStore(ctx, store)
 
 	// Load input and output state trees
-	actorsIn, err := builtin.LoadTree(adtStore, actorsRootIn)
+	actorsIn, err := builtin.LoadLegacyTree(adtStore, actorsRootIn)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("loading state tree: %w", err)
 	}

builtin/v12/migration/top.go

@@ -25,7 +25,7 @@ func MigrateStateTree(ctx context.Context, store cbor.IpldStore, newManifestCID
 	adtStore := adt11.WrapStore(ctx, store)
 
 	// Load input and output state trees
-	actorsIn, err := builtin.LoadTree(adtStore, actorsRootIn)
+	actorsIn, err := builtin.LoadLegacyTree(adtStore, actorsRootIn)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("loading state tree: %w", err)
 	}

builtin/v13/migration/top.go

@@ -27,7 +27,7 @@ func MigrateStateTree(ctx context.Context, store cbor.IpldStore, newManifestCID
 	adtStore := adt13.WrapStore(ctx, store)
 
 	// Load input and output state trees
-	actorsIn, err := builtin.LoadTree(adtStore, actorsRootIn)
+	actorsIn, err := builtin.LoadLegacyTree(adtStore, actorsRootIn)
 	if err != nil {
 		return cid.Undef, xerrors.Errorf("loading state tree: %w", err)
 	}

builtin/v8/check.go

@@ -28,7 +28,7 @@ import (
 // Within this code, Go errors are not expected, but are often converted to messages so that execution
 // can continue to find more errors rather than fail with no insight.
 // Only errors thar are particularly troublesome to recover from should propagate as Go errors.
-func CheckStateInvariants(tree builtin.ActorTree, priorEpoch abi.ChainEpoch, actorCodes map[string]cid.Cid) (*builtin.MessageAccumulator, error) {
+func CheckStateInvariants(tree builtin.LegacyActorTree, priorEpoch abi.ChainEpoch, actorCodes map[string]cid.Cid) (*builtin.MessageAccumulator, error) {
 	acc := &builtin.MessageAccumulator{}
 	totalFIl := big.Zero()
 	var initSummary *init_.StateSummary

builtin/v9/check.go

@@ -30,7 +30,7 @@ import (
 // Within this code, Go errors are not expected, but are often converted to messages so that execution
 // can continue to find more errors rather than fail with no insight.
 // Only errors thar are particularly troublesome to recover from should propagate as Go errors.
-func CheckStateInvariants(tree builtin.ActorTree, priorEpoch abi.ChainEpoch, actorCodes map[string]cid.Cid) (*builtin.MessageAccumulator, error) {
+func CheckStateInvariants(tree builtin.LegacyActorTree, priorEpoch abi.ChainEpoch, actorCodes map[string]cid.Cid) (*builtin.MessageAccumulator, error) {
 	acc := &builtin.MessageAccumulator{}
 	totalFIl := big.Zero()
 	var initSummary *init_.StateSummary

builtin/v9/migration/test/migration_test.go

@@ -26,7 +26,7 @@ func TestMigration(t *testing.T) {
 
 	startRoot := makeInputTree(ctx, t, adtStore)
 
-	oldStateTree, err := builtin.LoadTree(adtStore, startRoot)
+	oldStateTree, err := builtin.LoadLegacyTree(adtStore, startRoot)
 	require.NoError(t, err)
 
 	oldSystemActor, found, err := oldStateTree.GetActorV4(builtin.SystemActorAddr)
@@ -60,7 +60,7 @@ func TestMigration(t *testing.T) {
 
 	// check that the system actor state was correctly updated
 
-	newStateTree, err := builtin.LoadTree(adtStore, cacheRoot)
+	newStateTree, err := builtin.LoadLegacyTree(adtStore, cacheRoot)
 	require.NoError(t, err)
 
 	newSystemActor, found, err := newStateTree.GetActorV4(builtin.SystemActorAddr)

builtin/v9/migration/test/miner_test.go

@@ -32,7 +32,7 @@ func TestMinerMigration(t *testing.T) {
 
 	startRoot := makeInputTree(ctx, t, adtStore)
 
-	oldStateTree, err := builtin.LoadTree(adtStore, startRoot)
+	oldStateTree, err := builtin.LoadLegacyTree(adtStore, startRoot)
 	require.NoError(t, err)
 
 	oldSystemActor, found, err := oldStateTree.GetActorV4(builtin.SystemActorAddr)
@@ -237,7 +237,7 @@ func TestMinerMigration(t *testing.T) {
 
 	// check that the actor states were correctly updated
 
-	newStateTree, err := builtin.LoadTree(adtStore, cacheRoot)
+	newStateTree, err := builtin.LoadLegacyTree(adtStore, cacheRoot)
 	require.NoError(t, err)
 
 	// miner 1 is just empty precommits
@@ -324,7 +324,7 @@ func TestFip0029MinerMigration(t *testing.T) {
 
 	startRoot := makeInputTree(ctx, t, adtStore)
 
-	oldStateTree, err := builtin.LoadTree(adtStore, startRoot)
+	oldStateTree, err := builtin.LoadLegacyTree(adtStore, startRoot)
 	require.NoError(t, err)
 
 	oldSystemActor, found, err := oldStateTree.GetActorV4(builtin.SystemActorAddr)
@@ -386,7 +386,7 @@ func TestFip0029MinerMigration(t *testing.T) {
 
 	// check that the actor states were correctly updated
 
-	newStateTree, err := builtin.LoadTree(adtStore, cacheRoot)
+	newStateTree, err := builtin.LoadLegacyTree(adtStore, cacheRoot)
 	require.NoError(t, err)
 
 	newMinerActor, ok, err := newStateTree.GetActorV4(addr)