cache address resolutions in state tree

chain/state/statetree.go

@@ -30,7 +30,19 @@ type StateTree struct {
 }
 
 type stateSnaps struct {
-	layers []map[address.Address]streeOp
+	layers []*stateSnapLayer
+}
+
+type stateSnapLayer struct {
+	actors       map[address.Address]streeOp
+	resolveCache map[address.Address]address.Address
+}
+
+func newStateSnapLayer() *stateSnapLayer {
+	return &stateSnapLayer{
+		actors:       make(map[address.Address]streeOp),
+		resolveCache: make(map[address.Address]address.Address),
+	}
 }
 
 type streeOp struct {
@@ -39,13 +51,13 @@ type streeOp struct {
 }
 
 func newStateSnaps() *stateSnaps {
-	return &stateSnaps{
+	ss := &stateSnaps{}
-		layers: []map[address.Address]streeOp{make(map[address.Address]streeOp)},
+	ss.addLayer()
-	}
+	return ss
 }
 
 func (ss *stateSnaps) addLayer() {
-	ss.layers = append(ss.layers, make(map[address.Address]streeOp))
+	ss.layers = append(ss.layers, newStateSnapLayer())
 }
 
 func (ss *stateSnaps) dropLayer() {
@@ -57,16 +69,34 @@ func (ss *stateSnaps) mergeLastLayer() {
 	last := ss.layers[len(ss.layers)-1]
 	nextLast := ss.layers[len(ss.layers)-2]
 
-	for k, v := range last {
+	for k, v := range last.actors {
-		nextLast[k] = v
+		nextLast.actors[k] = v
+	}
+
+	for k, v := range last.resolveCache {
+		nextLast.resolveCache[k] = v
 	}
 
 	ss.dropLayer()
 }
 
+func (ss *stateSnaps) resolveAddress(addr address.Address) (address.Address, bool) {
+	for i := len(ss.layers) - 1; i >= 0; i-- {
+		resa, ok := ss.layers[i].resolveCache[addr]
+		if ok {
+			return resa, true
+		}
+	}
+	return address.Undef, false
+}
+
+func (ss *stateSnaps) cacheResolveAddress(addr, resa address.Address) {
+	ss.layers[len(ss.layers)-1].resolveCache[addr] = resa
+}
+
 func (ss *stateSnaps) getActor(addr address.Address) (*types.Actor, error) {
 	for i := len(ss.layers) - 1; i >= 0; i-- {
-		act, ok := ss.layers[i][addr]
+		act, ok := ss.layers[i].actors[addr]
 		if ok {
 			if act.Delete {
 				return nil, types.ErrActorNotFound
@@ -79,11 +109,11 @@ func (ss *stateSnaps) getActor(addr address.Address) (*types.Actor, error) {
 }
 
 func (ss *stateSnaps) setActor(addr address.Address, act *types.Actor) {
-	ss.layers[len(ss.layers)-1][addr] = streeOp{Act: *act}
+	ss.layers[len(ss.layers)-1].actors[addr] = streeOp{Act: *act}
 }
 
 func (ss *stateSnaps) deleteActor(addr address.Address) {
-	ss.layers[len(ss.layers)-1][addr] = streeOp{Delete: true}
+	ss.layers[len(ss.layers)-1].actors[addr] = streeOp{Delete: true}
 }
 
 func NewStateTree(cst cbor.IpldStore) (*StateTree, error) {
@@ -125,6 +155,11 @@ func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
 		return addr, nil
 	}
 
+	resa, ok := st.snaps.resolveAddress(addr)
+	if ok {
+		return resa, nil
+	}
+
 	act, err := st.GetActor(builtin.InitActorAddr)
 	if err != nil {
 		return address.Undef, xerrors.Errorf("getting init actor: %w", err)
@@ -139,6 +174,9 @@ func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
 	if err != nil {
 		return address.Undef, xerrors.Errorf("resolve address %s: %w", addr, err)
 	}
+
+	st.snaps.cacheResolveAddress(addr, a)
+
 	return a, nil
 }
 
@@ -213,7 +251,7 @@ func (st *StateTree) Flush(ctx context.Context) (cid.Cid, error) {
 		return cid.Undef, xerrors.Errorf("tried to flush state tree with snapshots on the stack")
 	}
 
-	for addr, sto := range st.snaps.layers[0] {
+	for addr, sto := range st.snaps.layers[0].actors {
 		if sto.Delete {
 			if err := st.root.Delete(ctx, string(addr.Bytes())); err != nil {
 				return cid.Undef, err