Support for Enum Types in ABI Generation (#8)

* Introducing ENUM ABI parsing support

* Small fix in Enums.sol

abi_helpers.go

@@ -35,6 +35,15 @@ func isStructType(definedStructs map[string]MethodIO, typeName string) bool {
 	return exists
 }
 
+// isEnumType checks if a given type is an enumerated type.
+// It takes a map of defined enums and a type name as arguments.
+// The function returns true if the type name exists in the map of defined enums, indicating that it is an enumerated type.
+// Otherwise, it returns false.
+func isEnumType(definedEnums map[string]bool, typeName string) bool {
+	_, exists := definedEnums[typeName]
+	return exists
+}
+
 // parseMappingType parses a mapping type in Solidity ABI.
 // It takes a string of the form "mapping(keyType => valueType)" and returns three values:
 //   - A boolean indicating whether the parsing was successful. If the string is not a mapping type, this will be false.

abi_listener.go

@@ -19,6 +19,8 @@ func NewAbiListener() *AbiListener {
 	return &AbiListener{parser: &AbiParser{
 		abi:            ABI{},
 		definedStructs: make(map[string]MethodIO),
+		definedEnums:   make(map[string]bool), // map to keep track of defined enum types
+
 	}}
 }
 
@@ -106,6 +108,18 @@ func (l *AbiListener) ExitStructDefinition(ctx *parser.StructDefinitionContext)
 	}
 }
 
