Print ast to Solidity Source Code

ast/if.go

@@ -186,6 +186,9 @@ func (i *IfStatement) Parse(
 				body.ParseBlock(unit, contractNode, fnNode, statementCtx.Block())
 				break
 			}
+			// Edge case for single statement if
+			// Eg: if (a) b;
+			body.parseStatements(unit, contractNode, fnNode, statementCtx.GetChild(0))
 
 			i.Body = body
 		}

ast/parameter.go

@@ -438,5 +438,5 @@ func (p *Parameter) getStorageLocationFromCtx(ctx *parser.ParameterDeclarationCo
 		}
 	}
 
-	return ast_pb.StorageLocation_MEMORY
+	return ast_pb.StorageLocation_DEFAULT
 }

ast/source_unit.go

@@ -2,10 +2,11 @@ package ast
 
 import (
 	"fmt"
-	"github.com/goccy/go-json"
 	"path/filepath"
 	"regexp"
 
+	"github.com/goccy/go-json"
+
 	v3 "github.com/cncf/xds/go/xds/type/v3"
 	ast_pb "github.com/unpackdev/protos/dist/go/ast"
 	"github.com/unpackdev/solgo"

ast/state_variable.go

@@ -22,6 +22,7 @@ type StateVariableDeclaration struct {
 	Visibility      ast_pb.Visibility      `json:"visibility"`        // Visibility of the state variable declaration
 	StorageLocation ast_pb.StorageLocation `json:"storage_location"`  // Storage location of the state variable declaration
 	StateMutability ast_pb.Mutability      `json:"mutability"`        // State mutability of the state variable declaration
+	Override        bool                   `json:"is_override"`       // Indicates if the state variable is an override
 	TypeName        *TypeName              `json:"type_name"`         // Type name of the state variable
 	InitialValue    Node[NodeType]         `json:"initial_value"`     // Initial value of the state variable
 }
@@ -207,6 +208,8 @@ func (v *StateVariableDeclaration) Parse(
 		v.Constant = constantCtx != nil
 	}
 
+	v.Override = ctx.GetOverrideSpecifierSet()
+
 	typeName := NewTypeName(v.ASTBuilder)
 
 	typeName.Parse(unit, nil, v.Id, ctx.GetType_())
@@ -270,6 +273,8 @@ func (v *StateVariableDeclaration) ParseGlobal(
 		v.Constant = constantCtx != nil
 	}
 
+	v.Override = ctx.GetOverrideSpecifierSet()
+
 	typeName := NewTypeName(v.ASTBuilder)
 	typeName.Parse(nil, nil, v.Id, ctx.GetType_())
 	v.TypeName = typeName

ast/storage.go

@@ -57,7 +57,6 @@ func (t *TypeName) StorageSize() (int64, bool) {
 	// Add cases for other node types like struct, enum, etc., as needed.
 	default:
 		panic(fmt.Sprintf("Unhandled node type @ StorageSize: %s", t.NodeType))
-		return 0, false // Type not recognized or not handled yet.
 	}
 }
 

printer/ast_printer/and_operation.go

@@ -0,0 +1,19 @@
+package ast_printer
+
+import (
+	"strings"
+
+	"github.com/unpackdev/solgo/ast"
+)
+
+func printAndOperation(node *ast.AndOperation, sb *strings.Builder, depth int) bool {
+	expressions := []string{}
+	success := true
+	for _, exp := range node.GetExpressions() {
+		s, ok := Print(exp)
+		success = ok && success
+		expressions = append(expressions, s)
+	}
+	writeSeperatedList(sb, " && ", expressions)
+	return success
+}

printer/ast_printer/assignment.go

