Mako/parser/ast.go

package parser

import "fmt"

// TypeInfo represents type information for expressions
type TypeInfo struct {
	Type     string // "number", "string", "bool", "table", "function", "nil", "any"
	Inferred bool   // true if type was inferred, false if explicitly declared
}

// 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()
	GetType() *TypeInfo
	SetType(*TypeInfo)
}

// 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 with optional type hint
type AssignStatement struct {
	Name          Expression // Changed from *Identifier to Expression for member access
	TypeHint      *TypeInfo  // optional type hint
	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 "
	}

	var nameStr string
	if as.TypeHint != nil {
		nameStr = fmt.Sprintf("%s: %s", as.Name.String(), as.TypeHint.Type)
	} else {
		nameStr = as.Name.String()
	}

	return fmt.Sprintf("%s%s = %s", prefix, nameStr, as.Value.String())
}

// ExpressionStatement represents expressions used as statements
type ExpressionStatement struct {
	Expression Expression
}

func (es *ExpressionStatement) statementNode() {}
func (es *ExpressionStatement) String() string {
	return es.Expression.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
}

// FunctionParameter represents a function parameter with optional type hint
type FunctionParameter struct {
	Name     string
	TypeHint *TypeInfo
}

func (fp *FunctionParameter) String() string {
	if fp.TypeHint != nil {
		return fmt.Sprintf("%s: %s", fp.Name, fp.TypeHint.Type)
	}
	return fp.Name
}

// Identifier represents identifiers
type Identifier struct {
	Value    string
	typeInfo *TypeInfo
}

func (i *Identifier) expressionNode()     {}
func (i *Identifier) String() string      { return i.Value }
func (i *Identifier) GetType() *TypeInfo  { return i.typeInfo }
func (i *Identifier) SetType(t *TypeInfo) { i.typeInfo = t }

// NumberLiteral represents numeric literals
type NumberLiteral struct {
	Value    float64
	typeInfo *TypeInfo
}

func (nl *NumberLiteral) expressionNode()     {}
func (nl *NumberLiteral) String() string      { return fmt.Sprintf("%.2f", nl.Value) }
func (nl *NumberLiteral) GetType() *TypeInfo  { return nl.typeInfo }
func (nl *NumberLiteral) SetType(t *TypeInfo) { nl.typeInfo = t }

// StringLiteral represents string literals
type StringLiteral struct {
	Value    string
	typeInfo *TypeInfo
}

func (sl *StringLiteral) expressionNode()     {}
func (sl *StringLiteral) String() string      { return fmt.Sprintf(`"%s"`, sl.Value) }
func (sl *StringLiteral) GetType() *TypeInfo  { return sl.typeInfo }
func (sl *StringLiteral) SetType(t *TypeInfo) { sl.typeInfo = t }

// BooleanLiteral represents boolean literals
type BooleanLiteral struct {
	Value    bool
	typeInfo *TypeInfo
}

func (bl *BooleanLiteral) expressionNode() {}
func (bl *BooleanLiteral) String() string {
	if bl.Value {
		return "true"
	}
	return "false"
}
func (bl *BooleanLiteral) GetType() *TypeInfo  { return bl.typeInfo }
func (bl *BooleanLiteral) SetType(t *TypeInfo) { bl.typeInfo = t }

// NilLiteral represents nil literal
type NilLiteral struct {
	typeInfo *TypeInfo
}

func (nl *NilLiteral) expressionNode()     {}
func (nl *NilLiteral) String() string      { return "nil" }
func (nl *NilLiteral) GetType() *TypeInfo  { return nl.typeInfo }
func (nl *NilLiteral) SetType(t *TypeInfo) { nl.typeInfo = t }

// FunctionLiteral represents function literals with typed parameters
type FunctionLiteral struct {
	Parameters []FunctionParameter
	Variadic   bool
	ReturnType *TypeInfo // optional return type hint
	Body       []Statement
	typeInfo   *TypeInfo
}

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

	result := fmt.Sprintf("fn(%s)", params)
	if fl.ReturnType != nil {
		result += ": " + fl.ReturnType.Type
	}
	result += "\n"

	for _, stmt := range fl.Body {
		result += "\t" + stmt.String() + "\n"
	}
	result += "end"
	return result
}
func (fl *FunctionLiteral) GetType() *TypeInfo  { return fl.typeInfo }
func (fl *FunctionLiteral) SetType(t *TypeInfo) { fl.typeInfo = t }

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

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, ", "))
}
func (ce *CallExpression) GetType() *TypeInfo  { return ce.typeInfo }
func (ce *CallExpression) SetType(t *TypeInfo) { ce.typeInfo = t }

// PrefixExpression represents prefix operations like -x, not x
type PrefixExpression struct {
	Operator string
	Right    Expression
	typeInfo *TypeInfo
}

func (pe *PrefixExpression) expressionNode() {}
func (pe *PrefixExpression) String() string {
	// Add space for word operators
	if pe.Operator == "not" {
		return fmt.Sprintf("(%s %s)", pe.Operator, pe.Right.String())
	}
	return fmt.Sprintf("(%s%s)", pe.Operator, pe.Right.String())
}
func (pe *PrefixExpression) GetType() *TypeInfo  { return pe.typeInfo }
func (pe *PrefixExpression) SetType(t *TypeInfo) { pe.typeInfo = t }

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

func (ie *InfixExpression) expressionNode() {}
func (ie *InfixExpression) String() string {
	return fmt.Sprintf("(%s %s %s)", ie.Left.String(), ie.Operator, ie.Right.String())
}
func (ie *InfixExpression) GetType() *TypeInfo  { return ie.typeInfo }
func (ie *InfixExpression) SetType(t *TypeInfo) { ie.typeInfo = t }

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

func (ie *IndexExpression) expressionNode() {}
func (ie *IndexExpression) String() string {
	return fmt.Sprintf("%s[%s]", ie.Left.String(), ie.Index.String())
}
func (ie *IndexExpression) GetType() *TypeInfo  { return ie.typeInfo }
func (ie *IndexExpression) SetType(t *TypeInfo) { ie.typeInfo = t }

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

func (de *DotExpression) expressionNode() {}
func (de *DotExpression) String() string {
	return fmt.Sprintf("%s.%s", de.Left.String(), de.Key)
}
func (de *DotExpression) GetType() *TypeInfo  { return de.typeInfo }
func (de *DotExpression) SetType(t *TypeInfo) { de.typeInfo = t }

// 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
	typeInfo *TypeInfo
}

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, ", "))
}
func (tl *TableLiteral) GetType() *TypeInfo  { return tl.typeInfo }
func (tl *TableLiteral) SetType(t *TypeInfo) { tl.typeInfo = t }

// 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
}