+// EnterEnumDefinition is a method of the AbiParser struct that is called when the parser encounters an enum definition in the Solidity code.
+// It takes a context of type EnumDefinitionContext, which contains the information about the enum definition in the code.
+// The method uses this context to extract the enum's name and its values, and stores them in the AbiParser's map of defined enums for future reference.
+// This allows the parser to recognize the enum type when it is used later in the code.
+func (p *AbiListener) EnterEnumDefinition(ctx *parser.EnumDefinitionContext) {
+	if ctx.AllIdentifier() != nil {
+		for _, id := range ctx.AllIdentifier() {
+			p.parser.definedEnums[id.GetText()] = true
+		}
+	}
+}
+
 // EnterFallbackFunctionDefinition is called when the parser enters a fallback function definition.
 // It injects the fallback function into the ABI.
 func (l *AbiListener) EnterFallbackFunctionDefinition(ctx *parser.FallbackFunctionDefinitionContext) {

abi_listener_test.go

@@ -45,6 +45,11 @@ func TestAbiListener(t *testing.T) {
 			contract: ReadContractFileForTest(t, "Structs").Content,
 			expected: `[{"inputs":[{"internalType":"struct MyStructs.ClasicStruct","name":"ClasicStruct","type":"tuple","components":[{"internalType":"uint256","name":"one","type":"uint256"},{"internalType":"uint256","name":"two","type":"uint256"}]},{"internalType":"struct MyStructs.NestedStruct","name":"NestedStruct","type":"tuple","components":[{"internalType":"uint256","name":"one","type":"uint256"},{"internalType":"uint256","name":"two","type":"uint256"},{"internalType":"struct MyStructs.ClasicStruct","name":"myStruct","type":"tuple","components":[{"internalType":"uint256","name":"one","type":"uint256"},{"internalType":"uint256","name":"two","type":"uint256"}]}]},{"internalType":"uint256","name":"Integer","type":"uint256"}],"outputs":[{"internalType":"struct MyStructs.NestedStruct","name":"NestedStruct","type":"tuple","components":[{"internalType":"uint256","name":"one","type":"uint256"},{"internalType":"uint256","name":"two","type":"uint256"},{"internalType":"struct MyStructs.ClasicStruct","name":"myStruct","type":"tuple","components":[{"internalType":"uint256","name":"one","type":"uint256"},{"internalType":"uint256","name":"two","type":"uint256"}]}]}],"name":"nestedStructExample","type":"function","stateMutability":"pure"}]`,
 		},
+		{
+			name:     "Enums",
+			contract: ReadContractFileForTest(t, "Enums").Content,
+			expected: `[{"inputs":[],"outputs":[{"internalType":"enum EnumContract.State","name":"","type":"uint8"}],"name":"state","type":"function","stateMutability":"view"},{"inputs":[],"type":"constructor"},{"inputs":[],"outputs":[{"internalType":"bool","name":"","type":"bool"}],"name":"isWaiting","type":"function","stateMutability":"view"},{"inputs":[],"outputs":[{"internalType":"bool","name":"","type":"bool"}],"name":"isReady","type":"function","stateMutability":"view"},{"inputs":[],"outputs":[{"internalType":"bool","name":"","type":"bool"}],"name":"isActive","type":"function","stateMutability":"view"},{"inputs":[],"outputs":[],"name":"makeReady","type":"function","stateMutability":"nonpayable"},{"inputs":[],"outputs":[],"name":"makeActive","type":"function","stateMutability":"nonpayable"},{"inputs":[],"outputs":[],"name":"reset","type":"function","stateMutability":"nonpayable"}]`,
+		},
 	}
 
 	for _, testCase := range testCases {

abi_parser.go

@@ -17,6 +17,8 @@ type AbiParser struct {
 	abi            ABI
 	contractName   string
 	definedStructs map[string]MethodIO
+	definedEnums   map[string]bool // map to keep track of defined enum types
+
 }
 
 // InjectConstructor injects a constructor definition into the ABI.
@@ -183,10 +185,12 @@ func (p *AbiParser) InjectFunction(ctx *parser.FunctionDefinitionContext) error
 // InjectStateVariable injects a state variable definition into the ABI.
 // It takes a StateVariableDeclarationContext (from the parser) as input.
 func (p *AbiParser) InjectStateVariable(ctx *parser.StateVariableDeclarationContext) error {
-	if isMappingType(ctx.TypeName().GetText()) {
+	typeName := ctx.TypeName().GetText()
+
+	if isMappingType(typeName) {
 		inputs := make([]MethodIO, 0)
 		outputs := make([]MethodIO, 0)
-		matched, inputList, outputList := parseMappingType(ctx.TypeName().GetText())
+		matched, inputList, outputList := parseMappingType(typeName)
 
 		if matched {
 			for _, input := range inputList {
@@ -212,14 +216,32 @@ func (p *AbiParser) InjectStateVariable(ctx *parser.StateVariableDeclarationCont
 			Type:            "function",
 		})
 		return nil
+	} else if isEnumType(p.definedEnums, typeName) {
+		p.abi = append(p.abi, MethodVariable{
+			Inputs: make([]MethodIO, 0),
+			Outputs: []MethodIO{
+				{
+					Type: "uint8", // enums are represented as uint8 in the ABI
+					InternalType: fmt.Sprintf(
+						"enum %s.%s",
+						p.contractName,
+						typeName,
+					),
+				},
+			},
+			Name:            ctx.Identifier().GetText(),
+			StateMutability: "view",
+			Type:            "function",
+		})
+		return nil
 	}
 
 	p.abi = append(p.abi, MethodVariable{
 		Inputs: make([]MethodIO, 0),
 		Outputs: []MethodIO{
 			{
-				Type:         normalizeTypeName(ctx.TypeName().GetText()),
-				InternalType: normalizeTypeName(ctx.TypeName().GetText()),
+				Type:         normalizeTypeName(typeName),
+				InternalType: normalizeTypeName(typeName),
 			},
 		},
 		Name:            ctx.Identifier().GetText(),

data/tests/Enums.sol

@@ -0,0 +1,45 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+contract EnumContract {
+    // Define an enum
+    enum State { Waiting, Ready, Active }
+
+    // Declare a state variable of type State
+    State public state;
+
+    // Initialize the state
+    constructor() {
+        state = State.Waiting;
+    }
+
+    // Function to check if state is Waiting
+    function isWaiting() public view returns(bool) {
+        return state == State.Waiting;
+    }
+
+    // Function to check if state is Ready
+    function isReady() public view returns(bool) {
+        return state == State.Ready;
+    }
+
+    // Function to check if state is Active
+    function isActive() public view returns(bool) {
+        return state == State.Active;
+    }
+
+    // Function to set state to Ready
+    function makeReady() public {
+        state = State.Ready;
+    }
+
+    // Function to set state to Active
+    function makeActive() public {
+        state = State.Active;
+    }
+
+    // Function to reset state to Waiting
+    function reset() public {
+        state = State.Waiting;
+    }
+}