@@ -0,0 +1,65 @@
+package ast_printer
+
+import (
+	"strings"
+
+	ast_pb "github.com/unpackdev/protos/dist/go/ast"
+	"github.com/unpackdev/solgo/ast"
+)
+
+func getAssignOperatorString(op ast_pb.Operator) string {
+	switch op {
+	case ast_pb.Operator_EQUAL:
+		return "="
+	case ast_pb.Operator_PLUS_EQUAL:
+		return "+="
+	case ast_pb.Operator_MINUS_EQUAL:
+		return "-="
+	case ast_pb.Operator_MUL_EQUAL:
+		return "*="
+	case ast_pb.Operator_DIV_EQUAL:
+		return "/="
+	case ast_pb.Operator_MOD_EQUAL:
+		return "%="
+	case ast_pb.Operator_AND_EQUAL:
+		return "&="
+	case ast_pb.Operator_OR_EQUAL:
+		return "|="
+	case ast_pb.Operator_XOR_EQUAL:
+		return "^="
+	case ast_pb.Operator_SHIFT_LEFT_EQUAL:
+		return "<<="
+	case ast_pb.Operator_SHIFT_RIGHT_EQUAL:
+		return ">>="
+	case ast_pb.Operator_BIT_AND_EQUAL:
+		return "&="
+	case ast_pb.Operator_BIT_OR_EQUAL:
+		return "|="
+	case ast_pb.Operator_BIT_XOR_EQUAL:
+		return "^="
+	case ast_pb.Operator_POW_EQUAL:
+		return "**="
+	default:
+		return ""
+	}
+}
+
+func printAssignment(node *ast.Assignment, sb *strings.Builder, depth int) bool {
+	success := true
+	if node.Expression != nil {
+		return PrintRecursive(node.Expression, sb, depth)
+	}
+	if node.LeftExpression == nil || node.RightExpression == nil {
+		return false
+	}
+	op := getAssignOperatorString(node.Operator)
+	if op == "" {
+		success = false
+	}
+	success = PrintRecursive(node.LeftExpression, sb, depth) && success
+	sb.WriteString(" ")
+	sb.WriteString(op)
+	sb.WriteString(" ")
+	success = PrintRecursive(node.RightExpression, sb, depth) && success
+	return success
+}

printer/ast_printer/ast_printer.go

