From 85936de254a55920fb501c90af677d265a0d7e99 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 2 Aug 2021 22:38:41 +0300
Subject: [PATCH] vm: don't create reference counter when it's not needed

 * invocation stack doesn't need refcounting
 * exception stack doesn't need refcounting
 * evaluation stack always has VM-level refcounter
---
 pkg/vm/ref_counter.go |  6 ++++++
 pkg/vm/stack.go       |  7 +++++--
 pkg/vm/vm.go          | 17 +++++------------
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/pkg/vm/ref_counter.go b/pkg/vm/ref_counter.go
index 0d176cadf..182177f6c 100644
--- a/pkg/vm/ref_counter.go
+++ b/pkg/vm/ref_counter.go
@@ -18,6 +18,9 @@ func newRefCounter() *refCounter {
 
 // Add adds an item to the reference counter.
 func (r *refCounter) Add(item stackitem.Item) {
+	if r == nil {
+		return
+	}
 	r.size++
 
 	switch item.(type) {
@@ -41,6 +44,9 @@ func (r *refCounter) Add(item stackitem.Item) {
 
 // Remove removes item from the reference counter.
 func (r *refCounter) Remove(item stackitem.Item) {
+	if r == nil {
+		return
+	}
 	r.size--
 
 	switch item.(type) {
diff --git a/pkg/vm/stack.go b/pkg/vm/stack.go
index 1b1130f53..fca28215b 100644
--- a/pkg/vm/stack.go
+++ b/pkg/vm/stack.go
@@ -158,13 +158,16 @@ type Stack struct {
 
 // NewStack returns a new stack name by the given name.
 func NewStack(n string) *Stack {
+	return newStack(n, newRefCounter())
+}
+
+func newStack(n string, refc *refCounter) *Stack {
 	s := &Stack{
 		name: n,
+		refs: refc,
 	}
 	s.top.next = &s.top
 	s.top.prev = &s.top
-	s.len = 0
-	s.refs = newRefCounter()
 	return s
 }
 
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 1ff7176e0..be47455a7 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -99,7 +99,7 @@ func New() *VM {
 func NewWithTrigger(t trigger.Type) *VM {
 	vm := &VM{
 		state:   NoneState,
-		istack:  NewStack("invocation"),
+		istack:  newStack("invocation", nil),
 		refs:    newRefCounter(),
 		trigger: t,
 
@@ -107,17 +107,10 @@ func NewWithTrigger(t trigger.Type) *VM {
 		Invocations:    make(map[util.Uint160]int),
 	}
 
-	vm.estack = vm.newItemStack("evaluation")
+	vm.estack = newStack("evaluation", vm.refs)
 	return vm
 }
 
-func (v *VM) newItemStack(n string) *Stack {
-	s := NewStack(n)
-	s.refs = v.refs
-
-	return s
-}
-
 // SetPriceGetter registers the given PriceGetterFunc in v.
 // f accepts vm's Context, current instruction and instruction parameter.
 func (v *VM) SetPriceGetter(f func(opcode.Opcode, []byte) int64) {
@@ -288,9 +281,9 @@ func (v *VM) LoadScript(b []byte) {
 func (v *VM) LoadScriptWithFlags(b []byte, f callflag.CallFlag) {
 	v.checkInvocationStackSize()
 	ctx := NewContextWithParams(b, 0, -1, 0)
-	v.estack = v.newItemStack("estack")
+	v.estack = newStack("evaluation", v.refs)
 	ctx.estack = v.estack
-	ctx.tryStack = NewStack("exception")
+	ctx.tryStack = newStack("exception", nil)
 	ctx.callFlag = f
 	ctx.static = newSlot(v.refs)
 	ctx.callingScriptHash = v.GetCurrentScriptHash()
@@ -1527,7 +1520,7 @@ func (v *VM) call(ctx *Context, offset int) {
 	newCtx.RetCount = -1
 	newCtx.local = nil
 	newCtx.arguments = nil
-	newCtx.tryStack = NewStack("exception")
+	newCtx.tryStack = newStack("exception", nil)
 	newCtx.NEF = ctx.NEF
 	v.istack.PushVal(newCtx)
 	v.Jump(newCtx, offset)