From 8b28e30088bfdff94d5adcbe4b25cec01267ab6d Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Tue, 30 Nov 2021 11:34:44 +0000
Subject: [PATCH] Add `Type()` method to `node.Node`

---
 lib/trie/branch/type.go      | 12 ++++++++++
 lib/trie/branch/type_test.go | 46 ++++++++++++++++++++++++++++++++++++
 lib/trie/leaf/type.go        |  8 +++++++
 lib/trie/leaf/type_test.go   | 18 ++++++++++++++
 4 files changed, 84 insertions(+)
 create mode 100644 lib/trie/branch/type.go
 create mode 100644 lib/trie/branch/type_test.go
 create mode 100644 lib/trie/leaf/type.go
 create mode 100644 lib/trie/leaf/type_test.go

diff --git a/lib/trie/branch/type.go b/lib/trie/branch/type.go
new file mode 100644
index 00000000000..c680652801b
--- /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 00000000000..0b49c6db86e
--- /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 00000000000..d182f2dd225
--- /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 00000000000..f3713d88f8f
--- /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)
+}