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)