Add test

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>

chain/state/statetree.go

@@ -30,6 +30,7 @@ type StateTree struct {
 	version     types.StateTreeVersion
 	info        cid.Cid
 	Store       cbor.IpldStore
+	lookupIDFun func(address.Address) (address.Address, error)
 
 	snaps *stateSnaps
 }
@@ -173,13 +174,15 @@ func NewStateTree(cst cbor.IpldStore, ver types.StateTreeVersion) (*StateTree, e
 		return nil, err
 	}
 
-	return &StateTree{
+	s := &StateTree{
 		root:    root,
 		info:    info,
 		version: ver,
 		Store:   cst,
 		snaps:   newStateSnaps(),
-	}, nil
+	}
+	s.lookupIDFun = s.lookupIDinternal
+	return s, nil
 }
 
 func LoadStateTree(cst cbor.IpldStore, c cid.Cid) (*StateTree, error) {
@@ -203,13 +206,15 @@ func LoadStateTree(cst cbor.IpldStore, c cid.Cid) (*StateTree, error) {
 			return nil, err
 		}
 
-		return &StateTree{
+		s := &StateTree{
 			root:    nd,
 			info:    root.Info,
 			version: root.Version,
 			Store:   cst,
 			snaps:   newStateSnaps(),
-		}, nil
+		}
+		s.lookupIDFun = s.lookupIDinternal
+		return s, nil
 	default:
 		return nil, xerrors.Errorf("unsupported state tree version: %d", root.Version)
 	}
@@ -226,17 +231,7 @@ func (st *StateTree) SetActor(addr address.Address, act *types.Actor) error {
 	return nil
 }
 
-// LookupID gets the ID address of this actor's `addr` stored in the `InitActor`.
+func (st *StateTree) lookupIDinternal(addr address.Address) (address.Address, error) {
-func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
-	if addr.Protocol() == address.ID {
-		return addr, nil
-	}
-
-	resa, ok := st.snaps.resolveAddress(addr)
-	if ok {
-		return resa, nil
-	}
-
 	act, err := st.GetActor(init_.Address)
 	if err != nil {
 		return address.Undef, xerrors.Errorf("getting init actor: %w", err)
@@ -254,6 +249,23 @@ func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
 	if err != nil {
 		return address.Undef, xerrors.Errorf("resolve address %s: %w", addr, err)
 	}
+	return a, err
+}
+
+// LookupID gets the ID address of this actor's `addr` stored in the `InitActor`.
+func (st *StateTree) LookupID(addr address.Address) (address.Address, error) {
+	if addr.Protocol() == address.ID {
+		return addr, nil
+	}
+
+	resa, ok := st.snaps.resolveAddress(addr)
+	if ok {
+		return resa, nil
+	}
+	a, err := st.lookupIDFun(addr)
+	if err != nil {
+		return a, err
+	}
 
 	st.snaps.cacheResolveAddress(addr, a)
 

chain/state/statetree_test.go

@@ -73,6 +73,104 @@ func BenchmarkStateTreeSetFlush(b *testing.B) {
 	}
 }
 
+func TestResolveCache(t *testing.T) {
+	cst := cbor.NewMemCborStore()
+	st, err := NewStateTree(cst, VersionForNetwork(build.NewestNetworkVersion))
+	if err != nil {
+		t.Fatal(err)
+	}
+	nonId := address.NewForTestGetter()()
+	id, _ := address.NewIDAddress(1000)
+
+	st.lookupIDFun = func(a address.Address) (address.Address, error) {
+		if a == nonId {
+			return id, nil
+		} else {
+			return address.Undef, types.ErrActorNotFound
+		}
+	}
+
+	err = st.SetActor(nonId, &types.Actor{Nonce: 1})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	{
+		err = st.Snapshot(context.TODO())
+		if err != nil {
+			t.Fatal(err)
+		}
+		act, err := st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 1 {
+			t.Fatalf("expected nonce 1, got %d", act.Nonce)
+		}
+		err = st.SetActor(nonId, &types.Actor{Nonce: 2})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		act, err = st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 2 {
+			t.Fatalf("expected nonce 2, got %d", act.Nonce)
+		}
+
+		if err := st.Revert(); err != nil {
+			t.Fatal(err)
+		}
+		st.ClearSnapshot()
+	}
+
+	act, err := st.GetActor(nonId)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if act.Nonce != 1 {
+		t.Fatalf("expected nonce 1, got %d", act.Nonce)
+	}
+
+	{
+		err = st.Snapshot(context.TODO())
+		if err != nil {
+			t.Fatal(err)
+		}
+		act, err := st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 1 {
+			t.Fatalf("expected nonce 1, got %d", act.Nonce)
+		}
+		err = st.SetActor(nonId, &types.Actor{Nonce: 2})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		act, err = st.GetActor(nonId)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if act.Nonce != 2 {
+			t.Fatalf("expected nonce 2, got %d", act.Nonce)
+		}
+		st.ClearSnapshot()
+	}
+
+	act, err = st.GetActor(nonId)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if act.Nonce != 2 {
+		t.Fatalf("expected nonce 2, got %d", act.Nonce)
+	}
+
+}
+
 func BenchmarkStateTree10kGetActor(b *testing.B) {
 	cst := cbor.NewMemCborStore()
 	st, err := NewStateTree(cst, VersionForNetwork(build.NewestNetworkVersion))