codetrie: add ssz types for chunk sizes 24,32,40

Author: s1na

codetrie/codetrie.go

@@ -2,6 +2,7 @@ package codetrie
 
 import (
 	"encoding/binary"
+	"errors"
 	"math"
 
 	sszlib "github.com/ferranbt/fastssz"
@@ -25,6 +26,12 @@ type Trie interface {
 	Update(key, value []byte)
 }
 
+type CodeTrie interface {
+	HashTreeRoot() ([32]byte, error)
+	HashTreeRootWith(*sszlib.Hasher) error
+	GetTree() (*sszlib.Node, error)
+}
+
 type Chunk struct {
 	fio  uint8 // firstInstructionOffset
 	code []byte
@@ -72,6 +79,7 @@ func GetSSZTree(code []byte, chunkSize uint) (*sszlib.Node, error) {
 
 	return codeTrie.GetTree()
 }
+
 func MerkleizeSSZSha(code []byte, chunkSize uint) (common.Hash, error) {
 	codeTrie, err := prepareSSZ(code, chunkSize)
 	if err != nil {
@@ -104,20 +112,64 @@ func MerkleizeSSZKeccak(code []byte, chunkSize uint) (common.Hash, error) {
 	return common.BytesToHash(root[:]), err
 }
 
-func prepareSSZ(code []byte, chunkSize uint) (*ssz.CodeTrie, error) {
+func prepareSSZ(code []byte, chunkSize uint) (CodeTrie, error) {
 	rawChunks := Chunkify(code, chunkSize)
-	chunks := make([]*ssz.Chunk, len(rawChunks))
+	metadata := &ssz.Metadata{Version: 0, CodeHash: crypto.Keccak256(code), CodeLength: uint16(len(code))}
+	switch chunkSize {
+	case 24:
+		chunks := prepareSSZChunks24(rawChunks)
+		return &ssz.CodeTrie24{Metadata: metadata, Chunks: chunks}, nil
+	case 32:
+		chunks := prepareSSZChunks32(rawChunks)
+		return &ssz.CodeTrie32{Metadata: metadata, Chunks: chunks}, nil
+	case 40:
+		chunks := prepareSSZChunks40(rawChunks)
+		return &ssz.CodeTrie40{Metadata: metadata, Chunks: chunks}, nil
+	default:
+		return nil, errors.New("chunk size should be one of (24, 32, 40)")
+	}
+}
+
+func prepareSSZChunks24(rawChunks []*Chunk) []*ssz.Chunk24 {
+	chunks := make([]*ssz.Chunk24, len(rawChunks))
 	for i, rc := range rawChunks {
 		code := rc.code
-		if uint(len(code)) < chunkSize {
-			code = make([]byte, chunkSize)
+		if uint(len(code)) < 24 {
+			code = make([]byte, 24)
 			copy(code[:len(rc.code)], rc.code)
 		}
-		chunks[i] = &ssz.Chunk{FIO: rc.fio, Code: code}
+		chunks[i] = &ssz.Chunk24{FIO: rc.fio, Code: code}
 	}
 
-	metadata := &ssz.Metadata{Version: 0, CodeHash: crypto.Keccak256(code), CodeLength: uint16(len(code))}
-	return &ssz.CodeTrie{Metadata: metadata, Chunks: chunks}, nil
+	return chunks
+}
+
+func prepareSSZChunks32(rawChunks []*Chunk) []*ssz.Chunk32 {
+	chunks := make([]*ssz.Chunk32, len(rawChunks))
+	for i, rc := range rawChunks {
+		code := rc.code
+		if uint(len(code)) < 32 {
+			code = make([]byte, 32)
+			copy(code[:len(rc.code)], rc.code)
+		}
+		chunks[i] = &ssz.Chunk32{FIO: rc.fio, Code: code}
+	}
+
+	return chunks
+}
+
+func prepareSSZChunks40(rawChunks []*Chunk) []*ssz.Chunk40 {
+	chunks := make([]*ssz.Chunk40, len(rawChunks))
+	for i, rc := range rawChunks {
+		code := rc.code
+		if uint(len(code)) < 40 {
+			code = make([]byte, 40)
+			copy(code[:len(rc.code)], rc.code)
+		}
+		chunks[i] = &ssz.Chunk40{FIO: rc.fio, Code: code}
+	}
+
+	return chunks
 }
 
 func merkleize(code []byte, chunkSize uint, trie Trie) {

codetrie/ssz/types.go

@@ -8,12 +8,32 @@ type Metadata struct {
 	CodeLength uint16
 }
 
-type Chunk struct {
+type Chunk32 struct {
 	FIO  uint8
 	Code []byte `ssz-size:"32"` // Last chunk is right-padded with zeros
 }
 
-type CodeTrie struct {
+type Chunk24 struct {
+	FIO  uint8
+	Code []byte `ssz-size:"24"` // Last chunk is right-padded with zeros
+}
+
+type Chunk40 struct {
+	FIO  uint8
+	Code []byte `ssz-size:"40"` // Last chunk is right-padded with zeros
+}
+
+type CodeTrie32 struct {
+	Metadata *Metadata
+	Chunks   []*Chunk32 `ssz-max:"1024"`
+}
+
+type CodeTrie24 struct {
+	Metadata *Metadata
+	Chunks   []*Chunk24 `ssz-max:"1024"`
+}
+
+type CodeTrie40 struct {
 	Metadata *Metadata
-	Chunks   []*Chunk `ssz-max:"1024"`
+	Chunks   []*Chunk40 `ssz-max:"1024"`
 }