From 6ecd1ae437f94860bca9282ab6b5a186dc69c882 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 24 Jul 2020 11:54:12 +0300
Subject: [PATCH] vm: allow to initialize static slot in method

---
 pkg/vm/json_test.go |  2 +-
 pkg/vm/slot.go      | 30 ++++++++++++++++++++++--------
 pkg/vm/slot_test.go |  5 ++++-
 pkg/vm/vm.go        |  6 ++----
 pkg/vm/vm_test.go   |  8 ++++++++
 5 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/pkg/vm/json_test.go b/pkg/vm/json_test.go
index 7a915a64a..e97328362 100644
--- a/pkg/vm/json_test.go
+++ b/pkg/vm/json_test.go
@@ -228,7 +228,7 @@ func compareStacks(t *testing.T, expected []vmUTStackItem, actual *Stack) {
 }
 
 func compareSlots(t *testing.T, expected []vmUTStackItem, actual *Slot) {
-	if actual == nil && len(expected) == 0 {
+	if actual.storage == nil && len(expected) == 0 {
 		return
 	}
 	require.NotNil(t, actual)
diff --git a/pkg/vm/slot.go b/pkg/vm/slot.go
index 18a65ac85..22e4d8d5c 100644
--- a/pkg/vm/slot.go
+++ b/pkg/vm/slot.go
@@ -10,16 +10,25 @@ type Slot struct {
 	refs    *refCounter
 }
 
-// newSlot returns new slot of n items.
-func newSlot(n int, refs *refCounter) *Slot {
+// newSlot returns new slot with the provided reference counter.
+func newSlot(refs *refCounter) *Slot {
 	return &Slot{
-		storage: make([]stackitem.Item, n),
-		refs:    refs,
+		refs: refs,
 	}
 }
 
+// init sets static slot size to n. It is intended to be used only by INITSSLOT.
+func (s *Slot) init(n int) {
+	if s.storage != nil {
+		panic("already initialized")
+	}
+	s.storage = make([]stackitem.Item, n)
+}
+
 func (v *VM) newSlot(n int) *Slot {
-	return newSlot(n, v.refs)
+	s := newSlot(v.refs)
+	s.init(n)
+	return s
 }
 
 // Set sets i-th storage slot.
@@ -43,12 +52,17 @@ func (s *Slot) Get(i int) stackitem.Item {
 	return stackitem.Null{}
 }
 
-// Size returns slot size.
-func (s *Slot) Size() int { return len(s.storage) }
-
 // Clear removes all slot variables from reference counter.
 func (s *Slot) Clear() {
 	for _, item := range s.storage {
 		s.refs.Remove(item)
 	}
 }
+
+// Size returns slot size.
+func (s *Slot) Size() int {
+	if s.storage == nil {
+		panic("not initialized")
+	}
+	return len(s.storage)
+}
diff --git a/pkg/vm/slot_test.go b/pkg/vm/slot_test.go
index 19186cdee..434464476 100644
--- a/pkg/vm/slot_test.go
+++ b/pkg/vm/slot_test.go
@@ -9,8 +9,11 @@ import (
 )
 
 func TestSlot_Get(t *testing.T) {
-	s := newSlot(3, newRefCounter())
+	s := newSlot(newRefCounter())
 	require.NotNil(t, s)
+	require.Panics(t, func() { s.Size() })
+
+	s.init(3)
 	require.Equal(t, 3, s.Size())
 
 	// Null is the default
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index d5ef55f7f..558d8ccc4 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -282,6 +282,7 @@ func (v *VM) LoadScriptWithFlags(b []byte, f smartcontract.CallFlag) {
 	ctx.estack = v.estack
 	ctx.tryStack = NewStack("exception")
 	ctx.callFlag = f
+	ctx.static = newSlot(v.refs)
 	v.istack.PushVal(ctx)
 }
 
@@ -568,13 +569,10 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 		v.estack.PushVal(result)
 
 	case opcode.INITSSLOT:
-		if ctx.static != nil {
-			panic("already initialized")
-		}
 		if parameter[0] == 0 {
 			panic("zero argument")
 		}
-		ctx.static = v.newSlot(int(parameter[0]))
+		ctx.static.init(int(parameter[0]))
 
 	case opcode.INITSLOT:
 		if ctx.local != nil || ctx.arguments != nil {
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 3eb653283..f4a0a94ec 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -2596,6 +2596,14 @@ func TestSLOTOpcodes(t *testing.T) {
 		t.Run("Local", getTestFuncForVM(makeProgram(opcode.INITSLOT, 8, 0, opcode.STLOC, 7, opcode.LDLOC, 7), 42, 42))
 		t.Run("Argument", getTestFuncForVM(makeProgram(opcode.INITSLOT, 0, 2, opcode.STARG, 1, opcode.LDARG, 1), 42, 42, 1, 2))
 	})
+
+	t.Run("InitStaticSlotInMethod", func(t *testing.T) {
+		prog := makeProgram(
+			opcode.CALL, 4, opcode.LDSFLD0, opcode.RET,
+			opcode.INITSSLOT, 1, opcode.PUSH12, opcode.STSFLD0, opcode.RET,
+		)
+		runWithArgs(t, prog, 12)
+	})
 }
 
 func makeProgram(opcodes ...opcode.Opcode) []byte {