diff --git a/lib/trie/branch/type.go b/lib/trie/branch/type.go new file mode 100644 index 0000000000..c680652801 --- /dev/null +++ b/lib/trie/branch/type.go @@ -0,0 +1,12 @@ +package branch + +import "github.com/ChainSafe/gossamer/lib/trie/node" + +// Type returns node.BranchType if the branch value +// is nil, and node.BranchWithValueType otherwise. +func (b *Branch) Type() node.Type { + if b.Value == nil { + return node.BranchType + } + return node.BranchWithValueType +} diff --git a/lib/trie/branch/type_test.go b/lib/trie/branch/type_test.go new file mode 100644 index 0000000000..0b49c6db86 --- /dev/null +++ b/lib/trie/branch/type_test.go @@ -0,0 +1,46 @@ +// Copyright 2021 ChainSafe Systems (ON) +// SPDX-License-Identifier: LGPL-3.0-only + +package branch + +import ( + "testing" + + "github.com/ChainSafe/gossamer/lib/trie/node" + "github.com/stretchr/testify/assert" +) + +func Test_Branch_Type(t *testing.T) { + testCases := map[string]struct { + branch *Branch + Type node.Type + }{ + "nil value": { + branch: &Branch{}, + Type: node.BranchType, + }, + "empty value": { + branch: &Branch{ + Value: []byte{}, + }, + Type: node.BranchWithValueType, + }, + "non empty value": { + branch: &Branch{ + Value: []byte{1}, + }, + Type: node.BranchWithValueType, + }, + } + + for name, testCase := range testCases { + testCase := testCase + t.Run(name, func(t *testing.T) { + t.Parallel() + + Type := testCase.branch.Type() + + assert.Equal(t, testCase.Type, Type) + }) + } +} diff --git a/lib/trie/leaf/type.go b/lib/trie/leaf/type.go new file mode 100644 index 0000000000..d182f2dd22 --- /dev/null +++ b/lib/trie/leaf/type.go @@ -0,0 +1,8 @@ +package leaf + +import "github.com/ChainSafe/gossamer/lib/trie/node" + +// Type returns node.LeafType. +func (l *Leaf) Type() node.Type { + return node.LeafType +} diff --git a/lib/trie/leaf/type_test.go b/lib/trie/leaf/type_test.go new file mode 100644 index 0000000000..f3713d88f8 --- /dev/null +++ b/lib/trie/leaf/type_test.go @@ -0,0 +1,18 @@ +package leaf + +import ( + "testing" + + "github.com/ChainSafe/gossamer/lib/trie/node" + "github.com/stretchr/testify/assert" +) + +func Test_Leaf_Type(t *testing.T) { + t.Parallel() + + leaf := new(Leaf) + + Type := leaf.Type() + + assert.Equal(t, node.LeafType, Type) +}