Skip to content

Commit

Permalink
chore(trie): lib/trie/hash.go tests (#2049)
Browse files Browse the repository at this point in the history
  • Loading branch information
@qdm12
qdm12 committed Nov 25, 2021
1 parent 594fee1 commit c89d50c
Show file tree
Hide file tree
Showing 8 changed files with 1,446 additions and 90 deletions.
77 changes: 77 additions & 0 deletions lib/trie/bytesBuffer_mock_test.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

123 changes: 78 additions & 45 deletions lib/trie/hash.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,10 @@ func hashNode(n node, digestBuffer io.Writer) (err error) {
// if length of encoded leaf is less than 32 bytes, do not hash
if encodingBuffer.Len() < 32 {
_, err = digestBuffer.Write(encodingBuffer.Bytes())
return err
if err != nil {
return fmt.Errorf("cannot write encoded node to buffer: %w", err)
}
return nil
}

// otherwise, hash encoded node
Expand All @@ -72,16 +75,26 @@ func hashNode(n node, digestBuffer io.Writer) (err error) {
}

_, err = digestBuffer.Write(hasher.Sum(nil))
return err
if err != nil {
return fmt.Errorf("cannot write hash sum of node to buffer: %w", err)
}
return nil
}

var ErrNodeTypeUnsupported = errors.New("node type is not supported")

type bytesBuffer interface {
// note: cannot compose with io.Writer for mock generation
Write(p []byte) (n int, err error)
Len() int
Bytes() []byte
}

// encodeNode writes the encoding of the node to the buffer given.
// It is the high-level function wrapping the encoding for different
// node types. The encoding has the following format:
// NodeHeader | Extra partial key length | Partial Key | Value
func encodeNode(n node, buffer *bytes.Buffer, parallel bool) (err error) {
func encodeNode(n node, buffer bytesBuffer, parallel bool) (err error) {
switch n := n.(type) {
case *branch:
err := encodeBranch(n, buffer, parallel)
Expand All @@ -104,70 +117,58 @@ func encodeNode(n node, buffer *bytes.Buffer, parallel bool) (err error) {
copy(n.encoding, buffer.Bytes())
return nil
case nil:
buffer.Write([]byte{0})
_, err := buffer.Write([]byte{0})
if err != nil {
return fmt.Errorf("cannot encode nil node: %w", err)
}
return nil
default:
return fmt.Errorf("%w: %T", ErrNodeTypeUnsupported, n)
}
}

func encodeAndHash(n node) ([]byte, error) {
buffer := digestBufferPool.Get().(*bytes.Buffer)
buffer.Reset()
defer digestBufferPool.Put(buffer)

err := hashNode(n, buffer)
if err != nil {
return nil, err
}

scEncChild, err := scale.Marshal(buffer.Bytes())
if err != nil {
return nil, err
}
return scEncChild, nil
}

// encodeBranch encodes a branch with the encoding specified at the top of this package
// to the buffer given.
func encodeBranch(b *branch, buffer io.Writer, parallel bool) (err error) {
if !b.dirty && b.encoding != nil {
_, err = buffer.Write(b.encoding)
if err != nil {
return fmt.Errorf("cannot write stored encoded branch to buffer: %w", err)
return fmt.Errorf("cannot write stored encoding to buffer: %w", err)
}
return nil
}

encoding, err := b.header()
encodedHeader, err := b.header()
if err != nil {
return fmt.Errorf("cannot encode branch header: %w", err)
return fmt.Errorf("cannot encode header: %w", err)
}

_, err = buffer.Write(encoding)
_, err = buffer.Write(encodedHeader)
if err != nil {
return fmt.Errorf("cannot write encoded branch header to buffer: %w", err)
return fmt.Errorf("cannot write encoded header to buffer: %w", err)
}

_, err = buffer.Write(nibblesToKeyLE(b.key))
keyLE := nibblesToKeyLE(b.key)
_, err = buffer.Write(keyLE)
if err != nil {
return fmt.Errorf("cannot write encoded branch key to buffer: %w", err)
return fmt.Errorf("cannot write encoded key to buffer: %w", err)
}

_, err = buffer.Write(common.Uint16ToBytes(b.childrenBitmap()))
childrenBitmap := common.Uint16ToBytes(b.childrenBitmap())
_, err = buffer.Write(childrenBitmap)
if err != nil {
return fmt.Errorf("cannot write branch children bitmap to buffer: %w", err)
return fmt.Errorf("cannot write children bitmap to buffer: %w", err)
}

if b.value != nil {
bytes, err := scale.Marshal(b.value)
if err != nil {
return fmt.Errorf("cannot scale encode branch value: %w", err)
return fmt.Errorf("cannot scale encode value: %w", err)
}

_, err = buffer.Write(bytes)
if err != nil {
return fmt.Errorf("cannot write encoded branch value to buffer: %w", err)
return fmt.Errorf("cannot write encoded value to buffer: %w", err)
}
}

Expand Down Expand Up @@ -222,10 +223,15 @@ func encodeChildrenInParallel(children [16]node, buffer io.Writer) (err error) {
// write as many completed buffers to the result buffer.
for currentIndex < len(children) &&
resultBuffers[currentIndex] != nil {
// note buffer.Write copies the byte slice given as argument
_, writeErr := buffer.Write(resultBuffers[currentIndex].Bytes())
if writeErr != nil && err == nil {
err = writeErr
bufferSlice := resultBuffers[currentIndex].Bytes()
if len(bufferSlice) > 0 {
// note buffer.Write copies the byte slice given as argument
_, writeErr := buffer.Write(bufferSlice)
if writeErr != nil && err == nil {
err = fmt.Errorf(
"cannot write encoding of child at index %d: %w",
currentIndex, writeErr)
}
}

encodingBufferPool.Put(resultBuffers[currentIndex])
Expand All @@ -246,17 +252,26 @@ func encodeChildrenInParallel(children [16]node, buffer io.Writer) (err error) {
}

func encodeChildrenSequentially(children [16]node, buffer io.Writer) (err error) {
for _, child := range children {
for i, child := range children {
err = encodeChild(child, buffer)
if err != nil {
return err
return fmt.Errorf("cannot encode child at index %d: %w", i, err)
}
}
return nil
}

func encodeChild(child node, buffer io.Writer) (err error) {
if child == nil {
var isNil bool
switch impl := child.(type) {
case *branch:
isNil = impl == nil
case *leaf:
isNil = impl == nil
default:
isNil = child == nil
}
if isNil {
return nil
}

Expand All @@ -273,6 +288,23 @@ func encodeChild(child node, buffer io.Writer) (err error) {
return nil
}

func encodeAndHash(n node) (b []byte, err error) {
buffer := digestBufferPool.Get().(*bytes.Buffer)
buffer.Reset()
defer digestBufferPool.Put(buffer)

err = hashNode(n, buffer)
if err != nil {
return nil, fmt.Errorf("cannot hash node: %w", err)
}

scEncChild, err := scale.Marshal(buffer.Bytes())
if err != nil {
return nil, fmt.Errorf("cannot scale encode hashed node: %w", err)
}
return scEncChild, nil
}

// encodeLeaf encodes a leaf to the buffer given, with the encoding
// specified at the top of this package.
func encodeLeaf(l *leaf, buffer io.Writer) (err error) {
Expand All @@ -286,27 +318,28 @@ func encodeLeaf(l *leaf, buffer io.Writer) (err error) {
return nil
}

encoding, err := l.header()
encodedHeader, err := l.header()
if err != nil {
return fmt.Errorf("cannot encode header: %w", err)
}

_, err = buffer.Write(encoding)
_, err = buffer.Write(encodedHeader)
if err != nil {
return fmt.Errorf("cannot write encoded header to buffer: %w", err)
}

_, err = buffer.Write(nibblesToKeyLE(l.key))
keyLE := nibblesToKeyLE(l.key)
_, err = buffer.Write(keyLE)
if err != nil {
return fmt.Errorf("cannot write LE key to buffer: %w", err)
}

bytes, err := scale.Marshal(l.value) // TODO scale encoder to write to buffer
encodedValue, err := scale.Marshal(l.value) // TODO scale encoder to write to buffer
if err != nil {
return err
return fmt.Errorf("cannot scale marshal value: %w", err)
}

_, err = buffer.Write(bytes)
_, err = buffer.Write(encodedValue)
if err != nil {
return fmt.Errorf("cannot write scale encoded value to buffer: %w", err)
}
Expand Down

0 comments on commit c89d50c

Please sign in to comment.