%{
open Ast
open Util
%}

%token BOOL INT REAL
%token PROBE
%token TRUE FALSE
%token IF THEN ELSE
%token NOT AND OR
%token MOD
%token CONST NODE RETURNS
%token VAR LET TEL
%token PRE
%token ASSUME GUARANTEE
%token AUTOMATON STATE INITIAL UNTIL UNLESS RESUME

%token LPAREN RPAREN
%token LCURLY RCURLY
%token STAR PLUS MINUS DIVIDE
%token LESS GREATER LESS_EQUAL GREATER_EQUAL
%token DIFF EQUAL
%token SEMICOLON COLON COMMA
%token ARROW

%token <string> IDENT
%token <string> INTVAL
%token <string> REALVAL

%token EOF

(* Lowest priorities first *)
%left ARROW
%left OR
%left AND
%left EQUAL DIFF
%left LESS GREATER LESS_EQUAL GREATER_EQUAL
%left PLUS MINUS
%left STAR DIVIDE MOD

(* Root token *)
%start<Ast.toplevel list> file

%%

file:
| t=list(toplevel) EOF      { t }

toplevel:
| d=ext(const_decl)         { AST_const_decl d }
| d=ext(node_decl)          { AST_node_decl d }


primary_expr:
| LPAREN e=expr RPAREN      { e }
| e=ext(IDENT)              { AST_identifier e }
| e=ext(INTVAL)             { AST_int_const e }
| e=ext(REALVAL)            { AST_real_const e }
| TRUE                      { AST_bool_const true }
| FALSE                     { AST_bool_const false }
| e=ext(IDENT) LPAREN l=separated_list(COMMA,ext(expr)) RPAREN 
                            { AST_instance (e, l, fst e ^ uid ()) }

unary_expr:
| e=primary_expr                   { e }
| NOT e=ext(unary_expr)            { AST_not(e) }
| PRE e=ext(unary_expr)            { AST_pre(e, "pre"^uid()) }
| o=unary_op e=ext(unary_expr)     { AST_unary (o, e) }

%inline unary_op:
| PLUS          { AST_UPLUS }
| MINUS         { AST_UMINUS }


binary_expr:
| e=unary_expr                                        { e }
| e=ext(binary_expr) ARROW f=ext(binary_expr)         { AST_arrow(e, f) }
| e=ext(binary_expr) o=binary_op f=ext(binary_expr)   { AST_binary (o, e, f) }
| e=ext(binary_expr) o=binary_rel f=ext(binary_expr)  { AST_binary_rel (o, e, f) }
| e=ext(binary_expr) o=binary_bool f=ext(binary_expr) { AST_binary_bool (o, e, f) }

%inline binary_op:
| STAR           { AST_MUL }
| DIVIDE         { AST_DIV }
| MOD            { AST_MOD }
| PLUS           { AST_PLUS }
| MINUS          { AST_MINUS }
%inline binary_rel:
| LESS           { AST_LT }
| GREATER        { AST_GT }
| LESS_EQUAL     { AST_LE }
| GREATER_EQUAL  { AST_GE }
| EQUAL          { AST_EQ}
| DIFF           { AST_NE }
%inline binary_bool:
| AND            { AST_AND }
| OR             { AST_OR }

if_expr:
| IF c=ext(expr) THEN t=ext(expr) ELSE e=ext(expr)    { AST_if(c, t, e) }
| e=binary_expr                                       { e }

expr:
| e=if_expr      { e }


(* Equations (can be asserts, automata, ...) *)
automaton:
| AUTOMATON a=IDENT? s=list(ext(state)) RETURNS r=separated_list(COMMA, IDENT)
    { ((match a with Some n -> n | None -> "aut") ^ uid (), s, r) }

state:
| i=boption(INITIAL) STATE n=IDENT
  unless=trans(UNLESS)
  v=option(var_decl)
  b=body
  until=trans(UNTIL)
{ if unless <> [] then failwith "UNLESS transitions not supported.";
  { initial = i;
    name = n;
    locals = (match v with Some v -> v | None -> []);
    body = b;
    until = until;
} }

trans(TT):
| TT t=nonempty_list(terminated(transition, SEMICOLON))   { t }
|                                                         { [] }

transition:
| IF e=ext(expr) RESUME s=ext(IDENT)                      { (e, s) }


eqn:
| i=separated_list(COMMA, ext(IDENT)) EQUAL e=ext(expr)         
                                                { AST_assign(i, e) }
| ASSUME i=ext(IDENT) COLON e=ext(expr)         { AST_assume(i, e) }
| GUARANTEE i=ext(IDENT) COLON e=ext(expr)      { AST_guarantee(i, e) }
| a=automaton                                   { AST_automaton(a) }

typ:
| INT                                           { AST_TINT }
| BOOL                                          { AST_TBOOL }
| REAL                                          { AST_TREAL }

(* Declarations *)
dbody:
| e=ext(eqn) SEMICOLON                          { [e] }
| l=body                                        { l }

body:
| LET l=nonempty_list(terminated(ext(eqn), SEMICOLON)) TEL
  { l }

var:
|   p=boption(PROBE) i=ext(IDENT)   { (p, i) }

vari:
|   vn=separated_list(COMMA, var) COLON t=typ
    { List.map (fun (p, i) -> (p, i, t)) vn }

vars:
| v=separated_list(SEMICOLON, vari)
    { List.flatten v }

const_decl:
| CONST i=IDENT COLON t=typ EQUAL e=ext(expr) SEMICOLON
    { {
        name = i;
        typ = t;
        value = e;
    } }

var_decl:
| VAR l=nonempty_list(terminated(vari, SEMICOLON))
    { List.flatten l }

node_decl:
| NODE id=IDENT LPAREN v=vars RPAREN
    RETURNS LPAREN rv=vars RPAREN
    e = dbody
{ { name = id;
    args = v;
    ret = rv;
    var = [];
    body = e;
} }
| NODE id=IDENT LPAREN v=vars RPAREN
    RETURNS LPAREN rv=vars RPAREN
    lv=var_decl
    b=body
{ { name = id;
    args = v;
    ret = rv;
    var = lv;
    body = b;
} }

(* Utility : add extent information *)
%inline ext(X): 
| x=X { x, ($startpos, $endpos) }

%%