mapper.ml

let call fn loc = function
  | PStr [ {pstr_desc = Pstr_eval (
    { pexp_desc = Pexp_constant(Const_string(sym, _))}, _)} ] ->
      with_default_loc loc (fun () -> fn sym)
  | _ -> raise (Location.Error(Location.error ~loc (
    "sqlexpr extension accepts a string")))

let call_sqlcheck loc = function   
  | PStr [ {pstr_desc = Pstr_eval ({ pexp_desc =
    Pexp_constant(Const_string("sqlite", None))}, _)}] ->
      with_default_loc loc sqlcheck_sqlite
  | _ -> raise (Location.Error(Location.error ~loc (
    "sqlcheck extension accepts \"sqlite\"")))

let shared_exprs = Hashtbl.create 25

let register_shared_expr expr =
  let id = new_id () in
  Hashtbl.add shared_exprs id expr;
  id

let get_shared_expr = Hashtbl.find shared_exprs

let shared_expr (bindings : value_binding list) : value_binding =
  let f {pvb_expr; _} =
    match pvb_expr with
    | {pexp_desc = Pexp_ident ({txt;})} ->
        let id = Longident.last txt in
        Hashtbl.mem shared_exprs id
    | _ -> false in
  List.find f bindings
  
let vb_expr_ident = function
  | {pvb_expr = {pexp_desc = Pexp_ident ({txt})}} -> Longident.last txt
  | _ -> raise Not_found

let vb_expr {pvb_expr} = pvb_expr

let new_mapper argv = Ast_mapper.({
  default_mapper with
  expr = (fun mapper expr ->
    match expr with
    (* is this an extension node? *)
    | {pexp_desc = Pexp_extension ({txt = "sql"; loc}, pstr)} ->
        call gen_sql loc pstr
    | {pexp_desc = Pexp_extension ({txt = "sqlc"; loc}, pstr)} ->
        let expr = call (gen_sql ~cacheable:true) loc pstr in
        let id = register_shared_expr expr in
        let expr = Exp.ident {txt=Longident.Lident id; loc} in
        with_default_loc loc (fun () -> expr)
    | {pexp_desc = Pexp_extension ({txt = "sqlinit"; loc}, pstr)} ->
        call (gen_sql ~init:true) loc pstr
    | {pexp_desc = Pexp_extension ({txt = "sqlcheck"; loc}, pstr)} ->
        call_sqlcheck loc pstr
    (* Delegate to the default mapper *)
    | x -> default_mapper.expr mapper x);
  structure_item = fun mapper structure_item ->
    match structure_item with
    | {pstr_desc = Pstr_value (rec_flag, value_bindings)} -> begin try
        let list_replace e1 e2 = List.map(fun x -> if x = e1 then e2 else x) in
        let value_binding = shared_expr value_bindings in
        let id = vb_expr_ident value_binding in
        let expr = vb_expr value_binding in
        let e = [%expr let [%p AC.pvar id] = [%e get_shared_expr id] in [%e expr]] in
        let vb = {value_binding with pvb_expr = e} in
        let value_bindings = list_replace value_binding vb value_bindings in
        { structure_item with pstr_desc = Pstr_value (rec_flag, value_bindings)}
      with _ -> structure_item
    end
    | x -> default_mapper.structure_item mapper x
})

let () =
  Random.self_init ();
  Ast_mapper.register "sqlexpr" new_mapper;
  ()