options {
  STATIC = false;
  LOOKAHEAD = 2;
  UNICODE_INPUT = true;
  NODE_PREFIX = "AST";
}

PARSER_BEGIN(ArrowParser)
package coffee.liz.lambda.parser;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import coffee.liz.lambda.ast.Expression.AbstractionExpression;
import coffee.liz.lambda.ast.Expression.ApplicationExpression;
import coffee.liz.lambda.ast.Expression.IdentifierExpression;
import coffee.liz.lambda.ast.LambdaProgram;
import coffee.liz.lambda.ast.Macro;
import coffee.liz.lambda.ast.Expression;
import coffee.liz.lambda.ast.SourceSpan;
import coffee.liz.lambda.ast.SourceComment;

public class ArrowParser {
  public static void main(final String[] args) throws Exception {
    final ArrowParser parser = new ArrowParser(System.in);
    final LambdaProgram program = parser.Program();
    System.out.println(program);
  }

  private static SourceSpan spanFrom(final Token start, final Token end) {
    return new SourceSpan(start.beginLine, start.beginColumn, end.endLine, end.endColumn);
  }

  private static SourceSpan spanFrom(final Token start, final Expression end) {
    return new SourceSpan(start.beginLine, start.beginColumn, end.span().endLine(), end.span().endColumn());
  }

  private static SourceSpan spanFromInts(final int startLine, final int startColumn, final Token end) {
    return new SourceSpan(startLine, startColumn, end.endLine, end.endColumn);
  }

  private static SourceSpan spanFromInts(final int startLine, final int startColumn, final Expression end) {
    return new SourceSpan(startLine, startColumn, end.span().endLine(), end.span().endColumn());
  }

  private static Expression withComment(final Expression expr, final Optional<SourceComment> comment) {
    if (comment.isEmpty()) {
      return expr;
    }
    if (expr instanceof AbstractionExpression) {
      final AbstractionExpression a = (AbstractionExpression) expr;
      return new AbstractionExpression(comment, a.span(), a.parameter(), a.body());
    }
    if (expr instanceof ApplicationExpression) {
      final ApplicationExpression a = (ApplicationExpression) expr;
      return new ApplicationExpression(comment, a.span(), a.applicable(), a.argument());
    }
    if (expr instanceof IdentifierExpression) {
      final IdentifierExpression a = (IdentifierExpression) expr;
      return new IdentifierExpression(comment, a.span(), a.name());
    }
    return expr;
  }

  private static Expression withSpan(final Expression expr, final SourceSpan span) {
    if (expr instanceof AbstractionExpression) {
      final AbstractionExpression a = (AbstractionExpression) expr;
      return new AbstractionExpression(a.comment(), span, a.parameter(), a.body());
    }
    if (expr instanceof ApplicationExpression) {
      final ApplicationExpression a = (ApplicationExpression) expr;
      return new ApplicationExpression(a.comment(), span, a.applicable(), a.argument());
    }
    if (expr instanceof IdentifierExpression) {
      final IdentifierExpression a = (IdentifierExpression) expr;
      return new IdentifierExpression(a.comment(), span, a.name());
    }
    return expr;
  }

  private static Expression withCommentAndSpan(final Expression expr, final Optional<SourceComment> comment, final SourceSpan span) {
    if (expr instanceof AbstractionExpression) {
      final AbstractionExpression a = (AbstractionExpression) expr;
      return new AbstractionExpression(comment, span, a.parameter(), a.body());
    }
    if (expr instanceof ApplicationExpression) {
      final ApplicationExpression a = (ApplicationExpression) expr;
      return new ApplicationExpression(comment, span, a.applicable(), a.argument());
    }
    if (expr instanceof IdentifierExpression) {
      final IdentifierExpression a = (IdentifierExpression) expr;
      return new IdentifierExpression(comment, span, a.name());
    }
    return expr;
  }

