Mako/parser/ast.go

package parser

import "fmt"

// Note: Type definitions moved to types.go for proper separation of concerns

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

// Statement represents statement nodes that can appear at the top level or in blocks
type Statement interface {
	Node
	statementNode()
}

// Expression represents expression nodes that produce values and have types
type Expression interface {
	Node
	expressionNode()
	TypeInfo() TypeInfo // Returns type by value, not pointer
}

// Program represents the root of the AST containing all top-level statements.
// Tracks exit code for script termination and owns the statement list.
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
}

// StructField represents a field definition within a struct.
// Contains field name and required type annotation for compile-time checking.
type StructField struct {
	Name     string
	TypeHint TypeInfo // Required for struct fields, embeds directly
}

func (sf *StructField) String() string {
	return fmt.Sprintf("%s: %s", sf.Name, typeToString(sf.TypeHint))
}

// StructStatement represents struct type definitions with named fields.
// Defines new types that can be instantiated and used for type checking.
type StructStatement struct {
	Name   string
	Fields []StructField
	ID     uint16 // Unique identifier for fast lookup
}

func (ss *StructStatement) statementNode() {}
func (ss *StructStatement) String() string {
	var fields string
	for i, field := range ss.Fields {
		if i > 0 {
			fields += ",\n\t"
		}
		fields += field.String()
	}
	return fmt.Sprintf("struct %s {\n\t%s\n}", ss.Name, fields)
}

// MethodDefinition represents method definitions attached to struct types.
// Links a function implementation to a specific struct via struct ID.
type MethodDefinition struct {
	StructID   uint16 // Index into struct table for fast lookup
	MethodName string
	Function   *FunctionLiteral
}

func (md *MethodDefinition) statementNode() {}
func (md *MethodDefinition) String() string {
	return fmt.Sprintf("fn <struct>.%s%s", md.MethodName, md.Function.String()[2:])
}

// StructConstructor represents struct instantiation with field initialization.
// Uses struct ID for fast type resolution and validation during parsing.
type StructConstructor struct {
	StructID uint16      // Index into struct table
	Fields   []TablePair // Reuses table pair structure for field assignments
	typeInfo TypeInfo    // Cached type info for this constructor
}

func (sc *StructConstructor) expressionNode() {}
func (sc *StructConstructor) String() string {
	var pairs []string
	for _, pair := range sc.Fields {
		pairs = append(pairs, pair.String())
	}
	return fmt.Sprintf("<struct>{%s}", joinStrings(pairs, ", "))
}
func (sc *StructConstructor) TypeInfo() TypeInfo { return sc.typeInfo }

// Assignment represents both variable assignment statements and assignment expressions.
// Unified design reduces AST node count and simplifies type checking logic.
type Assignment struct {
	Target        Expression // Target (identifier, dot, or index expression)
	Value         Expression // Value being assigned
	TypeHint      TypeInfo   // Optional explicit type hint, embeds directly
	IsDeclaration bool       // True if declaring new variable in current scope
	IsExpression  bool       // True if used as expression (wrapped in parentheses)
}

func (a *Assignment) statementNode()  {}
func (a *Assignment) expressionNode() {}
func (a *Assignment) String() string {
	prefix := ""
	if a.IsDeclaration {
		prefix = "local "
	}

	var nameStr string
	if a.TypeHint.Type != TypeUnknown {
		nameStr = fmt.Sprintf("%s: %s", a.Target.String(), typeToString(a.TypeHint))
	} else {
		nameStr = a.Target.String()
	}

	result := fmt.Sprintf("%s%s = %s", prefix, nameStr, a.Value.String())
	if a.IsExpression {
		return "(" + result + ")"
	}
	return result
}
func (a *Assignment) TypeInfo() TypeInfo { return a.Value.TypeInfo() }

// ExpressionStatement wraps expressions used as statements.
// Allows function calls and other expressions at statement level.
type ExpressionStatement struct {
	Expression Expression
}

func (es *ExpressionStatement) statementNode() {}
func (es *ExpressionStatement) String() string {
	return es.Expression.String()
}

// EchoStatement represents output statements for displaying values.
// Simple debugging and output mechanism built into the language.
type EchoStatement struct {
	Value Expression
}

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

// BreakStatement represents loop exit statements.
// Simple marker node with no additional data needed.
type BreakStatement struct{}

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

// ExitStatement represents script termination with optional exit code.
// Value expression is nil for plain "exit", non-nil for "exit <code>".
type ExitStatement struct {
	Value Expression // Optional exit code expression
}

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 function return with optional value.
// Value expression is nil for plain "return", non-nil for "return <value>".
type ReturnStatement struct {
	Value Expression // Optional return value expression
}

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 conditional branches in if statements.
// Contains condition expression and body statements for this branch.
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 execution with optional elseif and else branches.
// Supports multiple elseif clauses and an optional final else clause.
type IfStatement struct {
	Condition Expression     // Main condition
	Body      []Statement    // Statements to execute if condition is true
	ElseIfs   []ElseIfClause // Optional elseif branches
	Else      []Statement    // Optional else branch
}

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

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

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

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

	result += "end"
	return result
}

