From be407332b9cf47977ccdba98674a2e9433176e3d Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Thu, 19 Mar 2020 18:21:56 +0300
Subject: [PATCH] vm: add TryInteger() to StackItem interface

Conversion should be done in a StackItem, not in an Element.
---
 pkg/vm/context.go    |  6 ++++++
 pkg/vm/stack.go      | 16 ++++------------
 pkg/vm/stack_item.go | 45 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/pkg/vm/context.go b/pkg/vm/context.go
index 041df6b8e..3d3717104 100644
--- a/pkg/vm/context.go
+++ b/pkg/vm/context.go
@@ -3,6 +3,7 @@ package vm
 import (
 	"encoding/binary"
 	"errors"
+	"math/big"
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
@@ -178,6 +179,11 @@ func (c *Context) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Context to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (c *Context) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Context to Integer")
+}
+
 // Equals implements StackItem interface.
 func (c *Context) Equals(s StackItem) bool {
 	return c == s
diff --git a/pkg/vm/stack.go b/pkg/vm/stack.go
index d0f833db0..6730c4456 100644
--- a/pkg/vm/stack.go
+++ b/pkg/vm/stack.go
@@ -7,7 +7,6 @@ import (
 	"math/big"
 
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract"
-	"github.com/nspcc-dev/neo-go/pkg/vm/emit"
 )
 
 // Stack implementation for the neo-go virtual machine. The stack implements
@@ -74,18 +73,11 @@ func (e *Element) Value() interface{} {
 // BigInt attempts to get the underlying value of the element as a big integer.
 // Will panic if the assertion failed which will be caught by the VM.
 func (e *Element) BigInt() *big.Int {
-	switch t := e.value.(type) {
-	case *BigIntegerItem:
-		return t.value
-	case *BoolItem:
-		if t.value {
-			return big.NewInt(1)
-		}
-		return big.NewInt(0)
-	default:
-		b := t.Value().([]uint8)
-		return emit.BytesToInt(b)
+	val, err := e.value.TryInteger()
+	if err != nil {
+		panic(err)
 	}
+	return val
 }
 
 // TryBool attempts to get the underlying value of the element as a boolean.
diff --git a/pkg/vm/stack_item.go b/pkg/vm/stack_item.go
index 3500281bd..93fbf9037 100644
--- a/pkg/vm/stack_item.go
+++ b/pkg/vm/stack_item.go
@@ -22,6 +22,8 @@ type StackItem interface {
 	Dup() StackItem
 	// TryBytes converts StackItem to a byte slice.
 	TryBytes() ([]byte, error)
+	// TryInteger converts StackItem to an integer.
+	TryInteger() (*big.Int, error)
 	// Equals checks if 2 StackItems are equal.
 	Equals(s StackItem) bool
 	// ToContractParameter converts StackItem to smartcontract.Parameter
@@ -134,6 +136,11 @@ func (i *StructItem) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Struct to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (i *StructItem) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Struct to Integer")
+}
+
 // Equals implements StackItem interface.
 func (i *StructItem) Equals(s StackItem) bool {
 	if i == s {
@@ -210,6 +217,11 @@ func (i NullItem) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Null to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (i NullItem) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Null to Integer")
+}
+
 // Equals implements StackItem interface.
 func (i NullItem) Equals(s StackItem) bool {
 	_, ok := s.(NullItem)
@@ -245,6 +257,11 @@ func (i *BigIntegerItem) TryBytes() ([]byte, error) {
 	return i.Bytes(), nil
 }
 
+// TryInteger implements StackItem interface.
+func (i *BigIntegerItem) TryInteger() (*big.Int, error) {
+	return i.value, nil
+}
+
 // Equals implements StackItem interface.
 func (i *BigIntegerItem) Equals(s StackItem) bool {
 	if i == s {
@@ -334,6 +351,14 @@ func (i *BoolItem) TryBytes() ([]byte, error) {
 	return i.Bytes(), nil
 }
 
+// TryInteger implements StackItem interface.
+func (i *BoolItem) TryInteger() (*big.Int, error) {
+	if i.value {
+		return big.NewInt(1), nil
+	}
+	return big.NewInt(0), nil
+}
+
 // Equals implements StackItem interface.
 func (i *BoolItem) Equals(s StackItem) bool {
 	if i == s {
@@ -388,6 +413,11 @@ func (i *ByteArrayItem) TryBytes() ([]byte, error) {
 	return i.value, nil
 }
 
+// TryInteger implements StackItem interface.
+func (i *ByteArrayItem) TryInteger() (*big.Int, error) {
+	return emit.BytesToInt(i.value), nil
+}
+
 // Equals implements StackItem interface.
 func (i *ByteArrayItem) Equals(s StackItem) bool {
 	if i == s {
@@ -445,6 +475,11 @@ func (i *ArrayItem) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Array to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (i *ArrayItem) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Array to Integer")
+}
+
 // Equals implements StackItem interface.
 func (i *ArrayItem) Equals(s StackItem) bool {
 	return i == s
@@ -505,6 +540,11 @@ func (i *MapItem) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Map to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (i *MapItem) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Map to Integer")
+}
+
 // Equals implements StackItem interface.
 func (i *MapItem) Equals(s StackItem) bool {
 	return i == s
@@ -616,6 +656,11 @@ func (i *InteropItem) TryBytes() ([]byte, error) {
 	return nil, errors.New("can't convert Interop to ByteArray")
 }
 
+// TryInteger implements StackItem interface.
+func (i *InteropItem) TryInteger() (*big.Int, error) {
+	return nil, errors.New("can't convert Interop to Integer")
+}
+
 // Equals implements StackItem interface.
 func (i *InteropItem) Equals(s StackItem) bool {
 	if i == s {