Source file ast_traverse.ml

open! Import

class map =
  object
    inherit Ppxlib_traverse_builtins.map
    inherit Ast.map
  end

class iter =
  object
    inherit Ppxlib_traverse_builtins.iter
    inherit Ast.iter
  end

class ['acc] fold =
  object
    inherit ['acc] Ppxlib_traverse_builtins.fold
    inherit ['acc] Ast.fold
  end

class ['acc] fold_map =
  object
    inherit ['acc] Ppxlib_traverse_builtins.fold_map
    inherit ['acc] Ast.fold_map
  end

class ['ctx] map_with_context =
  object
    inherit ['ctx] Ppxlib_traverse_builtins.map_with_context
    inherit ['ctx] Ast.map_with_context
  end

class virtual ['res] lift =
  object
    inherit ['res] Ppxlib_traverse_builtins.lift
    inherit ['res] Ast.lift
  end

let module_name = function None -> "_" | Some name -> name
let enter name path = if String.is_empty path then name else path ^ "." ^ name
let enter_opt name_opt path = enter (module_name name_opt) path

class map_with_path =
  object
    inherit [string] map_with_context as super

    (* WAS:
       method! structure_item_desc path x =
       match x with
       | Pstr_module mb -> super#structure_item_desc (enter mb.pmb_name.txt path) x
       | _ -> super#structure_item_desc path x

       Overriding [module_binding] seems to be OK because it does not catch
       local module bindings because at the moment the parsetree doesn't make
       use of [module_binding] for local modules, but that might change in the
       future, so this might be something to keep in mind.

       The following:

           module A = struct .. end
           module A = struct .. end

       is disallowed, but

           let _ = .. let module A = struct .. end in ..
           module A = struct .. end
           let _ = .. let module A = struct .. end in ..

       isn't, and the "path" constructed here would be able to differentiate
       between them. *)
    method! module_binding path mb =
      super#module_binding (enter_opt mb.pmb_name.txt path) mb

    method! module_declaration path md =
      super#module_declaration (enter_opt md.pmd_name.txt path) md

    method! module_type_declaration path mtd =
      super#module_type_declaration (enter mtd.pmtd_name.txt path) mtd
  end

let var_names_of =
  object
    inherit [string list] fold as super

    method! pattern p acc =
      let acc = super#pattern p acc in
      match p.ppat_desc with Ppat_var { txt; _ } -> txt :: acc | _ -> acc
  end

let ec_enter_module_opt ~loc name_opt ctxt =
  Expansion_context.Base.enter_module ~loc (module_name name_opt) ctxt

class map_with_expansion_context =
  object (self)
    inherit [Expansion_context.Base.t] map_with_context as super

    method! expression ctxt expr =
      super#expression (Expansion_context.Base.enter_expr ctxt) expr

    method! module_binding ctxt mb =
      super#module_binding
        (ec_enter_module_opt ~loc:mb.pmb_loc mb.pmb_name.txt ctxt)
        mb

    method! module_declaration ctxt md =
      super#module_declaration
        (ec_enter_module_opt ~loc:md.pmd_loc md.pmd_name.txt ctxt)
        md

    method! module_type_declaration ctxt mtd =
      super#module_type_declaration
        (Expansion_context.Base.enter_module ~loc:mtd.pmtd_loc mtd.pmtd_name.txt
           ctxt)
        mtd

    method! value_description ctxt vd =
      super#value_description
        (Expansion_context.Base.enter_value ~loc:vd.pval_loc vd.pval_name.txt
           ctxt)
        vd

    method! value_binding ctxt { pvb_pat; pvb_expr; pvb_attributes; pvb_loc } =
      let all_var_names = var_names_of#pattern pvb_pat [] in
      let var_name = Stdppx.List.last all_var_names in
      let in_binding_ctxt =
        match var_name with
        | None -> ctxt
        | Some var_name ->
            Expansion_context.Base.enter_value ~loc:pvb_loc var_name ctxt
      in
      let pvb_pat = self#pattern ctxt pvb_pat in
      let pvb_expr = self#expression in_binding_ctxt pvb_expr in
      let pvb_attributes = self#attributes in_binding_ctxt pvb_attributes in
      let pvb_loc = self#location ctxt pvb_loc in
      { pvb_pat; pvb_expr; pvb_attributes; pvb_loc }
  end

class sexp_of =
  object
    inherit [Sexp.t] Ast.lift
    method int = sexp_of_int
    method string = sexp_of_string
    method bool = sexp_of_bool
    method char = sexp_of_char
    method float = sexp_of_float
    method int32 = sexp_of_int32
    method int64 = sexp_of_int64
    method nativeint = sexp_of_nativeint
    method unit = sexp_of_unit
    method option = sexp_of_option
    method list = sexp_of_list
    method array : 'a. ('a -> Sexp.t) -> 'a array -> Sexp.t = sexp_of_array
    method other : 'a. 'a -> Sexp.t = fun _ -> Sexp.Atom "_"

    method record fields =
      List
        (List.map fields ~f:(fun (label, sexp) ->
             Sexp.List [ Atom label; sexp ]))

    method constr tag args =
      match args with [] -> Atom tag | _ -> List (Atom tag :: args)

    method tuple l = List l
  end

let sexp_of = new sexp_of