// WhileStatement represents condition-based loops that execute while condition is true.
// Contains condition expression and body statements to repeat.
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 with start, end, and optional step.
// Variable is automatically scoped to the loop body.
type ForStatement struct {
	Variable *Identifier // Loop variable (automatically number type)
	Start    Expression  // Starting value expression
	End      Expression  // Ending value expression
	Step     Expression  // Optional step expression (nil means step of 1)
	Body     []Statement // Loop body statements
}

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-based loops over tables, arrays, or other iterables.
// Supports both single variable (for v in iter) and key-value (for k,v in iter) forms.
type ForInStatement struct {
	Key      *Identifier // Optional key variable (nil for single variable iteration)
	Value    *Identifier // Value variable (required)
	Iterable Expression  // Expression to iterate over
	Body     []Statement // Loop body statements
}

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 parameter in function definitions.
// Contains parameter name and optional type hint for type checking.
type FunctionParameter struct {
	Name     string
	TypeHint TypeInfo // Optional type constraint, embeds directly
}

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

// Identifier represents variable references and names.
// Stores resolved type information for efficient type checking.
type Identifier struct {
	Value    string
	typeInfo TypeInfo // Resolved type, embeds directly
}

func (i *Identifier) expressionNode() {}
func (i *Identifier) String() string  { return i.Value }
func (i *Identifier) TypeInfo() TypeInfo {
	if i.typeInfo.Type == TypeUnknown {
		return AnyType
	}
	return i.typeInfo
}

// NumberLiteral represents numeric constants including integers, floats, hex, and binary.
// Always has number type, so no additional type storage needed.
type NumberLiteral struct {
	Value float64 // All numbers stored as float64 for simplicity
}

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

// StringLiteral represents string constants and multiline strings.
// Always has string type, so no additional type storage needed.
type StringLiteral struct {
	Value string // String content without quotes
}

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

// BooleanLiteral represents true and false constants.
// Always has bool type, so no additional type storage needed.
type BooleanLiteral struct {
	Value bool
}

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

// NilLiteral represents the nil constant value.
// Always has nil type, so no additional type storage needed.
type NilLiteral struct{}

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

// FunctionLiteral represents function definitions with parameters, body, and optional return type.
// Always has function type, stores additional return type information separately.
type FunctionLiteral struct {
	Parameters []FunctionParameter // Function parameters with optional types
	Body       []Statement         // Function body statements
	ReturnType TypeInfo            // Optional return type hint, embeds directly
	Variadic   bool                // True if function accepts variable arguments
}

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.Type != TypeUnknown {
		result += ": " + typeToString(fl.ReturnType)
	}
	result += "\n"

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

// CallExpression represents function calls with arguments.
// Stores inferred return type from function signature analysis.
type CallExpression struct {
	Function  Expression   // Function expression to call
	Arguments []Expression // Argument expressions
	typeInfo  TypeInfo     // Inferred return type, embeds directly
}

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) TypeInfo() TypeInfo { return ce.typeInfo }

// PrefixExpression represents unary operations like negation and logical not.
// Stores result type based on operator and operand type analysis.
type PrefixExpression struct {
	Operator string     // Operator symbol ("-", "not")
	Right    Expression // Operand expression
	typeInfo TypeInfo   // Result type, embeds directly
}

func (pe *PrefixExpression) expressionNode() {}
func (pe *PrefixExpression) String() string {
	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) TypeInfo() TypeInfo { return pe.typeInfo }

// InfixExpression represents binary operations between two expressions.
// Stores result type based on operator and operand type compatibility.
type InfixExpression struct {
	Left     Expression // Left operand
	Right    Expression // Right operand
	Operator string     // Operator symbol ("+", "-", "==", "and", etc.)
	typeInfo TypeInfo   // Result type, embeds directly
}

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) TypeInfo() TypeInfo { return ie.typeInfo }

// IndexExpression represents bracket-based member access (table[key]).
// Stores inferred element type based on container type analysis.
type IndexExpression struct {
	Left     Expression // Container expression
	Index    Expression // Index/key expression
	typeInfo TypeInfo   // Element type, embeds directly
}

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

// DotExpression represents dot-based member access (table.key).
// Stores inferred member type based on container type and field analysis.
type DotExpression struct {
	Left     Expression // Container expression
	Key      string     // Member name
	typeInfo TypeInfo   // Member type, embeds directly
}

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

// TablePair represents key-value pairs in table literals and struct constructors.
// Key is nil for array-style elements, non-nil for object-style elements.
type TablePair struct {
	Key   Expression // Key expression (nil for array elements)
	Value Expression // 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/array/object literals with key-value pairs.
// Always has table type, provides methods to check if it's array-style.
type TableLiteral struct {
	Pairs []TablePair // Key-value pairs (key nil for array elements)
}

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) TypeInfo() TypeInfo { return TableType }

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

// Helper function to convert TypeInfo to string representation
func typeToString(t TypeInfo) string {
	switch t.Type {
	case TypeNumber:
		return "number"
	case TypeString:
		return "string"
	case TypeBool:
		return "bool"
	case TypeNil:
		return "nil"
	case TypeTable:
		return "table"
	case TypeFunction:
		return "function"
	case TypeAny:
		return "any"
	case TypeStruct:
		return fmt.Sprintf("struct<%d>", t.StructID)
	default:
		return "unknown"
	}
}

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