Mako/parser/ast.go

package parser

import "fmt"

// Node represents any node in the AST
type Node interface {
	String() string
}

// Statement represents statement nodes
type Statement interface {
	Node
	statementNode()
}

// Expression represents expression nodes
type Expression interface {
	Node
	expressionNode()
}

// Program represents the root of the AST
type Program struct {
	Statements []Statement
	ExitCode   int
}

func (p *Program) String() string {
	var result string
	for _, stmt := range p.Statements {
		result += stmt.String() + "\n"
	}
	return result
}

// AssignStatement represents variable assignment
type AssignStatement struct {
	Name          Expression // Changed from *Identifier to Expression for member access
	Value         Expression
	IsDeclaration bool // true if this is the first assignment in current scope
}

func (as *AssignStatement) statementNode() {}
func (as *AssignStatement) String() string {
	prefix := ""
	if as.IsDeclaration {
		prefix = "local "
	}
	return fmt.Sprintf("%s%s = %s", prefix, as.Name.String(), as.Value.String())
}

// EchoStatement represents echo output statements
type EchoStatement struct {
	Value Expression
}

func (es *EchoStatement) statementNode() {}
func (es *EchoStatement) String() string {
	return fmt.Sprintf("echo %s", es.Value.String())
}

// BreakStatement represents break statements to exit loops
type BreakStatement struct{}

func (bs *BreakStatement) statementNode() {}
func (bs *BreakStatement) String() string {
	return "break"
}

// ExitStatement represents exit statements to quit the script
type ExitStatement struct {
	Value Expression // optional, can be nil
}

func (es *ExitStatement) statementNode() {}
func (es *ExitStatement) String() string {
	if es.Value == nil {
		return "exit"
	}
	return fmt.Sprintf("exit %s", es.Value.String())
}

// ReturnStatement represents return statements
type ReturnStatement struct {
	Value Expression // optional, can be nil
}

func (rs *ReturnStatement) statementNode() {}
func (rs *ReturnStatement) String() string {
	if rs.Value == nil {
		return "return"
	}
	return fmt.Sprintf("return %s", rs.Value.String())
}

// ElseIfClause represents an elseif condition
type ElseIfClause struct {
	Condition Expression
	Body      []Statement
}

func (eic *ElseIfClause) String() string {
	var body string
	for _, stmt := range eic.Body {
		body += "\t" + stmt.String() + "\n"
	}
	return fmt.Sprintf("elseif %s then\n%s", eic.Condition.String(), body)
}

// IfStatement represents conditional statements
type IfStatement struct {
	Condition Expression
	Body      []Statement
	ElseIfs   []ElseIfClause
	Else      []Statement
}

func (is *IfStatement) statementNode() {}
func (is *IfStatement) String() string {
	var result string

	// If clause
	result += fmt.Sprintf("if %s then\n", is.Condition.String())
	for _, stmt := range is.Body {
		result += "\t" + stmt.String() + "\n"
	}

	// ElseIf clauses
	for _, elseif := range is.ElseIfs {
		result += elseif.String()
	}

	// Else clause
	if len(is.Else) > 0 {
		result += "else\n"
		for _, stmt := range is.Else {
			result += "\t" + stmt.String() + "\n"
		}
	}

	result += "end"
	return result
}

// WhileStatement represents while loops: while condition do ... end
type WhileStatement struct {
	Condition Expression
	Body      []Statement
}

func (ws *WhileStatement) statementNode() {}
func (ws *WhileStatement) String() string {
	var result string
	result += fmt.Sprintf("while %s do\n", ws.Condition.String())

	for _, stmt := range ws.Body {
		result += "\t" + stmt.String() + "\n"
	}

	result += "end"
	return result
}

// ForStatement represents numeric for loops: for i = start, end, step do ... end
type ForStatement struct {
	Variable *Identifier
	Start    Expression
	End      Expression
	Step     Expression // optional, nil means step of 1
	Body     []Statement
}

func (fs *ForStatement) statementNode() {}
func (fs *ForStatement) String() string {
	var result string
	if fs.Step != nil {
		result += fmt.Sprintf("for %s = %s, %s, %s do\n",
			fs.Variable.String(), fs.Start.String(), fs.End.String(), fs.Step.String())
	} else {
		result += fmt.Sprintf("for %s = %s, %s do\n",
			fs.Variable.String(), fs.Start.String(), fs.End.String())
	}

	for _, stmt := range fs.Body {
		result += "\t" + stmt.String() + "\n"
	}

	result += "end"
	return result
}

// ForInStatement represents iterator for loops: for k, v in expr do ... end
type ForInStatement struct {
	Key      *Identifier // optional, nil for single variable iteration
	Value    *Identifier
	Iterable Expression
	Body     []Statement
}

