From 64491a4d83ab99f81ee10c708723e47fb3c54838 Mon Sep 17 00:00:00 2001
From: BlockChainDev <kevtheappdev@gmail.com>
Date: Wed, 27 Feb 2019 21:38:11 +0000
Subject: [PATCH] VM: Add Sub, Mul, Mod LSH, RSH

---
 pkg/vm/stack/Int.go      | 35 ++++++++++++++++++++++++
 pkg/vm/stack/int_test.go | 57 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)
 create mode 100644 pkg/vm/stack/int_test.go

diff --git a/pkg/vm/stack/Int.go b/pkg/vm/stack/Int.go
index 6ba864b90..ae81a2120 100644
--- a/pkg/vm/stack/Int.go
+++ b/pkg/vm/stack/Int.go
@@ -27,11 +27,46 @@ func (i *Int) Equal(s *Int) bool {
 
 // Add will add two stackIntegers together
 func (i *Int) Add(s *Int) (*Int, error) {
+	return &Int{
+		val: new(big.Int).Add(i.val, s.val),
+	}, nil
+}
+
+// Sub will subtract two stackIntegers together
+func (i *Int) Sub(s *Int) (*Int, error) {
 	return &Int{
 		val: new(big.Int).Sub(i.val, s.val),
 	}, nil
 }
 
+// Mul will multiply two stackIntegers together
+func (i *Int) Mul(s *Int) (*Int, error) {
+	return &Int{
+		val: new(big.Int).Mul(i.val, s.val),
+	}, nil
+}
+
+// Mod will take the mod of two stackIntegers together
+func (i *Int) Mod(s *Int) (*Int, error) {
+	return &Int{
+		val: new(big.Int).Mod(i.val, s.val),
+	}, nil
+}
+
+// Rsh will shift the integer b to the right by `n` bits
+func (i *Int) Rsh(n *Int) (*Int, error) {
+	return &Int{
+		val: new(big.Int).Rsh(i.val, uint(n.val.Int64())),
+	}, nil
+}
+
+// Lsh will shift the integer b to the left by `n` bits
+func (i *Int) Lsh(n *Int) (*Int, error) {
+	return &Int{
+		val: new(big.Int).Lsh(i.val, uint(n.val.Int64())),
+	}, nil
+}
+
 // Integer will overwrite the default implementation
 // to allow go to cast this item as an integer.
 func (i *Int) Integer() (*Int, error) {
diff --git a/pkg/vm/stack/int_test.go b/pkg/vm/stack/int_test.go
new file mode 100644
index 000000000..40ef3c8b2
--- /dev/null
+++ b/pkg/vm/stack/int_test.go
@@ -0,0 +1,57 @@
+package stack
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAdd(t *testing.T) {
+	a := testMakeStackInt(t, 10)
+	b := testMakeStackInt(t, 20)
+	expected := testMakeStackInt(t, 30)
+	c, err := a.Add(b)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}
+func TestSub(t *testing.T) {
+	a := testMakeStackInt(t, 30)
+	b := testMakeStackInt(t, 200)
+	expected := testMakeStackInt(t, 170)
+	c, err := b.Sub(a)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}
+func TestMul(t *testing.T) {
+	a := testMakeStackInt(t, 10)
+	b := testMakeStackInt(t, 20)
+	expected := testMakeStackInt(t, 200)
+	c, err := a.Mul(b)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}
+func TestMod(t *testing.T) {
+	a := testMakeStackInt(t, 10)
+	b := testMakeStackInt(t, 20)
+	expected := testMakeStackInt(t, 10)
+	c, err := a.Mod(b)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}
+func TestLsh(t *testing.T) {
+	a := testMakeStackInt(t, 23)
+	b := testMakeStackInt(t, 8)
+	expected := testMakeStackInt(t, 5888)
+	c, err := a.Lsh(b)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}
+
+func TestRsh(t *testing.T) {
+	a := testMakeStackInt(t, 128)
+	b := testMakeStackInt(t, 3)
+	expected := testMakeStackInt(t, 16)
+	c, err := a.Rsh(b)
+	assert.Nil(t, err)
+	assert.Equal(t, true, expected.Equal(c))
+}