Proper parsing of contract definitions

ast/contract.go

@@ -30,7 +30,13 @@ type Contract struct {
 // NewContractDefinition creates a new instance of Contract.
 func NewContractDefinition(b *ASTBuilder) *Contract {
 	return &Contract{
-		ASTBuilder: b,
+		ASTBuilder:              b,
+		NodeType:                ast_pb.NodeType_CONTRACT_DEFINITION,
+		Kind:                    ast_pb.NodeType_KIND_CONTRACT,
+		LinearizedBaseContracts: make([]int64, 0),
+		ContractDependencies:    make([]int64, 0),
+		BaseContracts:           make([]*BaseContract, 0),
+		Nodes:                   make([]Node[NodeType], 0),
 	}
 }
 

ast/interface.go

@@ -3,6 +3,7 @@
 package ast
 
 import (
+	"encoding/json"
 	"fmt"
 
 	v3 "github.com/cncf/xds/go/xds/type/v3"
@@ -32,9 +33,12 @@ type Interface struct {
 func NewInterfaceDefinition(b *ASTBuilder) *Interface {
 	return &Interface{
 		ASTBuilder:              b,
+		NodeType:                ast_pb.NodeType_CONTRACT_DEFINITION,
+		Kind:                    ast_pb.NodeType_KIND_INTERFACE,
 		LinearizedBaseContracts: make([]int64, 0),
 		ContractDependencies:    make([]int64, 0),
 		BaseContracts:           make([]*BaseContract, 0),
+		Nodes:                   make([]Node[NodeType], 0),
 	}
 }
 
@@ -223,6 +227,106 @@ func (l *Interface) GetReceive() *Receive {
 	return nil
 }
 
+func (l *Interface) UnmarshalJSON(data []byte) error {
+	var tempMap map[string]json.RawMessage
+	if err := json.Unmarshal(data, &tempMap); err != nil {
+		return err
+	}
+
+	if id := tempMap["id"]; id == nil {
+		if err := json.Unmarshal(tempMap["id"], &l.Id); err != nil {
+			return err
+		}
+	}
+
+	if name, ok := tempMap["name"]; ok {
+		if err := json.Unmarshal(name, &l.Name); err != nil {
+			return err
+		}
+	}
+
+	if nodeType, ok := tempMap["node_type"]; ok {
+		if err := json.Unmarshal(nodeType, &l.NodeType); err != nil {
+			return err
+		}
+	}
+
+	if src, ok := tempMap["src"]; ok {
+		if err := json.Unmarshal(src, &l.Src); err != nil {
+			return err
+		}
+	}
+
+	if nameLocation, ok := tempMap["name_location"]; ok {
+		if err := json.Unmarshal(nameLocation, &l.NameLocation); err != nil {
+			return err
+		}
+	}
+
+	if abstract, ok := tempMap["abstract"]; ok {
+		if err := json.Unmarshal(abstract, &l.Abstract); err != nil {
+			return err
+		}
+	}
+
+	if kind, ok := tempMap["kind"]; ok {
+		if err := json.Unmarshal(kind, &l.Kind); err != nil {
+			return err
+		}
+	}
+
+	if fullyImplemented, ok := tempMap["fully_implemented"]; ok {
+		if err := json.Unmarshal(fullyImplemented, &l.FullyImplemented); err != nil {
+			return err
+		}
+	}
+
+	if baseContracts, ok := tempMap["base_contracts"]; ok {
+		if err := json.Unmarshal(baseContracts, &l.BaseContracts); err != nil {
+			return err
+		}
+	}
+
+	if lbc, ok := tempMap["linearized_base_contracts"]; ok {
+		if err := json.Unmarshal(lbc, &l.LinearizedBaseContracts); err != nil {
+			return err
+		}
+	}
+
+	if cd, ok := tempMap["contract_dependencies"]; ok {
+		if err := json.Unmarshal(cd, &l.ContractDependencies); err != nil {
+			return err
+		}
+	}
+
+	if n, ok := tempMap["nodes"]; ok {
+		var nodes []json.RawMessage
+		if err := json.Unmarshal(n, &nodes); err != nil {
+			return err
+		}
+
+		for _, tempNode := range nodes {
+			var tempNodeMap map[string]json.RawMessage
+			if err := json.Unmarshal(tempNode, &tempNodeMap); err != nil {
+				return err
+			}
+
+			var tempNodeType ast_pb.NodeType
+			if err := json.Unmarshal(tempNodeMap["node_type"], &tempNodeType); err != nil {
+				return err
+			}
+
+			node, err := unmarshalNode(tempNode, tempNodeType)
+			if err != nil {
+				return err
+			}
+			l.Nodes = append(l.Nodes, node)
+		}
+	}
+
+	return nil
+}
+
 // ToProto converts the Interface node to its corresponding protocol buffer representation.
 func (l *Interface) ToProto() NodeType {
 	proto := ast_pb.Contract{

ast/library.go

@@ -1,6 +1,7 @@
 package ast
 
 import (
+	"encoding/json"
 	"fmt"
 	"reflect"
 
@@ -34,7 +35,13 @@ type Library struct {
 // It returns a pointer to the created Library.
 func NewLibraryDefinition(b *ASTBuilder) *Library {
 	return &Library{
-		ASTBuilder: b,
+		ASTBuilder:              b,
+		NodeType:                ast_pb.NodeType_CONTRACT_DEFINITION,
+		Kind:                    ast_pb.NodeType_KIND_LIBRARY,
+		LinearizedBaseContracts: make([]int64, 0),
+		ContractDependencies:    make([]int64, 0),
+		BaseContracts:           make([]*BaseContract, 0),
+		Nodes:                   make([]Node[NodeType], 0),
 	}
 }
 
@@ -228,6 +235,106 @@ func (l *Library) GetReceive() *Receive {
 	return nil
 }
 
+func (l *Library) UnmarshalJSON(data []byte) error {
+	var tempMap map[string]json.RawMessage
+	if err := json.Unmarshal(data, &tempMap); err != nil {
+		return err
+	}
+
+	if id := tempMap["id"]; id == nil {
+		if err := json.Unmarshal(tempMap["id"], &l.Id); err != nil {
+			return err
+		}
+	}
+
+	if name, ok := tempMap["name"]; ok {
+		if err := json.Unmarshal(name, &l.Name); err != nil {
+			return err
+		}
+	}
+
+	if nodeType, ok := tempMap["node_type"]; ok {
+		if err := json.Unmarshal(nodeType, &l.NodeType); err != nil {
+			return err
+		}
+	}
+
+	if src, ok := tempMap["src"]; ok {
+		if err := json.Unmarshal(src, &l.Src); err != nil {
+			return err
+		}
+	}
+
+	if nameLocation, ok := tempMap["name_location"]; ok {
+		if err := json.Unmarshal(nameLocation, &l.NameLocation); err != nil {
+			return err
+		}
+	}
+
+	if abstract, ok := tempMap["abstract"]; ok {
+		if err := json.Unmarshal(abstract, &l.Abstract); err != nil {
+			return err
+		}
+	}
+
+	if kind, ok := tempMap["kind"]; ok {
+		if err := json.Unmarshal(kind, &l.Kind); err != nil {
+			return err
+		}
+	}
+
+	if fullyImplemented, ok := tempMap["fully_implemented"]; ok {
+		if err := json.Unmarshal(fullyImplemented, &l.FullyImplemented); err != nil {
+			return err
+		}
+	}
+
+	if baseContracts, ok := tempMap["base_contracts"]; ok {
+		if err := json.Unmarshal(baseContracts, &l.BaseContracts); err != nil {
+			return err
+		}
+	}
+
+	if lbc, ok := tempMap["linearized_base_contracts"]; ok {
+		if err := json.Unmarshal(lbc, &l.LinearizedBaseContracts); err != nil {
+			return err
+		}
+	}
+
+	if cd, ok := tempMap["contract_dependencies"]; ok {
+		if err := json.Unmarshal(cd, &l.ContractDependencies); err != nil {
+			return err
+		}
+	}
+
+	if n, ok := tempMap["nodes"]; ok {
+		var nodes []json.RawMessage
+		if err := json.Unmarshal(n, &nodes); err != nil {
+			return err
+		}
+
+		for _, tempNode := range nodes {
+			var tempNodeMap map[string]json.RawMessage
+			if err := json.Unmarshal(tempNode, &tempNodeMap); err != nil {
+				return err
+			}
+
+			var tempNodeType ast_pb.NodeType
+			if err := json.Unmarshal(tempNodeMap["node_type"], &tempNodeType); err != nil {
+				return err
+			}
+
+			node, err := unmarshalNode(tempNode, tempNodeType)
+			if err != nil {
+				return err
+			}
+			l.Nodes = append(l.Nodes, node)
+		}
+	}
+
+	return nil
+}
+
 // ToProto converts the Library to a protocol buffer representation.
 // Currently, it returns an empty Contract and needs to be implemented.
 func (l *Library) ToProto() NodeType {

ast/unmarshaler.go

@@ -149,9 +149,45 @@ func unmarshalNode(data []byte, nodeType ast_pb.NodeType) (Node[NodeType], error
 		return toReturn, nil
 
 	case ast_pb.NodeType_CONTRACT_DEFINITION:
+
+		var tempMap map[string]json.RawMessage
+		if err := json.Unmarshal(data, &tempMap); err != nil {
+			return nil, err
+		}
+
+		if kind, ok := tempMap["kind"]; ok {
+			var fKind ast_pb.NodeType
+			if err := json.Unmarshal(kind, &fKind); err != nil {
+				return nil, err
+			}
+
+			switch fKind {
+			case ast_pb.NodeType_KIND_CONTRACT:
+				var toReturn *Contract
+				if err := json.Unmarshal(data, &toReturn); err != nil {
+					panic("contract definition")
+					return nil, err
+				}
+				return toReturn, nil
+			case ast_pb.NodeType_KIND_LIBRARY:
+				var toReturn *Library
+				if err := json.Unmarshal(data, &toReturn); err != nil {
+					panic("library definition")
+					return nil, err
+				}
+				return toReturn, nil
+			case ast_pb.NodeType_KIND_INTERFACE:
+				var toReturn *Interface
+				if err := json.Unmarshal(data, &toReturn); err != nil {
+					panic("interface definition")
+					return nil, err
+				}
+				return toReturn, nil
+			}
+		}
+
 		var toReturn *Contract
 		if err := json.Unmarshal(data, &toReturn); err != nil {
-			panic("contract definition")
 			return nil, err
 		}
 		return toReturn, nil