  private static Optional<SourceComment> extractComment(final Token t) {
    if (t == null || t.specialToken == null) {
      return Optional.empty();
    }
    Token st = t.specialToken;
    while (st.specialToken != null) {
      st = st.specialToken;
    }
    final Token firstComment = st;
    final StringBuilder sb = new StringBuilder();
    Token lastComment = st;
    int lastEndLine = 0;
    while (st != null) {
      if (sb.length() > 0) {
        final int blankLines = st.beginLine - lastEndLine - 1;
        for (int i = 0; i < blankLines; i++) {
          sb.append("\n");
        }
        sb.append("\n");
      }
      String image = st.image;
      while (image.endsWith("\n") || image.endsWith("\r")) {
        image = image.substring(0, image.length() - 1);
      }
      sb.append(image);
      lastEndLine = st.endLine;
      lastComment = st;
      st = st.next;
    }
    final SourceSpan span = new SourceSpan(
      firstComment.beginLine, firstComment.beginColumn,
      lastComment.endLine, lastComment.endColumn
    );
    return Optional.of(new SourceComment(sb.toString(), span));
  }
}
PARSER_END(ArrowParser)

SKIP : {
  " " | "\t" | "\r" | "\n"
}

SPECIAL_TOKEN : {
  < COMMENT: "--" (~["\n","\r"])* ("\n"|"\r"|"\r\n")? >
}

TOKEN : {
    < LET: "let" >
  | < ARROW: "->" >
  | < EQ: "=" >
  | < SEMI: ";" >
  | < LPAREN: "(" >
  | < RPAREN: ")" >
  | < IDENT: (["a"-"z","A"-"Z","0"-"9","_"])+ >
}

LambdaProgram Program() :
{
  final List<Macro> macros = new ArrayList<Macro>();
  Macro m;
  Expression body;
  Token eofToken;
  Optional<SourceComment> eofComment;
}
{
  (
    m = Macro()
    { macros.add(m); }
  )*
  body = Expression()
  eofToken = <EOF>
  { eofComment = extractComment(eofToken); }
  {
    final SourceSpan span;
    if (!macros.isEmpty()) {
      span = spanFromInts(macros.get(0).span().startLine(), macros.get(0).span().startColumn(), eofToken);
    } else {
      span = new SourceSpan(body.span().startLine(), body.span().startColumn(), eofToken.endLine, eofToken.endColumn);
    }
    if (eofComment.isPresent() && body.comment().isEmpty()) {
      body = withComment(body, eofComment);
    }
    return new LambdaProgram(span, macros, body);
  }
}

Macro Macro() :
{
  Token letToken;
  Token name;
  Token semiToken;
  Expression value;
  Optional<SourceComment> comment;
}
{
  letToken = <LET>
  { comment = extractComment(letToken); }
  name = <IDENT>
  <EQ>
  value = Expression()
  semiToken = <SEMI>
  {
    return new Macro(comment, spanFrom(letToken, semiToken), name.image, value);
  }
}

Expression Expression() :
{
  Token paramToken;
  Expression e;
  Expression body;
  Optional<SourceComment> comment;
}
{
  (
    LOOKAHEAD(<IDENT> <ARROW>)
    paramToken = <IDENT>
    { comment = extractComment(paramToken); }
    <ARROW>
    body = Expression()
    {
      e = new AbstractionExpression(comment, spanFrom(paramToken, body), paramToken.image, body);
    }
  |
    e = Application()
  )
  {
    return e;
  }
}

Expression Application() :
{
  Expression e;
  Expression arg;
  Token rparen;
}
{
  e = Atom()
  (
    <LPAREN>
    arg = Expression()
    rparen = <RPAREN>
    {
      e = new ApplicationExpression(Optional.empty(), spanFromInts(e.span().startLine(), e.span().startColumn(), rparen), e, arg);
    }
  )*
  {
    return e;
  }
}

Expression Atom() :
{
  Token id;
  Token lparen;
  Token rparen;
  Expression e;
  Optional<SourceComment> comment;
}
{
  id = <IDENT>
  { comment = extractComment(id); }
  {
    return new IdentifierExpression(comment, spanFrom(id, id), id.image);
  }
|
  lparen = <LPAREN>
  { comment = extractComment(lparen); }
  e = Expression()
  rparen = <RPAREN>
  {
    final SourceSpan parenSpan = spanFrom(lparen, rparen);
    if (comment.isPresent() && e.comment().isEmpty()) {
      return withCommentAndSpan(e, comment, parenSpan);
    }
    return withSpan(e, parenSpan);
  }
}