From cd6efc57bb1da51adae67165a64a2f0678c7d034 Mon Sep 17 00:00:00 2001
From: Ryan Chandler <ryangjchandler@gmail.com>
Date: Mon, 12 Sep 2022 00:24:04 +0100
Subject: [PATCH] parser: fix comments at end of class not being skipped

Closes #29.
---
 trunk_parser/src/parser/mod.rs | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/trunk_parser/src/parser/mod.rs b/trunk_parser/src/parser/mod.rs
index aac2b80..c278de3 100644
--- a/trunk_parser/src/parser/mod.rs
+++ b/trunk_parser/src/parser/mod.rs
@@ -971,11 +971,16 @@ impl Parser {
         self.lbrace()?;
 
         let mut body = Vec::new();
-        self.gather_comments();
         while self.current.kind != TokenKind::RightBrace && !self.is_eof() {
+            self.gather_comments();
+
+            if self.current.kind == TokenKind::RightBrace {
+                self.clear_comments();
+                break;
+            }
+
             body.push(self.class_statement()?);
         }
-
         self.rbrace()?;
 
         Ok(Statement::Class {
@@ -988,8 +993,6 @@ impl Parser {
     }
 
     fn class_statement(&mut self) -> ParseResult<Statement> {
-        self.gather_comments();
-
         match self.current.kind {
             TokenKind::Use => {
                 self.next();
@@ -2086,7 +2089,7 @@ impl Display for ParseError {
 mod tests {
     use super::Parser;
     use crate::{
-        ast::{Arg, ArrayItem, ElseIf, InfixOp, MethodFlag},
+        ast::{Arg, ArrayItem, ElseIf, InfixOp, MethodFlag, PropertyFlag},
         Expression, Identifier, Param, Statement, Type,
     };
     use trunk_lexer::Lexer;
@@ -2981,6 +2984,19 @@ mod tests {
         assert_ast("<?php ;", &[Statement::Noop]);
     }
 
+    #[test]
+    fn comment_at_end_of_class() {
+        assert_ast("<?php
+        class MyClass {
+            protected $a;
+            // my comment
+        }", &[
+            Statement::Class { name: "MyClass".into(), extends: None, implements: vec![], body: vec![
+                Statement::Property { var: "a".into(), value: None, r#type: None, flags: vec![PropertyFlag::Protected] }
+            ], flag: None }
+        ]);
+    }
+
     fn assert_ast(source: &str, expected: &[Statement]) {
         let mut lexer = Lexer::new(None);
         let tokens = lexer.tokenize(source).unwrap();