func (fis *ForInStatement) statementNode() {}
func (fis *ForInStatement) String() string {
	var result string
	if fis.Key != nil {
		result += fmt.Sprintf("for %s, %s in %s do\n",
			fis.Key.String(), fis.Value.String(), fis.Iterable.String())
	} else {
		result += fmt.Sprintf("for %s in %s do\n",
			fis.Value.String(), fis.Iterable.String())
	}

	for _, stmt := range fis.Body {
		result += "\t" + stmt.String() + "\n"
	}

	result += "end"
	return result
}

// Identifier represents identifiers
type Identifier struct {
	Value string
}

func (i *Identifier) expressionNode() {}
func (i *Identifier) String() string  { return i.Value }

// NumberLiteral represents numeric literals
type NumberLiteral struct {
	Value float64
}

func (nl *NumberLiteral) expressionNode() {}
func (nl *NumberLiteral) String() string  { return fmt.Sprintf("%.2f", nl.Value) }

// StringLiteral represents string literals
type StringLiteral struct {
	Value string
}

func (sl *StringLiteral) expressionNode() {}
func (sl *StringLiteral) String() string  { return fmt.Sprintf(`"%s"`, sl.Value) }

// BooleanLiteral represents boolean literals
type BooleanLiteral struct {
	Value bool
}

func (bl *BooleanLiteral) expressionNode() {}
func (bl *BooleanLiteral) String() string {
	if bl.Value {
		return "true"
	}
	return "false"
}

// NilLiteral represents nil literal
type NilLiteral struct{}

func (nl *NilLiteral) expressionNode() {}
func (nl *NilLiteral) String() string  { return "nil" }

// FunctionLiteral represents function literals: fn(a, b, ...) ... end
type FunctionLiteral struct {
	Parameters []string
	Variadic   bool
	Body       []Statement
}

func (fl *FunctionLiteral) expressionNode() {}
func (fl *FunctionLiteral) String() string {
	var params string
	for i, param := range fl.Parameters {
		if i > 0 {
			params += ", "
		}
		params += param
	}
	if fl.Variadic {
		if len(fl.Parameters) > 0 {
			params += ", "
		}
		params += "..."
	}

	result := fmt.Sprintf("fn(%s)\n", params)
	for _, stmt := range fl.Body {
		result += "\t" + stmt.String() + "\n"
	}
	result += "end"
	return result
}

// CallExpression represents function calls: func(arg1, arg2, ...)
type CallExpression struct {
	Function  Expression
	Arguments []Expression
}

func (ce *CallExpression) expressionNode() {}
func (ce *CallExpression) String() string {
	var args []string
	for _, arg := range ce.Arguments {
		args = append(args, arg.String())
	}
	return fmt.Sprintf("%s(%s)", ce.Function.String(), joinStrings(args, ", "))
}

// PrefixExpression represents prefix operations like -x
type PrefixExpression struct {
	Operator string
	Right    Expression
}

func (pe *PrefixExpression) expressionNode() {}
func (pe *PrefixExpression) String() string {
	return fmt.Sprintf("(%s%s)", pe.Operator, pe.Right.String())
}

// InfixExpression represents binary operations
type InfixExpression struct {
	Left     Expression
	Operator string
	Right    Expression
}

func (ie *InfixExpression) expressionNode() {}
func (ie *InfixExpression) String() string {
	return fmt.Sprintf("(%s %s %s)", ie.Left.String(), ie.Operator, ie.Right.String())
}

// IndexExpression represents table[key] access
type IndexExpression struct {
	Left  Expression
	Index Expression
}

func (ie *IndexExpression) expressionNode() {}
func (ie *IndexExpression) String() string {
	return fmt.Sprintf("%s[%s]", ie.Left.String(), ie.Index.String())
}

// DotExpression represents table.key access
type DotExpression struct {
	Left Expression
	Key  string
}

func (de *DotExpression) expressionNode() {}
func (de *DotExpression) String() string {
	return fmt.Sprintf("%s.%s", de.Left.String(), de.Key)
}

// TablePair represents a key-value pair in a table
type TablePair struct {
	Key   Expression // nil for array-style elements
	Value Expression
}

func (tp *TablePair) String() string {
	if tp.Key == nil {
		return tp.Value.String()
	}
	return fmt.Sprintf("%s = %s", tp.Key.String(), tp.Value.String())
}

// TableLiteral represents table literals {}
type TableLiteral struct {
	Pairs []TablePair
}

func (tl *TableLiteral) expressionNode() {}
func (tl *TableLiteral) String() string {
	var pairs []string
	for _, pair := range tl.Pairs {
		pairs = append(pairs, pair.String())
	}
	return fmt.Sprintf("{%s}", joinStrings(pairs, ", "))
}

// IsArray returns true if this table contains only array-style elements
func (tl *TableLiteral) IsArray() bool {
	for _, pair := range tl.Pairs {
		if pair.Key != nil {
			return false
		}
	}
	return true
}

// joinStrings joins string slice with separator
func joinStrings(strs []string, sep string) string {
	if len(strs) == 0 {
		return ""
	}
	if len(strs) == 1 {
		return strs[0]
	}

	var result string
	for i, s := range strs {
		if i > 0 {
			result += sep
		}
		result += s
	}
	return result
}