@@ -0,0 +1,195 @@
+package ast_printer
+
+import (
+	"strings"
+
+	ast_pb "github.com/unpackdev/protos/dist/go/ast"
+	"github.com/unpackdev/solgo/ast"
+	"go.uber.org/zap"
+)
+
+const INDENT_SIZE = 2
+
+// Print is a function that prints the AST nodes to source code
+func Print(node ast.Node[ast.NodeType]) (string, bool) {
+	sb := strings.Builder{}
+	success := PrintRecursive(node, &sb, 0)
+	return sb.String(), success
+}
+
+// PrintRecursive is a function that prints the AST nodes to source code recursively
+func PrintRecursive(node ast.Node[ast.NodeType], sb *strings.Builder, depth int) bool {
+	if node == nil {
+		zap.S().Error("Node is nil")
+		return false
+	}
+	switch node := node.(type) {
+	case *ast.AndOperation:
+		return printAndOperation(node, sb, depth)
+	case *ast.BodyNode:
+		return printBody(node, sb, depth)
+	case *ast.Conditional:
+		return printConditional(node, sb, depth)
+	case *ast.Constructor:
+		return printConstructor(node, sb, depth)
+	case *ast.Pragma:
+		return printPragma(node, sb, depth)
+	case *ast.Contract:
+		return printContract(node, sb, depth)
+	case *ast.Function:
+		return printFunction(node, sb, depth)
+	case *ast.Parameter:
+		return printParameter(node, sb, depth)
+	case *ast.Assignment:
+		return printAssignment(node, sb, depth)
+	case *ast.TypeName:
+		return printTypeName(node, sb, depth)
+	case *ast.BinaryOperation:
+		return printBinaryOperation(node, sb, depth)
+	case *ast.StateVariableDeclaration:
+		return printStateVariableDeclaration(node, sb, depth)
+	case *ast.Emit:
+		return printEmit(node, sb, depth)
+	case *ast.ForStatement:
+		return printFor(node, sb, depth)
+	case *ast.PrimaryExpression:
+		return printPrimaryExpression(node, sb, depth)
+	case *ast.FunctionCall:
+		return printFunctionCall(node, sb, depth)
+	case *ast.Import:
+		return printImport(node, sb, depth)
+	case *ast.MemberAccessExpression:
+		return printMemberAccessExpression(node, sb, depth)
+	case *ast.VariableDeclaration:
+		return printVariableDeclaration(node, sb, depth)
+	case *ast.Declaration:
+		return printDeclaration(node, sb, depth)
+	case *ast.UnaryPrefix:
+		return printUnaryPrefix(node, sb, depth)
+	case *ast.UnarySuffix:
+		return printUnarySuffix(node, sb, depth)
+	case *ast.IndexAccess:
+		return printIndexAccess(node, sb, depth)
+	case *ast.ReturnStatement:
+		return printReturn(node, sb, depth)
+	case *ast.TupleExpression:
+		return printTupleExpression(node, sb, depth)
+	case *ast.StructDefinition:
+		return printStructDefinition(node, sb, depth)
+	case *ast.IfStatement:
+		return printIfStatement(node, sb, depth)
+	case *ast.EnumDefinition:
+		return printEnumDefinition(node, sb, depth)
+	case *ast.ModifierDefinition:
+		return printModifierDefinition(node, sb, depth)
+	case *ast.EventDefinition:
+		return printEventDefinition(node, sb, depth)
+	case *ast.ErrorDefinition:
+		return printErrorDefinition(node, sb, depth)
+	case *ast.PayableConversion:
+		return printPayableConversion(node, sb, depth)
+	case *ast.RevertStatement:
+		return printRevertStatement(node, sb, depth)
+	case *ast.ContinueStatement:
+		return printContinueStatement(node, sb, depth)
+	case *ast.InlineArray:
+		return printInlineArray(node, sb, depth)
+	default:
+		if node.GetType() == ast_pb.NodeType_SOURCE_UNIT {
+			return printSourceUnit(node, sb, depth)
+		}
+		zap.S().Errorf("Unknown node type: %T\n", node)
+		return false
+	}
+}
+
+func writeSeperatedStrings(sb *strings.Builder, seperator string, s ...string) {
+	count := 0
+	for _, item := range s {
+		// Skip empty strings
+		if item == "" {
+			continue
+		}
+
+		if count > 0 {
+			sb.WriteString(seperator)
+			sb.WriteString(item)
+		} else {
+			sb.WriteString(item)
+		}
+		count++
+	}
+}
+
+func writeSeperatedList(sb *strings.Builder, seperator string, s []string) {
+	count := 0
+	for _, item := range s {
+		// Skip empty strings
+		if item == "" {
+			continue
+		}
+
+		if count > 0 {
+			sb.WriteString(seperator)
+			sb.WriteString(item)
+		} else {
+			sb.WriteString(item)
+		}
+		count++
+	}
+}
+
+func writeStrings(sb *strings.Builder, s ...string) {
+	for _, item := range s {
+		sb.WriteString(item)
+	}
+}
+
+func indentString(s string, depth int) string {
+	return strings.Repeat(" ", depth*INDENT_SIZE) + s
+}
+
+func getStorageLocationString(storage ast_pb.StorageLocation) string {
+	switch storage {
+	case ast_pb.StorageLocation_DEFAULT:
+		return ""
+	case ast_pb.StorageLocation_MEMORY:
+		return "memory"
+	case ast_pb.StorageLocation_STORAGE:
+		return "storage"
+	case ast_pb.StorageLocation_CALLDATA:
+		return "calldata"
+	default:
+		return ""
+	}
+}
+
+func getVisibilityString(visibility ast_pb.Visibility) string {
+	switch visibility {
+	case ast_pb.Visibility_INTERNAL:
+		return "internal"
+	case ast_pb.Visibility_PUBLIC:
+		return "public"
+	case ast_pb.Visibility_EXTERNAL:
+		return "external"
+	case ast_pb.Visibility_PRIVATE:
+		return "private"
+	default:
+		return ""
+	}
+}
+
+func getStateMutabilityString(mut ast_pb.Mutability) string {
+	switch mut {
+	case ast_pb.Mutability_PURE:
+		return "pure"
+	case ast_pb.Mutability_VIEW:
+		return "view"
+	case ast_pb.Mutability_NONPAYABLE:
+		return ""
+	case ast_pb.Mutability_PAYABLE:
+		return "payable"
+	default:
+		return ""
+	}
+}