New memory account structure (#190)

* Update init memory instruction args

* Hanlde unmarshalling both version of memory accounts

* Reduce duplicated code for unmarshalling memory account structs

* Rename simple allocator to slice allocator

Author: jeffyanta

pkg/solana/cvm/accounts_memory_account.go

@@ -14,26 +14,31 @@ const (
 )
 
 type MemoryAccount struct {
-	Vm     ed25519.PublicKey
-	Name   string
-	Bump   uint8
-	Layout MemoryLayout
+	Vm          ed25519.PublicKey
+	Name        string
+	Bump        uint8
+	Version     MemoryVersion
+	AccountSize uint16
+	NumAccounts uint32
 }
 
 type MemoryAccountWithData struct {
-	Vm     ed25519.PublicKey
-	Name   string
-	Bump   uint8
-	Layout MemoryLayout
-	Data   SimpleMemoryAllocator // todo: support other implementations
+	Vm          ed25519.PublicKey
+	Name        string
+	Bump        uint8
+	Version     MemoryVersion
+	AccountSize uint16
+	NumAccounts uint32
+	Data        SliceAllocator // todo: support other implementations
 }
 
 const MemoryAccountSize = (8 + // discriminator
 	32 + // vm
 	MaxMemoryAccountNameLength + // name
 	1 + // bump
-	6 + // padding
-	1) // layout
+	1 + // version
+	2 + // account_size
+	4) // num_accounts
 
 var MemoryAccountDiscriminator = []byte{byte(AccountTypeMemory), 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}
 
