open Lexing


type loc = Lexing.position * Lexing.position

let dummy_loc = (Lexing.dummy_pos, Lexing.dummy_pos)

exception Atd_error of string

type full_module = module_head * module_body

and module_head = loc * annot

and module_body = module_item list

and annot = annot_section list

and annot_section = string * (loc * annot_field list)

and annot_field = string * (loc * string option)

and type_def = loc * (string * type_param * annot) * type_expr

and module_item =
    [ `Type of type_def ]

and type_param = string list

and type_expr =
    [ `Sum of (loc * variant list * annot)
    | `Record of (loc * field list * annot)
    | `Tuple of (loc * cell list * annot)
    | `List of (loc * type_expr * annot)
    | `Option of (loc * type_expr * annot)
    | `Nullable of (loc * type_expr * annot)
    | `Shared of (loc * type_expr * annot)
    | `Wrap of (loc * type_expr * annot)
    | `Name of (loc * type_inst * annot)
    | `Tvar of (loc * string)
    ]
      (* `List, `Option, `Nullable, `Shared and `Wrap are
         the only predefined types with a type
         parameter (and no special syntax). *)

and type_inst = loc * string * type_expr list

and variant =
    [ `Variant of (loc * (string * annot) * type_expr option)
    | `Inherit of (loc * type_expr) ]

and cell = loc * type_expr * annot

