core: vm: sqlvm: runtime: implement opRepeatPK

2 files changed, 118 insertions, 21 deletions

diff --git a/core/vm/sqlvm/runtime/instructions.go b/core/vm/sqlvm/runtime/instructions.go
index d61bacdde..6fa72d61a 100644
--- a/core/vm/sqlvm/runtime/instructions.go
+++ b/core/vm/sqlvm/runtime/instructions.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"math/big"
 	"regexp"
 	"sort"
 	"strings"
@@ -98,6 +99,14 @@ func (op *Operand) toUint64() (result []uint64, err error) {
 	return
 }
 
+func (op *Operand) toTableRef() (schema.TableRef, error) {
+	t, err := op.toUint8()
+	if err != nil {
+		return 0, err
+	}
+	return schema.TableRef(t[0]), nil
+}
+
 func (op *Operand) toUint8() ([]uint8, error) {
 	result := make([]uint8, len(op.Data))
 	for i, tuple := range op.Data {
@@ -1881,3 +1890,30 @@ func opFunc(ctx *common.Context, ops, registers []*Operand, output uint) (err er
 	registers[output] = result
 	return
 }
+
+func uint64ToOperands(numbers []uint64) (*Operand, error) {
+	result := &Operand{
+		Meta: []ast.DataType{ast.ComposeDataType(ast.DataTypeMajorUint, 7)},
+		Data: []Tuple{},
+	}
+	result.Data = make([]Tuple, len(numbers))
+	for i, n := range numbers {
+		result.Data[i] = []*Raw{
+			{
+				Value: decimal.NewFromBigInt(new(big.Int).SetUint64(n), 0),
+				Bytes: nil,
+			},
+		}
+	}
+	return result, nil
+}
+
+func opRepeatPK(ctx *common.Context, input []*Operand, registers []*Operand, output int) (err error) {
+	tableRef, err := input[0].toTableRef()
+	if err != nil {
+		return err
+	}
+	IDs := ctx.Storage.RepeatPK(ctx.Contract.Address(), tableRef)
+	registers[output], err = uint64ToOperands(IDs)
+	return
+}
diff --git a/core/vm/sqlvm/runtime/instructions_test.go b/core/vm/sqlvm/runtime/instructions_test.go
index a399ad557..77ba0ff70 100644
--- a/core/vm/sqlvm/runtime/instructions_test.go
+++ b/core/vm/sqlvm/runtime/instructions_test.go
@@ -189,15 +189,6 @@ func hexToBytes(s string) []byte {
 	return b
 }
 
-type decodeTestCase struct {
-	dt             ast.DataType
-	expectData     *Raw
-	expectSlotHash dexCommon.Hash
-	shift          uint64
-	inputBytes     []byte
-	dBytes         []byte
-}
-
 type opLoadTestCase struct {
 	title          string
 	outputIdx      uint
@@ -205,7 +196,7 @@ type opLoadTestCase struct {
 	expectedErr    error
 	ids            []uint64
 	fields         []uint8
-	tableIdx       int8
+	tableRef       schema.TableRef
 }
 
 func (s *opLoadSuite) SetupTest() {
@@ -236,7 +227,7 @@ func (s *opLoadSuite) getOpLoadTestCases(raws []*raw) []opLoadTestCase {
 			expectedErr:    nil,
 			ids:            nil,
 			fields:         nil,
-			tableIdx:       0,
+			tableRef:       0,
 		},
 		{
 			title:          "NOT_EXIST_TABLE",
@@ -245,7 +236,7 @@ func (s *opLoadSuite) getOpLoadTestCases(raws []*raw) []opLoadTestCase {
 			expectedErr:    errors.ErrorCodeIndexOutOfRange,
 			ids:            nil,
 			fields:         nil,
-			tableIdx:       13,
+			tableRef:       13,
 		},
 		{
 			title:          "OK_CASE",
@@ -254,7 +245,7 @@ func (s *opLoadSuite) getOpLoadTestCases(raws []*raw) []opLoadTestCase {
 			expectedErr:    nil,
 			ids:            []uint64{123456, 654321},
 			fields:         s.getOKCaseFields(raws),
-			tableIdx:       1,
+			tableRef:       1,
 		},
 	}
 	return testCases
@@ -283,6 +274,15 @@ func (s *opLoadSuite) getOKCaseFields(raws []*raw) []uint8 {
 	return rValue
 }
 
+type decodeTestCase struct {
+	dt             ast.DataType
+	expectData     *Raw
+	expectSlotHash dexCommon.Hash
+	shift          uint64
+	inputBytes     []byte
+	dBytes         []byte
+}
+
 func (s *opLoadSuite) getDecodeTestCases(headHash dexCommon.Hash,
 	address dexCommon.Address, storage *common.Storage) []decodeTestCase {
 
@@ -304,9 +304,9 @@ func (s *opLoadSuite) getDecodeTestCases(headHash dexCommon.Hash,
 	return testCases
 }
 
-func (s *opLoadSuite) newRegisters(tableIdx int8, ids []uint64, fields []uint8) []*Operand {
+func (s *opLoadSuite) newRegisters(tableRef schema.TableRef, ids []uint64, fields []uint8) []*Operand {
 	o := make([]*Operand, 4)
-	o[1] = newTableNameOperand(tableIdx)
+	o[1] = newTableRefOperand(tableRef)
 	o[2] = newIDsOperand(ids)
 	o[3] = newFieldsOperand(fields)
 	return o
@@ -323,8 +323,8 @@ func newInput(nums []int) []*Operand {
 	return o
 }
 
-func newTableNameOperand(tableIdx int8) *Operand {
-	if tableIdx < 0 {
+func newTableRefOperand(tableRef schema.TableRef) *Operand {
+	if tableRef < 0 {
 		return nil
 	}
 	o := &Operand{
@@ -334,7 +334,7 @@ func newTableNameOperand(tableIdx int8) *Operand {
 		Data: []Tuple{
 			[]*Raw{
 				{
-					Value: decimal.New(int64(tableIdx), 0),
+					Value: decimal.New(int64(tableRef), 0),
 				},
 			},
 		},
@@ -404,7 +404,7 @@ func (s *opLoadSuite) TestOpLoad() {
 	testCases := s.getOpLoadTestCases(s.raws)
 	for _, t := range testCases {
 		input := newInput([]int{1, 2, 3})
-		reg := s.newRegisters(t.tableIdx, t.ids, t.fields)
+		reg := s.newRegisters(t.tableRef, t.ids, t.fields)
 
 		loadRegister(input, reg)
 		err := opLoad(s.ctx, input, reg, t.outputIdx)
@@ -417,11 +417,72 @@ func (s *opLoadSuite) TestOpLoad() {
 	}
 }
 
-func TestOpLoad(t *testing.T) {
-	suite.Run(t, new(opLoadSuite))
+type opRepeatPKSuite struct{ suite.Suite }
+
+type repeatPKTestCase struct {
+	tableRef    schema.TableRef
+	address     dexCommon.Address
+	title       string
+	expectedErr error
+	expectedIDs []uint64
+}
+
+func (s opRepeatPKSuite) newInput(tableRef schema.TableRef) []*Operand {
+	o := make([]*Operand, 1)
+	o[0] = newTableRefOperand(tableRef)
+	return o
+}
+
+func (s opRepeatPKSuite) newRegisters(tableRef schema.TableRef) []*Operand {
+	o := make([]*Operand, 2)
+	o[1] = newTableRefOperand(tableRef)
+	return o
+}
+
+func (s opRepeatPKSuite) getTestCases(storage *common.Storage) []repeatPKTestCase {
+	testCases := []repeatPKTestCase{
+		{
+			tableRef:    0,
+			address:     dexCommon.BytesToAddress([]byte("0")),
+			title:       "no IDs",
+			expectedErr: nil,
+			expectedIDs: []uint64{},
+		},
+		{
+			tableRef:    1,
+			address:     dexCommon.BytesToAddress([]byte("1")),
+			title:       "ok case",
+			expectedErr: nil,
+			expectedIDs: []uint64{1, 2, 3, 4},
+		},
+	}
+	for _, t := range testCases {
+		storage.SetPK(t.address, t.tableRef, t.expectedIDs)
+	}
+	return testCases
+}
+
+func (s opRepeatPKSuite) TestRepeatPK() {
+	ctx := &common.Context{}
+	ctx.Storage = newStorage()
+	testCases := s.getTestCases(ctx.Storage)
+	for _, t := range testCases {
+		address := t.address
+		ctx.Contract = vm.NewContract(vm.AccountRef(address),
+			vm.AccountRef(address), new(big.Int), uint64(0))
+		reg := s.newRegisters(t.tableRef)
+		input := newInput([]int{1})
+		loadRegister(input, reg)
+		err := opRepeatPK(ctx, input, reg, 0)
+		s.Require().Equalf(t.expectedErr, err, "testcase: [%v]", t.title)
+		result, _ := reg[0].toUint64()
+		s.Require().Equalf(t.expectedIDs, result, "testcase: [%v]", t.title)
+	}
 }
 
 func TestInstructions(t *testing.T) {
+	suite.Run(t, new(opLoadSuite))
+	suite.Run(t, new(opRepeatPKSuite))
 	suite.Run(t, new(instructionSuite))
 }