open Camlp4.PreCast
open Syntax

let types = Hashtbl.create 17

let declare loc a t =
  Hashtbl.add types a t;
  <:str_item@loc<  >>

exception Name_Not_found

let get_name patt =
  match patt with
    | <:patt< $lid:name$ >> -> name
    | _ -> raise Name_Not_found (* we can't assure we call this function
                                   only for function name in let bindings *)

let rec add t expr =
  match (t, expr) with
    | <:ctyp< $t1$ -> $t2$ >>, <:expr@_loc< fun $p$ -> $e$ >> -> 
        <:expr< fun ($p$ : $t1$) -> $add t2 e$ >>
        | <:ctyp< $t$ >>, <:expr< $e$ >> ->
          let _loc = Ast.loc_of_expr e in
            (* Actually, <:expr< $e$ >> match everything, so we can't extract
               directly the location and we need to use loc_of_expr to get one *)
        <:expr< ($e$ : $t$) >>
 
let add_types bi =
  try
    match bi with
      | <:binding@_loc< $patt$ = $expr$ >> -> (
          let name = get_name patt in
            try
              let t = Hashtbl.find types name in
                <:binding< $patt$ = $add t expr$ >>
            with
                _ -> bi
        )
      | _ -> bi
  with
      _ -> bi

DELETE_RULE Gram str_item: "let"; opt_rec; binding END;
EXTEND Gram
str_item:
[
  [ "let"; r = opt_rec; bi = binding ->
      <:str_item< value $rec:r$ $add_types bi$ >>
  ]
| [
	"declare"; u = a_LIDENT; ":"; t = ctyp -> declare _loc u t
  ]
];
END;;