%{
open Syntax

let here () = Location.span_of_lexing (Parsing.symbol_start_pos ()) (Parsing.symbol_end_pos ())

let symbol_span index =
  Location.span_of_lexing (Parsing.rhs_start_pos index) (Parsing.rhs_end_pos index)
%}

%token <int> INT
%token <bool> BOOL
%token <string> STRING
%token <string> IDENT
%token LET IN FUN IF THEN ELSE
%token TYPE INPUT COLLECTION QUERY FILTER MAP SUM COUNT
%token ARROW PIPE AND OR
%token EQ NE LT LE GT GE
%token PLUS MINUS STAR SLASH
%token LPAREN RPAREN COMMA DOT LBRACE RBRACE COLON SEMI
%token EOF

%start expr
%type <Syntax.expr> expr
%start declarations
%type <Syntax.declaration list> declarations

%%

expr:
  | pipeline { $1 }
  | LET IDENT EQ expr IN expr { make (ELet ($2, $4, $6)) (here ()) }
  | FUN IDENT ARROW expr { make (ELambda ($2, $4)) (here ()) }
  | IF expr THEN expr ELSE expr { make (EIf ($2, $4, $6)) (here ()) }

pipeline:
  | disjunction { $1 }
  | pipeline PIPE disjunction { make (EApp ($3, $1)) (here ()) }

disjunction:
  | conjunction { $1 }
  | disjunction OR conjunction { make (EBinop (Or, $1, $3)) (here ()) }

conjunction:
  | comparison { $1 }
  | conjunction AND comparison { make (EBinop (And, $1, $3)) (here ()) }

comparison:
  | sum { $1 }
  | sum EQ sum { make (EBinop (Eq, $1, $3)) (here ()) }
  | sum NE sum { make (EBinop (Ne, $1, $3)) (here ()) }
  | sum LT sum { make (EBinop (Lt, $1, $3)) (here ()) }
  | sum LE sum { make (EBinop (Le, $1, $3)) (here ()) }
  | sum GT sum { make (EBinop (Gt, $1, $3)) (here ()) }
  | sum GE sum { make (EBinop (Ge, $1, $3)) (here ()) }

sum:
  | product { $1 }
  | sum PLUS product { make (EBinop (Add, $1, $3)) (here ()) }
  | sum MINUS product { make (EBinop (Sub, $1, $3)) (here ()) }

product:
  | unary { $1 }
  | product STAR unary { make (EBinop (Mul, $1, $3)) (here ()) }
  | product SLASH unary { make (EBinop (Div, $1, $3)) (here ()) }

unary:
  | application { $1 }
  | MINUS unary { make (ENeg $2) (here ()) }

application:
  | atom { $1 }
  | application atom { make (EApp ($1, $2)) (here ()) }

atom:
  | INT { make (EInt $1) (here ()) }
  | BOOL { make (EBool $1) (here ()) }
  | STRING { make (EString $1) (here ()) }
  | IDENT { make (EVar $1) (here ()) }
  | LPAREN RPAREN { make EUnit (here ()) }
  | LPAREN expr RPAREN { $2 }
  | LPAREN expr COMMA expr_list RPAREN { make (ETuple ($2 :: $4)) (here ()) }
  | LBRACE record_fields RBRACE { make (ERecord $2) (here ()) }
  | atom DOT IDENT { make (EField ($1, $3)) (here ()) }
  | FILTER atom { make (EFilter $2) (here ()) }
  | MAP atom { make (EMap $2) (here ()) }
  | SUM { make ESum (here ()) }
  | COUNT { make ECount (here ()) }

expr_list:
  | expr { [ $1 ] }
  | expr COMMA expr_list { $1 :: $3 }

record_fields:
  | IDENT EQ expr { [ ($1, $3) ] }
  | IDENT EQ expr SEMI record_fields { ($1, $3) :: $5 }
  | IDENT EQ expr SEMI { [ ($1, $3) ] }

declarations:
  | decls EOF { List.rev $1 }

decls:
  | { [] }
  | decls declaration { $2 :: $1 }

declaration:
  | TYPE IDENT EQ LBRACE record_field_decls RBRACE {
      DRecord {
        rd_name = $2;
        rd_name_span = symbol_span 2;
        rd_fields = $5;
        rd_span = here ();
      }
    }
  | INPUT IDENT COLON type_expr {
      DInput { in_name = $2; in_name_span = symbol_span 2; in_ty = $4; in_span = here () }
    }
  | LET IDENT helper_params EQ expr {
      DHelper {
        h_name = $2;
        h_name_span = symbol_span 2;
        h_params = $3;
        h_body = $5;
        h_span = here ();
      }
    }
  | QUERY IDENT EQ expr {
      DQuery { q_name = $2; q_name_span = symbol_span 2; q_body = $4; q_span = here () }
    }

helper_params:
  | { [] }
  | helper_param_list { $1 }

helper_param_list:
  | IDENT { [ $1 ] }
  | IDENT helper_param_list { $1 :: $2 }

record_field_decls:
  | IDENT COLON type_expr { [ ($1, $3) ] }
  | IDENT COLON type_expr SEMI record_field_decls { ($1, $3) :: $5 }
  | IDENT COLON type_expr SEMI { [ ($1, $3) ] }

type_expr:
  | IDENT { make_ty (TyName $1) (here ()) }
  | COLLECTION type_expr { make_ty (TyCollection $2) (here ()) }
  | LPAREN type_expr RPAREN { $2 }
  | LPAREN type_expr COMMA type_expr_list RPAREN {
      make_ty (TyTuple ($2 :: $4)) (here ())
    }

type_expr_list:
  | type_expr { [ $1 ] }
  | type_expr COMMA type_expr_list { $1 :: $3 }