@@ -53,77 +58,81 @@ func (obj *MemoryAccount) Unmarshal(data []byte) error {
 	getKey(data, &obj.Vm, &offset)
 	getFixedString(data, &obj.Name, MaxMemoryAccountNameLength, &offset)
 	getUint8(data, &obj.Bump, &offset)
-	offset += 6 // padding
-	getMemoryLayout(data, &obj.Layout, &offset)
+	getMemoryVersion(data, &obj.Version, &offset)
+
+	switch obj.Version {
+	case MemoryVersionV0:
+		var layout MemoryLayout
+		offset += 5 // padding
+		getMemoryLayout(data, &layout, &offset)
+
+		switch layout {
+		case MemoryLayoutTimelock:
+			obj.AccountSize = uint16(GetVirtualAccountSizeInMemory(VirtualAccountTypeTimelock))
+		case MemoryLayoutNonce:
+			obj.AccountSize = uint16(GetVirtualAccountSizeInMemory(VirtualAccountTypeDurableNonce))
+		case MemoryLayoutRelay:
+			obj.AccountSize = uint16(GetVirtualAccountSizeInMemory(VirtualAccountTypeRelay))
+		default:
+			return errors.New("unsupported memory layout")
+		}
+		obj.NumAccounts = MemoryV0NumAccounts
+
+	case MemoryVersionV1:
+		getUint16(data, &obj.AccountSize, &offset)
+		getUint32(data, &obj.NumAccounts, &offset)
+
+	default:
+		return errors.New("invalid memory account version")
+	}
 
 	return nil
 }
 
 func (obj *MemoryAccount) String() string {
 	return fmt.Sprintf(
-		"MemoryAccount{vm=%s,bump=%d,name=%s,layout=%d}",
+		"MemoryAccount{vm=%s,bump=%d,name=%s,version=%d,num_accounts=%d,account_size=%d}",
 		base58.Encode(obj.Vm),
 		obj.Bump,
 		obj.Name,
-		obj.Layout,
+		obj.Version,
+		obj.NumAccounts,
+		obj.AccountSize,
 	)
 }
 
 func (obj *MemoryAccountWithData) Unmarshal(data []byte) error {
-	if len(data) < MemoryAccountSize {
-		return ErrInvalidAccountData
+	var memoryAccount MemoryAccount
+	if err := memoryAccount.Unmarshal(data); err != nil {
+		return err
 	}
 
-	var offset int
+	obj.Vm = memoryAccount.Vm
+	obj.Name = memoryAccount.Name
+	obj.Bump = memoryAccount.Bump
+	obj.Version = memoryAccount.Version
+	obj.AccountSize = memoryAccount.AccountSize
+	obj.NumAccounts = memoryAccount.NumAccounts
 
-	var discriminator []byte
-	getDiscriminator(data, &discriminator, &offset)
-	if !bytes.Equal(discriminator, MemoryAccountDiscriminator) {
+	if len(data) < MemoryAccountSize+GetSliceAllocatorSize(int(obj.NumAccounts), int(obj.AccountSize)) {
 		return ErrInvalidAccountData
 	}
 
-	getKey(data, &obj.Vm, &offset)
-	getFixedString(data, &obj.Name, MaxMemoryAccountNameLength, &offset)
-	getUint8(data, &obj.Bump, &offset)
-	offset += 6 // padding
-	getMemoryLayout(data, &obj.Layout, &offset)
-
-	switch obj.Layout {
-	case MemoryLayoutTimelock:
-		capacity := CompactStateItems
-		itemSize := int(GetVirtualAccountSizeInMemory(VirtualAccountTypeTimelock))
-		if len(data) < MemoryAccountSize+GetSimpleMemoryAllocatorSize(capacity, itemSize) {
-			return ErrInvalidAccountData
-		}
-		getSimpleMemoryAllocator(data, &obj.Data, capacity, itemSize, &offset)
-	case MemoryLayoutNonce:
-		capacity := CompactStateItems
-		itemSize := int(GetVirtualAccountSizeInMemory(VirtualAccountTypeDurableNonce))
-		if len(data) < MemoryAccountSize+GetSimpleMemoryAllocatorSize(capacity, itemSize) {
-			return ErrInvalidAccountData
-		}
-		getSimpleMemoryAllocator(data, &obj.Data, capacity, itemSize, &offset)
-	case MemoryLayoutRelay:
-		capacity := CompactStateItems
-		itemSize := int(GetVirtualAccountSizeInMemory(VirtualAccountTypeRelay))
-		if len(data) < MemoryAccountSize+GetSimpleMemoryAllocatorSize(capacity, itemSize) {
-			return ErrInvalidAccountData
-		}
-		getSimpleMemoryAllocator(data, &obj.Data, capacity, itemSize, &offset)
-	default:
-		return errors.New("unsupported memory layout")
-	}
+	offset := MemoryAccountSize
+	getSliceAllocator(data, &obj.Data, int(obj.NumAccounts), int(obj.AccountSize), &offset)
 
 	return nil
 }
 
 func (obj *MemoryAccountWithData) String() string {
 	return fmt.Sprintf(
-		"MemoryAccountWithData{vm=%s,name=%s,bump=%d,layout=%d,data=%s}",
+		"MemoryAccountWithData{vm=%s,name=%s,bump=%d,version=%d,num_accounts=%d,account_size=%d,data=%s}",
 		base58.Encode(obj.Vm),
 		obj.Name,
 		obj.Bump,
-		obj.Layout,
+		obj.Version,
+		obj.NumAccounts,
+		obj.AccountSize,
 		obj.Data.String(),
 	)
 }

pkg/solana/cvm/instructions_init_memory.go

@@ -8,13 +8,15 @@ import (
 
 const (
 	InitMemoryInstructionArgsSize = (MaxMemoryAccountNameLength + // name
-		1 + // layout
+		4 + // num_accounts
+		2 + // account_size
 		1) // vm_memory_bump
 )
 
 type InitMemoryInstructionArgs struct {
 	Name         string
-	Layout       MemoryLayout
+	NumAccounts  uint32
+	AccountSize  uint16
 	VmMemoryBump uint8
 }
 
@@ -35,7 +37,8 @@ func NewInitMemoryInstruction(
 
 	putCodeInstruction(data, CodeInstructionInitMemory, &offset)
 	putFixedString(data, args.Name, MaxMemoryAccountNameLength, &offset)
-	putMemoryLayout(data, args.Layout, &offset)
+	putUint32(data, args.NumAccounts, &offset)
+	putUint16(data, args.AccountSize, &offset)
 	putUint8(data, args.VmMemoryBump, &offset)
 
 	return solana.Instruction{

pkg/solana/cvm/types_memory_allocator.go

@@ -1,7 +1,7 @@
 package cvm
 
 const (
-	CompactStateItems = 32_000
+	MemoryV0NumAccounts = 32_000
 )
 
 type MemoryAllocator interface {

pkg/solana/cvm/types_memory_version.go

@@ -0,0 +1,17 @@
+package cvm
+
+type MemoryVersion uint8
+
+const (
+	MemoryVersionV0 MemoryVersion = iota
+	MemoryVersionV1
+)
+
+func putMemoryVersion(dst []byte, v MemoryVersion, offset *int) {
+	dst[*offset] = uint8(v)
+	*offset += 1
+}
+func getMemoryVersion(src []byte, dst *MemoryVersion, offset *int) {
+	*dst = MemoryVersion(src[*offset])
+	*offset += 1
+}

…ana/cvm/types_simple_memory_allocator.go pkg/solana/cvm/types_slice_allocator.gopkg/solana/cvm/types_simple_memory_allocator.go renamed to pkg/solana/cvm/types_slice_allocator.go pkg/solana/cvm/types_simple_memory_allocator.go renamed to pkg/solana/cvm/types_slice_allocator.go

@@ -7,13 +7,13 @@ import (
 	"github.com/mr-tron/base58"
 )
 
-type SimpleMemoryAllocator struct {
+type SliceAllocator struct {
 	State ItemStateArray
 	Data  [][]byte
 }
 
-func (obj *SimpleMemoryAllocator) Unmarshal(data []byte, capacity, itemSize int) error {
-	if len(data) < GetSimpleMemoryAllocatorSize(capacity, itemSize) {
+func (obj *SliceAllocator) Unmarshal(data []byte, capacity, itemSize int) error {
+	if len(data) < GetSliceAllocatorSize(capacity, itemSize) {
 		return ErrInvalidAccountData
 	}
 
@@ -30,19 +30,19 @@ func (obj *SimpleMemoryAllocator) Unmarshal(data []byte, capacity, itemSize int)
 	return nil
 }
 
-func getSimpleMemoryAllocator(src []byte, dst *SimpleMemoryAllocator, capacity, itemSize int, offset *int) {
+func getSliceAllocator(src []byte, dst *SliceAllocator, capacity, itemSize int, offset *int) {
 	dst.Unmarshal(src[*offset:], capacity, itemSize)
-	*offset += GetSimpleMemoryAllocatorSize(capacity, itemSize)
+	*offset += GetSliceAllocatorSize(capacity, itemSize)
 }
 
-func (obj *SimpleMemoryAllocator) IsAllocated(index int) bool {
+func (obj *SliceAllocator) IsAllocated(index int) bool {
 	if index >= len(obj.State) {
 		return false
 	}
 	return obj.State[index] == ItemStateAllocated
 }
 
-func (obj *SimpleMemoryAllocator) Read(index int) ([]byte, bool) {
+func (obj *SliceAllocator) Read(index int) ([]byte, bool) {
 	if !obj.IsAllocated(index) {
 		return nil, false
 	}
@@ -52,21 +52,21 @@ func (obj *SimpleMemoryAllocator) Read(index int) ([]byte, bool) {
 	return copied, true
 }
 
-func (obj *SimpleMemoryAllocator) String() string {
+func (obj *SliceAllocator) String() string {
 	dataStringValues := make([]string, len(obj.Data))
 	for i := 0; i < len(obj.Data); i++ {
 		dataStringValues[i] = base58.Encode(obj.Data[i])
 	}
 	dataString := fmt.Sprintf("[%s]", strings.Join(dataStringValues, ","))
 
 	return fmt.Sprintf(
-		"SimpleMemoryAllocator{state=%s,data=%s}",
+		"SliceAllocator{state=%s,data=%s}",
 		obj.State.String(),
 		dataString,
 	)
 }
 
-func GetSimpleMemoryAllocatorSize(capacity, itemSize int) int {
+func GetSliceAllocatorSize(capacity, itemSize int) int {
 	return (capacity*ItemStateSize + // state
 		capacity*itemSize) // data
 }

pkg/solana/cvm/utils.go

@@ -107,6 +107,10 @@ func putUint32(dst []byte, v uint32, offset *int) {
 	binary.LittleEndian.PutUint32(dst[*offset:], v)
 	*offset += 4
 }
+func getUint32(src []byte, dst *uint32, offset *int) {
+	*dst = binary.LittleEndian.Uint32(src[*offset:])
+	*offset += 4
+}
 
 func putUint64(dst []byte, v uint64, offset *int) {
 	binary.LittleEndian.PutUint64(dst[*offset:], v)