GroundPY/ground_ast.py

from __future__ import annotations
from dataclasses import dataclass
from tokenizer import Token, TokenType
from typing import Optional, Any
from error import traceback
from console import console


@dataclass
class RootNode:
	statements: list[Any]
@dataclass
class InstructionNode:
	instruction: str
	parent: FunctionNode | RootNode
	arguments: list[Any]
@dataclass
class StringNode:
	value: str
	def __repr__(self):
		return "String"
	def __str__(self):
		return self.value
@dataclass
class IntNode:
	value: float
	def __repr__(self):
		return "Int"
	def __str__(self):
		return str(self.value)
@dataclass
class FloatNode:
	value: float
	def __repr__(self):
		return "Float"
@dataclass
class VarRefNode:
	var_name: str
	def __repr__(self):
		return "VariableReference"
@dataclass
class VarPointerNode:
	var_name: str
	def __repr__(self):
		return "VariablePointer"
@dataclass
class FunctionCallNode:
	func_name: str
	def __repr__(self):
		return "FunctionCall"
@dataclass
class TypeNode:
	value: str
	def __repr__(self):
		return "Type"
	def __str__(self):
		return "-" + self.value
@dataclass
class ArgNode:
	arg_type: str
	name: str | None
	parent: FunctionNode
@dataclass
class FunctionNode:
	args: list[ArgNode]
	statements: list[Any]
	parent: FunctionNode | RootNode
	return_type: Optional[str] = None
	name: Optional[str] = None
@dataclass
class LabelDecNode:
	name: str
	def __repr__(self):
		return "LabelDecleration"
@dataclass
class LabelRefNode:
	name: str
	def __repr__(self):
		return "LabelReference"
@dataclass
class LineRefNode:
	line: int
	def __repr__(self):
		return "LineReference"
@dataclass
class BoolNode:
	value: bool
	def __repr__(self):
		return "Boolean"
	def __str__(self):
		return str(self.value).lower()

def generate_ast(tokens: list[Token], code: str) -> RootNode:
	root_node = RootNode([])

	current_node = None
	last_token = None
	current_node_type = None
	scope = root_node
	
	# todo: this is the absolute WORST way i could do this, but i could not care less lmao, its not even performant...... 
	for token in tokens:
		if token.type == TokenType.INSTRUCTION:
			
			if current_node:
				scope.statements.append(current_node)

			if token.value != "fun":
				if current_node_type == "func":
					scope = current_node

				current_node = InstructionNode(token.value, scope, [])
				current_node_type = "inst"

				if current_node.instruction == "endfun":
					if scope == root_node:
						traceback(code, "ScopeError", "You can't use endfun from global scope, did you mean to use the \"end\" instruction?")

					scope = scope.parent # go up one scope
					current_node.parent = scope
					scope.statements.append(current_node)
					current_node = None
				else:
					current_node.parent = scope
			else:
				current_node = FunctionNode([], [], scope)
				current_node_type = "func"
		elif token.type == TokenType.LABEL_DECLERATION:
			if current_node:
				scope.statements.append(current_node)
				if current_node_type == "func":
					scope = current_node
				current_node_type = None
				current_node = None
			scope.statements.append(LabelDecNode(token.value))
		
		if current_node:
			if token.type == TokenType.STRING:
				if current_node_type == "inst":
					current_node.arguments.append(StringNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not string.")

			elif token.type == TokenType.BOOL:
				if current_node_type == "inst":
					current_node.arguments.append(BoolNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not string.")

			elif token.type == TokenType.INTEGER:
				if current_node_type == "inst":
					current_node.arguments.append(IntNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not integer.")

			elif token.type == TokenType.FLOAT:
				if current_node_type == "inst":
					current_node.arguments.append(FloatNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not float.")

			elif token.type == TokenType.LINE_REFERENCE:
				if current_node_type == "inst":
					current_node.arguments.append(LineRefNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not line reference")
			
			elif token.type == TokenType.LABEL_REFERENCE:
				if current_node_type == "inst":
					current_node.arguments.append(LabelRefNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not label reference")

			elif token.type == TokenType.VARIABLE_POINTER:
				if current_node_type == "inst":
					current_node.arguments.append(VarPointerNode(token.value))
				elif last_token and last_token.type == TokenType.TYPE and current_node_type == "func":
					current_node.args[-1].name = token.value
				else:
					traceback(code, "SyntaxError", "Expected instruction, not variable pointer.")

			elif token.type == TokenType.VARIABLE_REFERENCE:
				if current_node_type == "inst":
					current_node.arguments.append(VarRefNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not variable reference.")

			elif token.type == TokenType.TYPE:
				if current_node_type == "inst":
					current_node.arguments.append(TypeNode(token.value))
				elif current_node_type == "func":
					if last_token and last_token.type == TokenType.FUNCTION_REFERENCE or current_node.return_type:
						current_node.args.append(ArgNode(
							arg_type=token.value,
							name=None,
							parent=current_node
						))
					else:
						current_node.return_type = token.value
				else:
					traceback(code, "SyntaxError", "Expected instruction, not type.")

			elif token.type == TokenType.FUNCTION_REFERENCE:
				if last_token and last_token.type == TokenType.TYPE and current_node_type == "func":
					current_node.name = token.value
				elif current_node_type == "inst":
					current_node.arguments.append(FunctionCallNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction or function return type, got function reference.")
			
			elif token.type == TokenType.COMMENT:
				continue

			elif token.type == TokenType.EOF:
				root_node.statements.append(current_node)
		last_token = token

	return root_node