Source file ag_mapping.ml

open Printf

open Ag_error

type loc = Atd_ast.loc

let annot_error loc =
  Ag_error.error loc "Invalid annotation"

type loc_id = string

(*
  Generic mapping, based on the core ATD types
*)
type ('a, 'b) mapping =
    [ `Unit of (loc * 'a * 'b)
    | `Bool of (loc * 'a * 'b)
    | `Int of (loc * 'a * 'b)
    | `Float of (loc * 'a * 'b)
    | `String of (loc * 'a * 'b)
    | `Sum of (loc * ('a, 'b) variant_mapping array * 'a * 'b)
    | `Record of (loc * ('a, 'b) field_mapping array * 'a * 'b)
    | `Tuple of (loc * ('a, 'b) cell_mapping array * 'a * 'b)
    | `List of (loc * ('a, 'b) mapping * 'a * 'b)
    | `Option of (loc * ('a, 'b) mapping * 'a * 'b)
    | `Nullable of (loc * ('a, 'b) mapping * 'a * 'b)
    | `Wrap of (loc * ('a, 'b) mapping * 'a * 'b)
    | `Name of (loc * string * ('a, 'b) mapping list * 'a option * 'b option)
    | `External of (loc * string * ('a, 'b) mapping list * 'a * 'b)
    | `Tvar of (loc * string) ]

and ('a, 'b) cell_mapping = {
  cel_loc : loc;
  cel_value : ('a, 'b) mapping;
  cel_arepr : 'a;
  cel_brepr : 'b
}

and ('a, 'b) field_mapping = {
  f_loc : loc;
  f_name : string;
  f_kind : Atd_ast.field_kind;
  f_value : ('a, 'b) mapping;
  f_arepr : 'a;
  f_brepr : 'b
}

and ('a, 'b) variant_mapping = {
  var_loc : loc;
  var_cons : string;
  var_arg : ('a, 'b) mapping option;
  var_arepr : 'a;
  var_brepr : 'b
}

type ('a, 'b) def = {
  def_loc : loc;
  def_name : string;
  def_param : string list;
  def_value : ('a, 'b) mapping option;
  def_arepr : 'a;
  def_brepr : 'b;
}


let as_abstract = function
    `Name (_, (loc, "abstract", l), a) ->
      if l <> [] then
        error loc "\"abstract\" takes no type parameters";
      Some (loc, a)
  | _ ->
      None

let is_abstract x = as_abstract x <> None


let loc_of_mapping x =
  match (x : (_, _) mapping) with
      `Unit (loc, _, _)
    | `Bool (loc, _, _)
    | `Int (loc, _, _)
    | `Float (loc, _, _)
    | `String (loc, _, _)
    | `Sum (loc, _, _, _)
    | `Record (loc, _, _, _)
    | `Tuple (loc, _, _, _)
    | `List (loc, _, _, _)
    | `Option (loc, _, _, _)
    | `Nullable (loc, _, _, _)
    | `Wrap (loc, _, _, _)
    | `Name (loc, _, _, _, _)
    | `External (loc, _, _, _, _)
    | `Tvar (loc, _) -> loc


module Env = Map.Make (String)

let rec subst env (x : (_, _) mapping) =
  match x with
      `Unit (loc, _, _)
    | `Bool (loc, _, _)
    | `Int (loc, _, _)
    | `Float (loc, _, _)
    | `String (loc, _, _) -> x
    | `Sum (loc, ar, a, b) ->
        `Sum (loc, Array.map (subst_variant env) ar, a, b)
    | `Record (loc, ar, a, b) ->
        `Record (loc, Array.map (subst_field env) ar, a, b)
    | `Tuple (loc, ar, a, b) ->
        `Tuple (loc, Array.map (subst_cell env) ar, a, b)
    | `List (loc, x, a, b) ->
        `List (loc, subst env x, a, b)
    | `Option (loc, x, a, b) ->
        `Option (loc, subst env x, a, b)
    | `Nullable (loc, x, a, b) ->
        `Nullable (loc, subst env x, a, b)
    | `Wrap (loc, x, a, b) ->
        `Wrap (loc, subst env x, a, b)
    | `Name (loc, name, args, a, b) ->
        `Name (loc, name, List.map (subst env) args, a, b)
    | `External (loc, name, args, a, b) ->
        `External (loc, name, List.map (subst env) args, a, b)
    | `Tvar (loc, s) ->
        try Env.find s env
        with Not_found ->
          invalid_arg (sprintf "Ag_mapping.subst_var: '%s" s)

and subst_variant env x =
  match x.var_arg with
      None -> x
    | Some v -> { x with var_arg = Some (subst env v) }

and subst_field env x =
  { x with f_value = subst env x.f_value }

and subst_cell env x =
  { x with cel_value = subst env x.cel_value }

(*
  Substitute type variables param in x by args
*)
let apply param x args =
  if List.length param <> List.length args then
    invalid_arg "Ag_mapping.apply";
  let env =
    List.fold_left2
      (fun env var value -> Env.add var value env)
      Env.empty param args
  in
  subst env x


let rec find_name loc env visited name =
  if List.mem name visited then
    error loc "Cyclic type definition"
  else
    let param, x = Env.find name env in
    (param, deref_expr env (name :: visited) x)

and deref_expr env visited x =
  match x with
      `Name (loc, name, args, _, _) ->
        (try
           let param, x = find_name loc env visited name in
           apply param x args
         with Not_found -> x)
    | _ -> x

let flatten l = List.flatten (List.map snd l)

let make_deref
    (l : (bool * ('a, 'b) def list) list) :
    (('a, 'b) mapping -> ('a, 'b) mapping) =

  let defs =
    List.fold_left
      (fun env d ->
         match d.def_value with
             None -> env
           | Some v -> Env.add d.def_name (d.def_param, v) env)
      Env.empty (flatten l) in

  fun x -> deref_expr defs [] x

(*
   Resolve names and unwrap `wrap` constructs
   (discarding annotations along the way)
*)
let rec unwrap (deref: ('a, 'b) mapping -> ('a, 'b) mapping) x =
  match deref x with
  | `Wrap (loc, x, a, b) -> unwrap deref x
  | x -> x

(* This is for debugging *)
let constructor : ('a, 'b) mapping -> string = function
  | `Unit _ -> "Unit"
  | `Bool _ -> "Bool"
  | `Int _ -> "Int"
  | `Float _ -> "Float"
  | `String _ -> "String"
  | `Sum _ -> "Sum"
  | `Record _ -> "Record"
  | `Tuple _ -> "Tuple"
  | `List _ -> "List"
  | `Option _ -> "Option"
  | `Nullable _ -> "Nullable"
  | `Wrap _ -> "Wrap"
  | `Name (loc, name, _, _, _) -> "Name " ^ name
  | `External _ -> "External"
  | `Tvar _ -> "Tvar"