From f79f62dab44daecbef29f031418893aedb637947 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Mon, 23 May 2022 12:02:37 +0300
Subject: [PATCH] vm: use single function instead of committer/discarder

---
 pkg/core/interop/context.go | 30 ++++++++++++++----------------
 pkg/vm/vm.go                | 27 ++++++++-------------------
 2 files changed, 22 insertions(+), 35 deletions(-)

diff --git a/pkg/core/interop/context.go b/pkg/core/interop/context.go
index 95283f6e0..3ff7f063c 100644
--- a/pkg/core/interop/context.go
+++ b/pkg/core/interop/context.go
@@ -324,29 +324,27 @@ func (ic *Context) SpawnVM() *vm.VM {
 		}
 		ic.DAO = ic.DAO.GetPrivate()
 	}
-	committer := func() error {
+	unwrapper := func(commit bool, ntfToRemove int) error {
+		if !commit {
+			have := len(ic.Notifications)
+			if have < ntfToRemove {
+				panic(fmt.Errorf("inconsistent notifications count: should remove %d, have %d", ntfToRemove, len(ic.Notifications)))
+			}
+			ic.Notifications = ic.Notifications[:have-ntfToRemove]
+		}
 		if ic.DAO == nil {
 			return nil
 		}
-		_, err := ic.DAO.Persist()
-		if err != nil {
-			return fmt.Errorf("failed to persist changes %w", err)
+		if commit {
+			_, err := ic.DAO.Persist()
+			if err != nil {
+				return fmt.Errorf("failed to persist changes %w", err)
+			}
 		}
 		ic.DAO = ic.DAO.GetUnwrapped()
 		return nil
 	}
-	reverter := func(ntfToRemove int) {
-		have := len(ic.Notifications)
-		if have < ntfToRemove {
-			panic(fmt.Errorf("inconsistent notifications count: should remove %d, have %d", ntfToRemove, len(ic.Notifications)))
-		}
-		ic.Notifications = ic.Notifications[:have-ntfToRemove]
-		if ic.DAO == nil {
-			return
-		}
-		ic.DAO = ic.DAO.GetUnwrapped() // Discard all changes made in this layer.
-	}
-	v.SetIsolationCallbacks(wrapper, committer, reverter)
+	v.SetIsolationCallbacks(wrapper, unwrapper)
 	ic.VM = v
 	return v
 }
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 11f72ca77..4733d777d 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -69,10 +69,8 @@ type VM struct {
 
 	// wraps DAO with private MemCachedStore
 	wrapDao func()
-	// commits DAO changes and unwraps DAO.
-	commitChanges func() error
-	// unwraps DAO and removes last notificationsCount notifications from the context
-	revertChanges func(notificationsCount int)
+	// either commits or discards changes made in the current context; performs DAO unwrapping.
+	unwrapDAO func(commit bool, notificationsCount int) error
 
 	istack Stack  // invocation stack.
 	estack *Stack // execution stack.
@@ -136,10 +134,9 @@ func (v *VM) EmitNotification() {
 // wrapper performs DAO cloning;
 // committer persists changes made in the upper snapshot to the underlying DAO;
 // reverter rolls back the whole set of changes made in the current snapshot.
-func (v *VM) SetIsolationCallbacks(wrapper func(), committer func() error, reverter func(ntfToRemove int)) {
+func (v *VM) SetIsolationCallbacks(wrapper func(), unwrapper func(commit bool, notificationsCount int) error) {
 	v.wrapDao = wrapper
-	v.commitChanges = committer
-	v.revertChanges = reverter
+	v.unwrapDAO = unwrapper
 }
 
 // SetPriceGetter registers the given PriceGetterFunc in v.
@@ -1635,18 +1632,10 @@ func (v *VM) unloadContext(ctx *Context) {
 	if ctx.static != nil && (currCtx == nil || ctx.static != currCtx.static) {
 		ctx.static.ClearRefs(&v.refs)
 	}
-	if ctx.isWrapped { // In case of CALL, CALLA, CALLL we don't need to commit/discard changes, unwrap DAO and change notificationsCount.
-		if v.uncaughtException == nil {
-			if v.commitChanges != nil {
-				if err := v.commitChanges(); err != nil {
-					// TODO: return an error instead?
-					panic(fmt.Errorf("failed to commit changes: %w", err))
-				}
-			}
-		} else {
-			if v.revertChanges != nil {
-				v.revertChanges(*ctx.notificationsCount)
-			}
+	if ctx.isWrapped && v.unwrapDAO != nil { // In case of CALL, CALLA, CALLL we don't need to commit/discard changes, unwrap DAO and change notificationsCount.
+		err := v.unwrapDAO(v.uncaughtException == nil, *ctx.notificationsCount)
+		if err != nil {
+			panic(fmt.Errorf("failed to unwrap DAO: %w", err))
 		}
 	}
 	if currCtx != nil && ctx.persistNotificationsCountOnUnloading && !(ctx.isWrapped && v.uncaughtException != nil) {