Support parsing simple return statements

src/parser.rs

@@ -5,6 +5,8 @@ enum ASTNode {
  IntegerLiteralExpression(Token),
  // TODO split declaration and assignment into two separate nodes.
  Declaration(Token, Token, Box<ASTNode>), // Type, Identifier, Expression
+ ReturnStatement(Token, Box<ASTNode>), // ReturnKeyword, Expression
+ Program(Vec<ASTNode>),
 }
 
 struct Parser {
@@ -18,7 +20,23 @@ impl Parser {
  }
 
  pub fn parse(&mut self) -> ASTNode {
- self.parse_declaration()
+ let mut result = Vec::new();
+ loop {
+ let c = self.current();
+ if c.token_type == TokenType::Eof {
+ break;
+ }
+ result.push(self.parse_statement());
+ }
+ ASTNode::Program(result)
+ }
+
+ fn parse_statement(&mut self) -> ASTNode {
+ match self.current().token_type {
+ TokenType::Type => self.parse_declaration(),
+ TokenType::ReturnKeyword => self.parse_return_statement(),
+ _ => panic!("Parser: Unexpected token {:?}", self.current()),
+ }
  }
 
  fn peak(&self, offset: usize) -> &Token {
@@ -67,6 +85,22 @@ impl Parser {
  ASTNode::Declaration(type_token, identifier, Box::new(literal))
  }
 
+ fn parse_return_statement(&mut self) -> ASTNode {
+ static EXPECTED_SEQUENCE: [TokenType; 3] = [
+ TokenType::ReturnKeyword,
+ TokenType::IntegerLiteral,
+ TokenType::SemiColon,
+ ];
+
+ self.expect_token_sequence(&EXPECTED_SEQUENCE);
+
+ let return_keyword = self.consume();
+ let literal = self.parse_int_literal();
+ self.advance(); // SemiColon
+
+ ASTNode::ReturnStatement(return_keyword, Box::new(literal))
+ }
+
  fn expect_token_sequence(&self, types: &[TokenType]) {
  for (i, token_type) in types.iter().enumerate() {
  if self.peak(i).token_type != *token_type {
@@ -91,16 +125,33 @@ mod tests {
  use crate::lexer::Lexer;
 
  #[rstest::rstest]
- #[case("int x = 55;", ASTNode::Declaration(
- Token{value: "int".to_string(), token_type: TokenType::Type, pos: 0},
- Token{value: "x".to_string(), token_type: TokenType::Identifier, pos: 4},
- Box::new(ASTNode::IntegerLiteralExpression(
- Token{value: "55".to_string(), token_type: TokenType::IntegerLiteral, pos: 8}
- ))
- ))]
+ #[case("int x = 55;", ASTNode::Program(
+ vec![ASTNode::Declaration(
+ Token{value: "int".to_string(), token_type: TokenType::Type, pos: 0},
+ Token{value: "x".to_string(), token_type: TokenType::Identifier, pos: 4},
+ Box::new(ASTNode::IntegerLiteralExpression(
+ Token{value: "55".to_string(), token_type: TokenType::IntegerLiteral, pos: 8}
+ ))
+ )])
+ )]
  fn test_parse_basic_declaration(#[case] test_case: String, #[case] expected: ASTNode) {
  let tokens = Lexer::new(test_case.clone()).lex();
  let result = Parser::new(tokens).parse();
  assert_eq!(expected, result);
  }
+
+ #[rstest::rstest]
+ #[case("return 123;", ASTNode::Program(
+ vec![ASTNode::ReturnStatement(
+ Token{value: "return".to_string(), token_type: TokenType::ReturnKeyword, pos: 0},
+ Box::new(ASTNode::IntegerLiteralExpression(
+ Token{value: "123".to_string(), token_type: TokenType::IntegerLiteral, pos: 7}
+ ))
+ )])
+ )]
+ fn test_parse_return_statement(#[case] test_case: String, #[case] expected: ASTNode) {
+ let tokens = Lexer::new(test_case).lex();
+ let result = Parser::new(tokens).parse();
+ assert_eq!(expected, result);
+ }
 }