Parse comment for docs?

[SeokminHong]

Hi, I'm trying to parse a simple pseudocode to generate codes. For the following pseudocode:

// A test function
function x(a, b) {
  // ...
  hello();
}

I want to generate the following parse tree:

- file
  - function
    - doc: "// A test function"
    - ident: "x"
    - args
      - ident: "a"
      - ident: "b"
    - COMMENT: "// ..."
    - expr
      - ident: "hello"
      - args: "()"
  - EOI: ""

Currently, I made a simple .pest file like the one below:

WHITESPACE = _{" " | "\t" | NEWLINE}
COMMENT = {"//" ~ (!NEWLINE ~ ANY)*}

ident_char = {ASCII_ALPHA | "_" | "$"}
ident = @{ident_char ~ (ASCII_DIGIT | ident_char)*}

args = {"(" ~ (ident ~ ("," ~ ident)* ~ ","?)? ~ ")"}


expr = {ident ~ args ~ ";" }

function = {
    "function"
    ~ ident
    ~ args
    ~ "{" ~ expr* ~ "}"
}

file = {
    SOI ~ (expr | function) ~ EOI
}

When I tried to add doc for function, it couldn't recognize the token because of the COMMENT. How can I parse doc comments?