AST is accepting functions!!!

plasma_parser.py

@@ -4,7 +4,7 @@ from typing import Callable
 from enum import Enum, auto
 
 from AST import Statement, Expression, Program
-from AST import ExpressionStatement, AssignmentStatement
+from AST import ExpressionStatement, AssignmentStatement, FunctionStatement, ReturnStatement, BlockStatement
 from AST import InfixExpression
 from AST import IntegerLiteral, FloatLiteral, IdentifierLiteral
 
@@ -39,6 +39,7 @@ class Parser:
 
 		self.prefix_parse_functions: dict[Token, Callable] = { # -1
 
+			TokenType.IDENT: self.__parse_identifier,
 			TokenType.INT: self.__parse_int_literal,
 			TokenType.FLOAT: self.__parse_float_literal,
 			TokenType.LPAREN: self.__parse_grouped_expression
@@ -114,9 +115,11 @@ class Parser:
 		match self.current_token.type:
 			case TokenType.IDENT:
 				return self.__parse_assignment_statement()
+			case TokenType.RETURN:
+				return self.__parse_return_statement()
 			case _:
 				return self.__parse_expression_statement()
-	
+
 	def __parse_expression_statement(self) -> ExpressionStatement:
 		expr = self.__parse_expression(PrecedenceType.P_LOWEST)
 
@@ -131,25 +134,88 @@ class Parser:
 		# x: Int = 10;
 		stmt: AssignmentStatement = AssignmentStatement(name=IdentifierLiteral(self.current_token.literal))
 
-		if not self.__expect_peek(TokenType.COLON):
-			return None
-		
-		if not self.__expect_peek(TokenType.TYPE):
-			return None
-		
-		stmt.value_type = self.current_token.literal
-		
-		if not self.__expect_peek(TokenType.EQ):
-			return None
-		
-		self.__next_token()
-
-		stmt.value = self.__parse_expression(PrecedenceType.P_LOWEST)
-
-		while not self.__current_token_is(TokenType.SEMICOLON) and not self.__current_token_is(TokenType.EOF):
+		if self.__peek_token_is(TokenType.EQ): # function definition
+			# x = Func(): Int { return 10; }
 			self.__next_token()
 
+			func_stmt: FunctionStatement = FunctionStatement(name=stmt.name)
+
+			if not self.__expect_peek(TokenType.TYPE): # Func word
+				return None
+			
+			if self.current_token.literal != "Func":
+				self.errors.append(f"Expected next token to be \"Func\", got {self.current_token.literal} instead.")
+				return None
+			
+			if not self.__expect_peek(TokenType.LPAREN):
+				return None
+			
+			func_stmt.parameters = []
+
+			if not self.__expect_peek(TokenType.RPAREN):
+				return None
+			
+			if not self.__expect_peek(TokenType.COLON):
+				return None
+			
+			if not self.__expect_peek(TokenType.TYPE):
+				return None
+			
+			func_stmt.return_type = self.current_token.literal
+
+			if not self.__expect_peek(TokenType.LBRACE):
+				return None
+			
+			func_stmt.body = self.__parse_block_statement()
+
+			return func_stmt
+
+		else:
+			if not self.__expect_peek(TokenType.COLON):
+				return None
+			
+			if not self.__expect_peek(TokenType.TYPE):
+				return None
+			
+			stmt.value_type = self.current_token.literal
+			
+			if not self.__expect_peek(TokenType.EQ):
+				return None
+			
+			self.__next_token()
+
+			stmt.value = self.__parse_expression(PrecedenceType.P_LOWEST)
+
+			while not self.__current_token_is(TokenType.SEMICOLON) and not self.__current_token_is(TokenType.EOF):
+				self.__next_token()
+
+			return stmt
+	
+	def __parse_return_statement(self) -> ReturnStatement:
+		stmt: ReturnStatement = ReturnStatement()
+
+		self.__next_token()
+
+		stmt.return_value = self.__parse_expression(PrecedenceType.P_LOWEST)
+
+		if not self.__expect_peek(TokenType.SEMICOLON):
+			return None
+		
 		return stmt
+
+	def __parse_block_statement(self) -> BlockStatement:
+		block_stmt: BlockStatement = BlockStatement()
+
+		self.__next_token()
+
+		while not self.__current_token_is(TokenType.RBRACE) and not self.__current_token_is(TokenType.EOF):
+			stmt: Statement = self.__parse_statement()
+			if stmt is not None:
+				block_stmt.statements.append(stmt)
+
+			self.__next_token()
+
+		return block_stmt
 	# endregion
 
 	# region Expression Methods
@@ -194,6 +260,9 @@ class Parser:
 	# endregion
 
 	# region Prefix Methods
+	def __parse_identifier(self) -> IdentifierLiteral:
+		return IdentifierLiteral(value=self.current_token.literal)
+
 	def __parse_int_literal(self) -> Expression:
 		int_lit: IntegerLiteral = IntegerLiteral()