From ce226f6b76277a6718678d2ebdbe7ca8ce38fb80 Mon Sep 17 00:00:00 2001
From: Anna Shaleva <shaleva.ann@nspcc.ru>
Date: Mon, 23 May 2022 08:42:10 +0300
Subject: [PATCH] vm: optimize context wrapping code

We can omit DAO wrapping for safe methods and for those methods that are not
wrapped into try-catch block. However, we still need to persist
notificationsCount changes for these methods to the parent context.
---
 pkg/vm/context.go |  3 +++
 pkg/vm/vm.go      | 28 +++++++++++++++++++++-------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/pkg/vm/context.go b/pkg/vm/context.go
index be859a4de..ebf392c4e 100644
--- a/pkg/vm/context.go
+++ b/pkg/vm/context.go
@@ -57,6 +57,9 @@ type Context struct {
 	// notificationsCount stores number of notifications emitted during current context
 	// handling.
 	notificationsCount *int
+	// persistNotificationsCountOnUnloading denotes whether notificationsCount should be
+	// persisted to the parent context on current context unloading.
+	persistNotificationsCountOnUnloading bool
 	// isWrapped tells whether the context's DAO was wrapped into another layer of
 	// MemCachedStore on creation and whether it should be unwrapped on context unloading.
 	isWrapped bool
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 74c3f0de5..36b289bd0 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -357,9 +357,9 @@ func (v *VM) loadScriptWithCallingHash(b []byte, exe *nef.File, caller util.Uint
 	ctx.scriptHash = hash
 	ctx.callingScriptHash = caller
 	ctx.NEF = exe
+	parent := v.Context()
 	if v.invTree != nil {
 		curTree := v.invTree
-		parent := v.Context()
 		if parent != nil {
 			curTree = parent.invTree
 		}
@@ -368,9 +368,22 @@ func (v *VM) loadScriptWithCallingHash(b []byte, exe *nef.File, caller util.Uint
 		ctx.invTree = newTree
 	}
 	if v.wrapDao != nil {
-		v.wrapDao()
-		ctx.isWrapped = true
+		needWrap := f&(callflag.All^callflag.ReadOnly) != 0 // If the method is safe, then it's read-only and doesn't perform storage changes or emit notifications.
+		if !needWrap && parent != nil {                     // If the method is not wrapped into try-catch block, then changes should be discarded anyway if exception occurs.
+			for i := 0; i < parent.tryStack.Len(); i++ {
+				eCtx := parent.tryStack.Peek(i).Value().(*exceptionHandlingContext)
+				if eCtx.State == eTry {
+					needWrap = true // TODO: is it correct to wrap it only once and break after the first occurrence?
+					break
+				}
+			}
+		}
+		if needWrap {
+			v.wrapDao()
+			ctx.isWrapped = true
+		}
 	}
+	ctx.persistNotificationsCountOnUnloading = true
 	v.istack.PushItem(ctx)
 }
 
@@ -1621,7 +1634,7 @@ func (v *VM) unloadContext(ctx *Context) {
 	if ctx.static != nil && (currCtx == nil || ctx.static != currCtx.static) {
 		ctx.static.ClearRefs(&v.refs)
 	}
-	if currCtx == nil || ctx.isWrapped { // In case of CALL, CALLA, CALLL we don't need to commit/discard changes, unwrap DAO and change notificationsCount.
+	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 {
@@ -1629,15 +1642,15 @@ func (v *VM) unloadContext(ctx *Context) {
 					panic(fmt.Errorf("failed to commit changes: %w", err))
 				}
 			}
-			if currCtx != nil {
-				*currCtx.notificationsCount += *ctx.notificationsCount
-			}
 		} else {
 			if v.revertChanges != nil {
 				v.revertChanges(*ctx.notificationsCount)
 			}
 		}
 	}
+	if currCtx != nil && ctx.persistNotificationsCountOnUnloading && !(ctx.isWrapped && v.uncaughtException != nil) {
+		*currCtx.notificationsCount += *ctx.notificationsCount
+	}
 }
 
 // getTryParams splits TRY(L) instruction parameter into offsets for catch and finally blocks.
@@ -1699,6 +1712,7 @@ func (v *VM) call(ctx *Context, offset int) {
 	// unloadContext without unnecessary DAO unwrapping and notificationsCount changes.
 	newCtx.notificationsCount = ctx.notificationsCount
 	newCtx.isWrapped = false
+	newCtx.persistNotificationsCountOnUnloading = false
 	v.istack.PushItem(newCtx)
 	newCtx.Jump(offset)
 }