Add ways to configure whether locations are shown

Signed-off-by: Nathan Rebours <[email protected]>

bin/pp_ast.ml

@@ -56,6 +56,8 @@ let pp_ast ~config ast =
 let input = ref None
 let kind = ref None
 let show_attrs = ref false
+let show_locs = ref false
+let loc_mode = ref `Short
 
 let set_input fn =
   match !input with
@@ -93,6 +95,12 @@ let args =
     ( "--show-attrs",
       Arg.Set show_attrs,
       "Show attributes in the pretty printed output" );
+    ( "--show-locs",
+      Arg.Set show_locs,
+      "Show locations in the pretty printed output" );
+    ( "--full-locs",
+      Arg.Unit (fun () -> loc_mode := `Full),
+      "Display locations in long form. Has no effect without --show-locs." );
   ]
 
 let main () =
@@ -118,7 +126,10 @@ let main () =
       in
       let input_name = match fn with "-" -> "<stdin>" | _ -> fn in
       let ast = load_input ~kind ~input_name fn in
-      let config = Pp_ast.Config.make ~show_attrs:!show_attrs () in
+      let config =
+        Pp_ast.Config.make ~show_attrs:!show_attrs ~show_locs:!show_locs
+          ~loc_mode:!loc_mode ()
+      in
       pp_ast ~config ast;
       Format.printf "%!\n"
 

src/pp_ast.ml

@@ -1,16 +1,26 @@
 open Import
 
 module Config = struct
-  type t = { show_attrs : bool }
+  type loc_mode = [ `Short | `Full ]
+  type t = { show_attrs : bool; show_locs : bool; loc_mode : loc_mode }
 
   module Default = struct
     let show_attrs = false
+    let show_locs = false
+    let loc_mode = `Short
   end
 
-  let default = { show_attrs = Default.show_attrs }
-  let make ?(show_attrs = Default.show_attrs) () = { show_attrs }
+  let default =
+    let open Default in
+    { show_attrs; show_locs; loc_mode }
+
+  let make ?(show_attrs = Default.show_attrs) ?(show_locs = Default.show_locs)
+      ?(loc_mode = Default.loc_mode) () =
+    { show_attrs; show_locs; loc_mode }
 end
 
+let cnum (pos : Lexing.position) = pos.pos_cnum - pos.pos_bol
+
 type simple_val =
   | Unit
   | Int of int
@@ -91,10 +101,36 @@ class lift_simple_val =
       Array (Array.map ~f:lift_a array |> Array.to_list)
 
     method other _a = Special "__"
-    method! location _loc = Special "__loc"
     method! location_stack _ls = Special "__lstack"
-    method! position _p = Special "__pos"
-    method! loc lift_a a_loc = lift_a a_loc.txt
+
+    method! position pos =
+      match (config.Config.show_locs, config.Config.loc_mode) with
+      | true, `Full -> super#position pos
+      | _, _ -> Special "__pos"
+
+    method! loc lift_a a_loc =
+      match config.Config.show_locs with
+      | true -> super#loc lift_a a_loc
+      | false -> lift_a a_loc.txt
+
+    method! location loc =
+      match (config.Config.show_locs, config.Config.loc_mode) with
+      | false, _ -> Special "__loc"
+      | true, `Full -> super#location loc
+      | true, `Short ->
+          let begin_line = loc.loc_start.pos_lnum in
+          let begin_char = cnum loc.loc_start in
+          let end_line = loc.loc_end.pos_lnum in
+          let end_char = cnum loc.loc_end in
+          let repr =
+            if Int.equal begin_line end_line then
+              Format.sprintf "l%ic%i..%i" begin_line begin_char end_char
+            else
+              Format.sprintf "l%ic%i..l%ic%i" begin_line begin_char end_line
+                end_char
+          in
+          let with_ghost = if loc.loc_ghost then repr ^ "(g)" else repr in
+          Special with_ghost
 
     method! attributes attrs =
       match config.Config.show_attrs with
@@ -110,9 +146,9 @@ class lift_simple_val =
           'record ->
           simple_val =
       fun ~lift_desc ~lift_record ~desc ~attrs x ->
-        match (config.show_attrs, attrs) with
-        | true, [] | false, _ -> lift_desc desc
-        | true, _ -> lift_record x
+        match (config.show_locs, config.show_attrs, attrs) with
+        | false, false, _ | false, true, [] -> lift_desc desc
+        | _, true, _ | true, _, _ -> lift_record x
 
     method! core_type ct =
       self#lift_record_with_desc ~lift_desc:self#core_type_desc

src/pp_ast.mli

@@ -4,7 +4,12 @@ module Config : sig
   type t
   (** Type for AST pretty-printing config *)
 
-  val make : ?show_attrs:bool -> unit -> t
+  val make :
+    ?show_attrs:bool ->
+    ?show_locs:bool ->
+    ?loc_mode:[ `Short | `Full ] ->
+    unit ->
+    t
   (** Create a custom pretty-printing config.
 
       Default values are the ones that are used when no configuration is passed
@@ -15,7 +20,19 @@ module Config : sig
         When set to [true], records such as [expression] that have a [desc]
         field will only be printed if the list of attributes is non-empty,
         otherwise their [_desc] field will be printed directly instead, as it is
-        the case when [show_attrs] is [false]. *)
+        the case when [show_attrs] is [false].
+
+      @param show_loc
+        controls whether locations are shown or hidden. Defaults to [false].
+
+      @param loc_mode
+        controls how locations are shown if they are shown at
+        all.
+        - When set to [`Short], locations are displayed as ["l1c6..l2c2"] for
+          multiline locations and as ["l1c6..12"] for single line locations.
+          Ghost locations are suffixed with a ["(g)"].
+        - When set to [`Full], locations are displayed as any other record would
+          be. Defaults to [`Short]. *)
 end
 
 type 'node pp = ?config:Config.t -> Format.formatter -> 'node -> unit