Introducing Abstract Syntax Tree Enum parsing support

ast/builder.go

@@ -148,6 +148,19 @@ func (b *ASTBuilder) EnterStateVariableDeclaration(ctx *parser.StateVariableDecl
 	b.currentContract.StateVariables = append(b.currentContract.StateVariables, variable)
 }
 
+func (b *ASTBuilder) EnterEnumDefinition(ctx *parser.EnumDefinitionContext) {
+	enum := &EnumNode{
+		Name:         ctx.GetName().GetText(),
+		MemberValues: make([]*EnumMemberNode, len(ctx.GetEnumValues())),
+	}
+
+	for i, valueCtx := range ctx.GetEnumValues() {
+		enum.MemberValues[i] = &EnumMemberNode{Name: valueCtx.GetText()}
+	}
+
+	b.currentContract.Enums = append(b.currentContract.Enums, enum)
+}
+
 func (b *ASTBuilder) EnterConstructorDefinition(ctx *parser.ConstructorDefinitionContext) {
 	constructor := &ConstructorNode{
 		Parameters: make([]*VariableNode, 0),

ast/contract.go

@@ -4,6 +4,7 @@ package ast
 type ContractNode struct {
 	Name           string                `json:"name"`
 	StateVariables []*StateVariableNode  `json:"variables"`
+	Enums          []*EnumNode           `json:"enums"`
 	Structs        []*StructNode         `json:"structs"`
 	Events         []*EventNode          `json:"events"`
 	Errors         []*ErrorNode          `json:"errors"`
@@ -15,7 +16,7 @@ type ContractNode struct {
 }
 
 func (c *ContractNode) Children() []Node {
-	nodes := make([]Node, len(c.Functions)+len(c.StateVariables)+len(c.Events)+len(c.Errors)+len(c.Structs)+len(c.Using))
+	nodes := make([]Node, len(c.Functions)+len(c.StateVariables)+len(c.Events)+len(c.Errors)+len(c.Structs)+len(c.Using)+len(c.Enums))
 	if c.Constructor != nil {
 		nodes = append(nodes, c.Constructor)
 	}
@@ -37,5 +38,8 @@ func (c *ContractNode) Children() []Node {
 	for i, using := range c.Using {
 		nodes[i+len(c.Functions)+len(c.StateVariables)+len(c.Events)+len(c.Errors)+len(c.Structs)] = using
 	}
+	for i, enum := range c.Enums {
+		nodes[i+len(c.Functions)+len(c.StateVariables)+len(c.Events)+len(c.Errors)+len(c.Structs)+len(c.Using)] = enum
+	}
 	return nodes
 }

ast/enum.go

@@ -0,0 +1,22 @@
+package ast
+
+type EnumMemberNode struct {
+	Name string `json:"name"`
+}
+
+func (e *EnumMemberNode) Children() []Node {
+	return nil
+}
+
+type EnumNode struct {
+	Name         string            `json:"name"`         // Name of the enum
+	MemberValues []*EnumMemberNode `json:"memberValues"` // Values of the enum members
+}
+
+func (e *EnumNode) Children() []Node {
+	nodes := make([]Node, len(e.MemberValues))
+	for i, member := range e.MemberValues {
+		nodes[i] = member
+	}
+	return nodes
+}

ast/enum_test.go

@@ -0,0 +1,93 @@
+package ast
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestEnumNode_Children(t *testing.T) {
+	tests := []struct {
+		name         string
+		memberValues []*EnumMemberNode
+		expected     []Node
+	}{
+		{
+			name:         "Empty Enum",
+			memberValues: []*EnumMemberNode{},
+			expected:     []Node{},
+		},
+		{
+			name: "Enum with Members",
+			memberValues: []*EnumMemberNode{
+				{Name: "Value1"},
+				{Name: "Value2"},
+				{Name: "Value3"},
+			},
+			expected: []Node{
+				&EnumMemberNode{Name: "Value1"},
+				&EnumMemberNode{Name: "Value2"},
+				&EnumMemberNode{Name: "Value3"},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			enum := &EnumNode{
+				Name:         "Color",
+				MemberValues: test.memberValues,
+			}
+
+			children := enum.Children()
+
+			assert.Equal(t, test.expected, children)
+		})
+	}
+}
+
+func TestEnumNode_JSONMarshalling(t *testing.T) {
+	tests := []struct {
+		name     string
+		enum     *EnumNode
+		expected string
+	}{
+		{
+			name: "Empty Enum",
+			enum: &EnumNode{
+				Name:         "Color",
+				MemberValues: []*EnumMemberNode{},
+			},
+			expected: `{"name":"Color","memberValues":[]}`,
+		},
+		{
+			name: "Enum with Members",
+			enum: &EnumNode{
+				Name: "Color",
+				MemberValues: []*EnumMemberNode{
+					{Name: "Red"},
+					{Name: "Green"},
+					{Name: "Blue"},
+				},
+			},
+			expected: `{"name":"Color","memberValues":[{"name":"Red"},{"name":"Green"},{"name":"Blue"}]}`,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			jsonBytes, err := json.Marshal(test.enum)
+			require.NoError(t, err)
+
+			assert.JSONEq(t, test.expected, string(jsonBytes))
+
+			unmarshaledEnum := &EnumNode{}
+			err = json.Unmarshal(jsonBytes, unmarshaledEnum)
+			require.NoError(t, err)
+
+			assert.Equal(t, test.enum, unmarshaledEnum)
+		})
+	}
+}