and field_kind =
    [ `Required
    | `Optional
    | `With_default
    ]

and field =
    [ `Field of (loc * (string * field_kind * annot) * type_expr)
    | `Inherit of (loc * type_expr) ]


let loc_of_type_expr = function
    `Sum (loc, _, _)
  | `Record (loc, _, _)
  | `Tuple (loc, _, _)
  | `List (loc, _, _)
  | `Option (loc, _, _)
  | `Nullable (loc, _, _)
  | `Shared (loc, _, _)
  | `Wrap (loc, _, _)
  | `Name (loc, _, _)
  | `Tvar (loc, _) -> loc

let set_type_expr_loc loc = function
    `Sum (_, a, b) -> `Sum (loc, a, b)
  | `Record (_, a, b) -> `Record (loc, a, b)
  | `Tuple (_, a, b) -> `Tuple (loc, a, b)
  | `List (_, a, b) -> `List (loc, a, b)
  | `Option (_, a, b) -> `Option (loc, a, b)
  | `Nullable (_, a, b) -> `Nullable (loc, a, b)
  | `Shared (_, a, b) -> `Shared (loc, a, b)
  | `Wrap (_, a, b) -> `Wrap (loc, a, b)
  | `Name (_, a, b) -> `Name (loc, a, b)
  | `Tvar (_, a) -> `Tvar (loc, a)

let string_of_loc (pos1, pos2) =
  let line1 = pos1.pos_lnum
  and start1 = pos1.pos_bol in
  Printf.sprintf "File %S, line %i, characters %i-%i"
    pos1.pos_fname line1
    (pos1.pos_cnum - start1)
    (pos2.pos_cnum - start1)

let error s = raise (Atd_error s)

let error_at loc s = error (string_of_loc loc ^ "\n" ^ s)

let annot_of_type_expr = function
    `Sum (_, _, an)
  | `Record (_, _, an)
  | `Tuple (_, _, an)
  | `List (_, _, an)
  | `Option (_, _, an)
  | `Nullable (_, _, an)
  | `Shared (_, _, an)
  | `Wrap (_, _, an)
  | `Name (_, _, an) -> an
  | `Tvar (_, _) -> []



let map_annot f = function
    `Sum (loc, vl, a) ->  `Sum (loc, vl, f a)
  | `Record (loc, fl, a) -> `Record (loc, fl, f a)
  | `Tuple (loc, tl, a) -> `Tuple (loc, tl, f a)
  | `List (loc, t, a) -> `List (loc, t, f a)
  | `Option (loc, t, a) -> `Option (loc, t, f a)
  | `Nullable (loc, t, a) -> `Nullable (loc, t, f a)
  | `Shared (loc, t, a) -> `Shared (loc, t, f a)
  | `Wrap (loc, t, a) -> `Wrap (loc, t, f a)
  | `Tvar _ as x -> x
  | `Name (loc, (loc2, name, args), a) ->
      `Name (loc, (loc2, name, args), f a)

let rec amap_type_expr f (x : type_expr) =
  match x with
      `Sum (loc, vl, a) ->  `Sum (loc, List.map (amap_variant f) vl, f a)
    | `Record (loc, fl, a) -> `Record (loc, List.map (amap_field f) fl, f a)
    | `Tuple (loc, tl, a) -> `Tuple (loc, List.map (amap_cell f) tl, f a)
    | `List (loc, t, a) -> `List (loc, amap_type_expr f t, f a)
    | `Option (loc, t, a) -> `Option (loc, amap_type_expr f t, f a)
    | `Nullable (loc, t, a) -> `Nullable (loc, amap_type_expr f t, f a)
    | `Shared (loc, t, a) -> `Shared (loc, amap_type_expr f t, f a)
    | `Wrap (loc, t, a) -> `Wrap (loc, amap_type_expr f t, f a)
    | `Tvar _ as x -> x
    | `Name (loc, (loc2, name, args), a) ->
        `Name (loc, (loc2, name, List.map (amap_type_expr f) args), f a)

and amap_variant f = function
    `Variant (loc, (name, a), o) ->
      let o =
        match o with
            None -> None
          | Some x -> Some (amap_type_expr f x)
      in
      `Variant (loc, (name, f a), o)
  | `Inherit (loc, x) ->
      `Inherit (loc, amap_type_expr f x)

and amap_field f = function
    `Field (loc, (name, kind, a), x) ->
      `Field (loc, (name, kind, f a), amap_type_expr f x)
  | `Inherit (loc, x) ->
      `Inherit (loc, amap_type_expr f x)

and amap_cell f (loc, x, a) =
  (loc, amap_type_expr f x, f a)

let amap_module_item f (`Type (loc, (name, param, a), x)) =
  `Type (loc, (name, param, f a), amap_type_expr f x)

let amap_head f (loc, a) = (loc, f a)

let amap_body f l =
  List.map (amap_module_item f) l

let map_all_annot f ((head, body) : full_module) =
  (amap_head f head, amap_body f body)


let rec fold (f : type_expr -> 'a -> 'a) (x : type_expr) acc =
  let acc = f x acc in
  match x with
      `Sum (loc, variant_list, annot) ->
        List.fold_right (fold_variant f) variant_list acc

    | `Record (loc, field_list, annot) ->
        List.fold_right (fold_field f) field_list acc

    | `Tuple (loc, l, annot) ->
        List.fold_right (fun (loc, x, _) acc -> fold f x acc) l acc

    | `List (loc, type_expr, annot) ->
        fold f type_expr acc

    | `Option (loc, type_expr, annot) ->
        fold f type_expr acc

    | `Nullable (loc, type_expr, annot) ->
        fold f type_expr acc

    | `Shared (loc, type_expr, annot) ->
        fold f type_expr acc

    | `Wrap (loc, type_expr, annot) ->
        fold f type_expr acc

    | `Name (loc, (loc2, name, type_expr_list), annot) ->
        List.fold_right (fold f) type_expr_list acc

    | `Tvar (loc, string) ->
        acc

and fold_variant f x acc =
  match x with
      `Variant (loc, _, Some type_expr) -> fold f type_expr acc
    | `Variant _ -> acc
    | `Inherit (loc, type_expr) -> fold f type_expr acc

and fold_field f x acc =
  match x with
      `Field (loc, _, type_expr) -> fold f type_expr acc
    | `Inherit (loc, type_expr) -> fold f type_expr acc


module Type_names = Set.Make (String)

let union l =
  List.fold_left Type_names.union Type_names.empty l

let extract_type_names ?(ignorable = []) x =
  let ign s = List.mem s ignorable in

  let add s set =
    if ign s then set
    else Type_names.add s set
  in

  let acc =
    fold (
      fun x acc ->
        match x with
            `Name (loc, (loc2, name, l), a) -> add name acc
          | _ -> acc
    )
      x Type_names.empty
  in
  Type_names.elements acc

let is_parametrized x =
  fold (fun x b -> b || match x with `Tvar _ -> true | _ -> false) x false
