Menhir (LR) grammar for pi types with syntactic sugar `(x y z : a) -> b` and `a -> b`

parser_simple.mly

%{
(** the [Surface] module contains our AST. *)
open Surface

(** exception raised when a [nonempty_list] returns an empty list. *)
exception EmptyUnfolding

(** desugars `(x y z : a) -> b` into `(x : a) -> (y : a) -> (z : a) -> b`. *)
let rec unfoldPi (xs : string list) (dom : expr) (cod : expr) : expr =
  match xs with
  | [] -> raise EmptyUnfolding
  | [x] -> Pi (x, dom, cod)
  | x :: rest -> Pi (x, dom, unfoldPi rest dom cod)

(** folds a nonempty list of expressions to their correctly-associated application. *)
let rec unfoldApp (es : expr list) : expr =
  match es with
  | [] -> raise EmptyUnfolding
  | [e] -> e
  | e1 :: e2 :: rest -> unfoldApp ((App (e1, e2)) :: rest)

(** converts an application of the form `(((w x) y ) z)` of variables
    to a list of variable names. returns [None] if any of the expressions are not a variable. *)
let rec appToVars (es : expr) : (string list) option =
  match es with
  | Var x -> Some [x]
  | App (rest, Var x) ->
    begin match appToVars rest with
    | None -> None
    | Some xs -> Some (xs @ [x])
    end
  | _ -> None

(** if the LHS of an arrow is an application of variables with an annotation,
    desugar it using [unfoldPi]. in any other case treat it as a non-dependent arrow. *)
let postprocessPi (a : expr) (b : expr) : expr =
  match a with
  | Ann (es, t) ->
    begin match appToVars es with
    | Some xs -> unfoldPi xs t b
    | None -> Pi ("", a, b)
    end
  | _ -> Pi ("", a, b)

%}

%token EOF

%token <string> IDENT
%token LPAREN RPAREN COLON
%token ARROW
%token TYPE

(** [ARROW]s are right-associative, [COLON]s require disambiguation *)
%nonassoc COLON
%right ARROW
(** subtlety: [COLON] must be higher than [ARROW], so that
    `x : a -> b` == `x : (a -> b)`
                 /= `(x : a) -> b` *)


%start program

%type <expr> program
%type <expr> expr
%type <expr> atom

%%
(** the grammar itself is straightforward and very declarative *)
program:
  | e=expr; EOF { e }

expr:
  | es=nonempty_list(atom) { unfoldApp es }
  | e=expr; COLON; t=expr { Ann (e, t) }
  | a=expr; ARROW; b=expr { postprocessPi a b }

atom:
  | x=IDENT { Var x }
  | LPAREN; e=expr; RPAREN { e }
  | TYPE { Uni }