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


@dataclass
class RootNode:
	statements: list[Any]
@dataclass
class InstructionNode:
	instruction: str
	parent: FunctionNode | RootNode
	arguments: list[Any]
@dataclass
class StringNode:
	value: str
@dataclass
class NumberNode:
	value: float
@dataclass
class VarRefNode:
	var_name: str
@dataclass
class VarPointerNode:
	var_name: str
@dataclass
class FunctionCallNode:
	func_name: str
@dataclass
class TypeNode:
	value: str
@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

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:
		print(token)
		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":
					scope = scope.parent # go up one scope

				current_node.parent = scope
			else:
				current_node = FunctionNode([], [], scope)
				current_node_type = "func"
		
		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.INTEGER or token.type == TokenType.FLOAT:
				if current_node_type == "inst":
					current_node.arguments.append(NumberNode(token.value))
				else:
					traceback(code, "SyntaxError", "Expected instruction, not number.")

			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":
					print(current_node)
					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.EOF:
				root_node.statements.append(current_node)
		last_token = token

	return root_node