From 96806262bf39d985596d0a9eff384a232b23ab88 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Mon, 6 Apr 2020 10:14:26 +0300
Subject: [PATCH] vm: handle negative arguments in SHL/SHR

Do it as in reference implementation: a >> -b == a << b.
---
 pkg/vm/vm.go      | 11 ++++++++++-
 pkg/vm/vm_test.go | 15 ++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 555deb8ce..aeb8b0180 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -806,8 +806,17 @@ func (v *VM) execute(ctx *Context, op opcode.Opcode, parameter []byte) (err erro
 		a := v.estack.Pop().BigInt()
 		v.checkBigIntSize(a)
 
+		newOp := op
+		if b < 0 {
+			b = -b
+			if op == opcode.SHR {
+				newOp = opcode.SHL
+			} else {
+				newOp = opcode.SHR
+			}
+		}
 		var item big.Int
-		if op == opcode.SHL {
+		if newOp == opcode.SHL {
 			item.Lsh(a, uint(b))
 		} else {
 			item.Rsh(a, uint(b))
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 26cebd624..23d816820 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -844,7 +844,7 @@ func TestMULBigResult(t *testing.T) {
 	checkVMFailed(t, vm)
 }
 
-func TestDivMod(t *testing.T) {
+func TestArithNegativeArguments(t *testing.T) {
 	runCase := func(op opcode.Opcode, p, q, result int64) func(t *testing.T) {
 		return func(t *testing.T) {
 			vm := load(makeProgram(op))
@@ -869,6 +869,19 @@ func TestDivMod(t *testing.T) {
 		t.Run("negative/negative", runCase(opcode.MOD, -5, -2, -1))
 	})
 
+	t.Run("SHR", func(t *testing.T) {
+		t.Run("positive/positive", runCase(opcode.SHR, 5, 2, 1))
+		t.Run("positive/negative", runCase(opcode.SHR, 5, -2, 20))
+		t.Run("negative/positive", runCase(opcode.SHR, -5, 2, -2))
+		t.Run("negative/negative", runCase(opcode.SHR, -5, -2, -20))
+	})
+
+	t.Run("SHL", func(t *testing.T) {
+		t.Run("positive/positive", runCase(opcode.SHL, 5, 2, 20))
+		t.Run("positive/negative", runCase(opcode.SHL, 5, -2, 1))
+		t.Run("negative/positive", runCase(opcode.SHL, -5, 2, -20))
+		t.Run("negative/negative", runCase(opcode.SHL, -5, -2, -2))
+	})
 }
 
 func TestSub(t *testing.T) {