Abstract Syntax Tree (AST) Assembly Support

ast/body.go

@@ -393,7 +393,7 @@ func (b *BodyNode) parseStatements(
 			unit, contractNode, fnNode, b, childCtx,
 		))
 	case *parser.AssemblyStatementContext:
-		statement := NewAssemblyStatement(b.ASTBuilder)
+		statement := NewYul(b.ASTBuilder)
 		b.Statements = append(b.Statements, statement.Parse(
 			unit, contractNode, fnNode, b, childCtx,
 		))

ast/builder_test.go

@@ -118,6 +118,8 @@ func TestAstBuilderFromSourceAsString(t *testing.T) {
 					assert.Equal(t, sourceUnit.GetName(), testCase.sources.EntrySourceUnitName)
 				}
 
+				assert.NotNil(t, astRoot.GetGlobalNodes())
+
 				assert.True(t, astRoot.HasSourceUnits())
 				assert.NotNil(t, astRoot.GetSourceUnitByName(sourceUnit.GetName()))
 				assert.NotNil(t, astRoot.GetSourceUnitById(sourceUnit.GetId()))

ast/pragma.go

@@ -108,16 +108,6 @@ func CreatePragmaFromCtx(b *ASTBuilder, unit *SourceUnit[Node[ast_pb.SourceUnit]
 	}
 }
 
-// EnterPragmaDirective is called when production pragmaDirective is entered.
-// However, it won't return pragma directives properly. For example, if we have
-// experimental pragma, it won't return it. It will return only the pragma.
-// Because of it, we are parsing pragmas in EnterSourceUnit to be able capture all of the
-// pragmas and assign them based on the contract they belong to. Source file can have multiple
-// contracts and multiple files and therefore we need to be able to assign pragmas to the
-// correct contract.
-// @WARN: DO NOT USE THIS METHOD.
-func (b *ASTBuilder) EnterPragmaDirective(ctx *parser.PragmaDirectiveContext) {}
-
 // FindPragmasForSourceUnit traverses the children of the provided source unit until
 // it finds the library, contract, or interface definition. It collects all pragma
 // directives encountered along the way and returns them as a slice of Node.

ast/unary_prefix.go

@@ -14,7 +14,7 @@ type UnaryPrefix struct {
 
 	Id                    int64            `json:"id"`
 	NodeType              ast_pb.NodeType  `json:"node_type"`
-	Kind                  NodeType         `json:"kind"`
+	Kind                  ast_pb.NodeType  `json:"kind"`
 	Src                   SrcNode          `json:"src"`
 	Operator              ast_pb.Operator  `json:"operator"`
 	Prefix                bool             `json:"prefix"`
@@ -55,7 +55,7 @@ func (u *UnaryPrefix) GetType() ast_pb.NodeType {
 }
 
 // GetKind returns the node type of the UnaryPrefix.
-func (u *UnaryPrefix) GetKind() NodeType {
+func (u *UnaryPrefix) GetKind() ast_pb.NodeType {
 	return u.Kind
 }
 
@@ -132,6 +132,12 @@ func (u *UnaryPrefix) UnmarshalJSON(data []byte) error {
 		}
 	}
 
+	if kind, ok := tempMap["kind"]; ok {
+		if err := json.Unmarshal(kind, &u.Kind); err != nil {
+			return err
+		}
+	}
+
 	if src, ok := tempMap["src"]; ok {
 		if err := json.Unmarshal(src, &u.Src); err != nil {
 			return err
@@ -214,6 +220,7 @@ func (u *UnaryPrefix) ToProto() NodeType {
 	proto := ast_pb.UnaryPrefix{
 		Id:                    u.GetId(),
 		NodeType:              u.GetType(),
+		Kind:                  u.GetKind(),
 		Src:                   u.GetSrc().ToProto(),
 		Operator:              u.GetOperator(),
 		Prefix:                u.GetPrefix(),

ast/unary_suffix.go

@@ -132,6 +132,12 @@ func (u *UnarySuffix) UnmarshalJSON(data []byte) error {
 		}
 	}
 
+	if kind, ok := tempMap["kind"]; ok {
+		if err := json.Unmarshal(kind, &u.Kind); err != nil {
+			return err
+		}
+	}
+
 	if src, ok := tempMap["src"]; ok {
 		if err := json.Unmarshal(src, &u.Src); err != nil {
 			return err
@@ -214,6 +220,7 @@ func (u *UnarySuffix) ToProto() NodeType {
 	proto := ast_pb.UnarySuffix{
 		Id:                    u.GetId(),
 		NodeType:              u.GetType(),
+		Kind:                  u.GetKind(),
 		Src:                   u.GetSrc().ToProto(),
 		Operator:              u.GetOperator(),
 		Prefix:                u.GetPrefix(),

ast/unmarshaler.go

@@ -390,11 +390,11 @@ func unmarshalNode(data []byte, nodeType ast_pb.NodeType) (Node[NodeType], error
 		return toReturn, nil
 
 	//
-	// FUTURE IMPLEMENTATIONS....
+	// ASSEMBLY / YUL NODES
 	//
 
 	case ast_pb.NodeType_ASSEMBLY_STATEMENT:
-		var toReturn *AssemblyStatement
+		var toReturn *Yul
 		if err := json.Unmarshal(data, &toReturn); err != nil {
 			return nil, err
 		}
@@ -405,6 +405,96 @@ func unmarshalNode(data []byte, nodeType ast_pb.NodeType) (Node[NodeType], error
 			return nil, err
 		}
 		return toReturn, nil
+	case ast_pb.NodeType_YUL_VARIABLE_DECLARATION:
+		var toReturn *YulAssignment
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_ASSIGNMENT:
+		var toReturn *YulAssignment
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_BLOCK:
+		var toReturn *YulBlockStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_FOR:
+		var toReturn *YulForStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_IF:
+		var toReturn *YulIfStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_SWITCH:
+		var toReturn *YulSwitchStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_SWITCH_CASE:
+		var toReturn *YulSwitchCaseStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_LITERAL:
+		var toReturn *YulLiteralStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_FUNCTION_CALL:
+		var toReturn *YulFunctionCallStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_EXPRESSION:
+		var toReturn *YulExpressionStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_FUNCTION_DEFINITION:
+		var toReturn *YulFunctionDefinition
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_BREAK:
+		var toReturn *YulBreakStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_CONTINUE:
+		var toReturn *YulContinueStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_LEAVE:
+		var toReturn *YulLeaveStatement
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
+	case ast_pb.NodeType_YUL_IDENTIFIER:
+		var toReturn *YulIdentifier
+		if err := json.Unmarshal(data, &toReturn); err != nil {
+			return nil, err
+		}
+		return toReturn, nil
 	default:
 		zap.L().Error(
 			"unknown node type while importing JSON",