From 810a8baecefe6b80b1b9b5be34174a4e8d858d3e Mon Sep 17 00:00:00 2001
From: Guillaume Petiot <guillaume@tarides.com>
Date: Wed, 21 Aug 2019 22:37:19 +0700
Subject: [PATCH] Fix formatting of custom indexing operators (#975)

---
 src/Ast.ml                    | 20 ++++++++++++++++----
 src/Ast.mli                   |  4 ++--
 src/Fmt_ast.ml                | 12 ++++++------
 test/passing/index_op.ml      |  6 ++++++
 test/passing/js_source.ml.ref |  6 +++---
 test/passing/source.ml.ref    |  4 +---
 6 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/src/Ast.ml b/src/Ast.ml
index f732e4be23..e049f0af2d 100644
--- a/src/Ast.ml
+++ b/src/Ast.ml
@@ -96,10 +96,22 @@ let index_op_get i =
 let index_op_set i =
   match String.chop_suffix i ~suffix:"<-" with
   | None -> None
-  | Some i -> (
-    match (i.[0], parens_kind i) with
-    | '.', Some (s, o, c) -> Some (s, o, c)
-    | _ -> None )
+  | Some i -> index_op_get i
+
+let index_op_lid i ~f =
+  match List.rev (Longident.flatten i) with
+  | [] -> impossible "not produced by parser"
+  | last :: rev_firsts -> (
+    match f last with
+    | None -> None
+    | Some (s, o, c) -> (
+      match List.rev rev_firsts with
+      | [] -> Some (s, o, c)
+      | firsts -> Some (String.concat ("." :: firsts) ^ s, o, c) ) )
+
+let index_op_get_lid = index_op_lid ~f:index_op_get
+
+let index_op_set_lid = index_op_lid ~f:index_op_set
 
 let index_op_string = (".", '[', ']')
 
diff --git a/src/Ast.mli b/src/Ast.mli
index f32ed5b265..024fb8418d 100644
--- a/src/Ast.mli
+++ b/src/Ast.mli
@@ -26,9 +26,9 @@ val is_infix_id : string -> bool
 val is_infix : expression -> bool
 (** Holds of infix symbol expressions. *)
 
-val index_op_get : string -> (string * Char.t * Char.t) option
+val index_op_get_lid : Longident.t -> (string * Char.t * Char.t) option
 
-val index_op_set : string -> (string * Char.t * Char.t) option
+val index_op_set_lid : Longident.t -> (string * Char.t * Char.t) option
 
 val index_op_get_sugar :
      Longident.t Location.loc
diff --git a/src/Fmt_ast.ml b/src/Fmt_ast.ml
index 9dc98f73d1..fc2493a48f 100644
--- a/src/Fmt_ast.ml
+++ b/src/Fmt_ast.ml
@@ -1721,23 +1721,23 @@ and fmt_expression c ?(box = true) ?pro ?epi ?eol ?parens ?(indent_wrap = 0)
                  (has_cmts, fmt_before_cmts, fmt_after_cmts, (fmt_op, args))
              | None -> (false, noop, noop, (noop, args))))
   | Pexp_apply
-      ( { pexp_desc= Pexp_ident {txt= Lident id; loc}
+      ( { pexp_desc= Pexp_ident {txt= id; loc}
         ; pexp_loc
         ; pexp_attributes= _
         ; _ }
       , (Nolabel, s) :: (Nolabel, i) :: _ )
-    when Option.is_some (index_op_get id) ->
-      let index_op = Option.value_exn (index_op_get id) in
+    when Option.is_some (index_op_get_lid id) ->
+      let index_op = Option.value_exn (index_op_get_lid id) in
       Cmts.relocate c.cmts ~src:pexp_loc ~before:loc ~after:loc ;
       fmt_index_op c ctx ~parens {txt= index_op; loc} s [i]
   | Pexp_apply
-      ( { pexp_desc= Pexp_ident {txt= Lident id; loc}
+      ( { pexp_desc= Pexp_ident {txt= id; loc}
         ; pexp_loc
         ; pexp_attributes= _
         ; _ }
       , (Nolabel, s) :: (Nolabel, i) :: (Nolabel, e) :: _ )
-    when Option.is_some (index_op_set id) ->
-      let index_op = Option.value_exn (index_op_set id) in
+    when Option.is_some (index_op_set_lid id) ->
+      let index_op = Option.value_exn (index_op_set_lid id) in
       Cmts.relocate c.cmts ~src:pexp_loc ~before:loc ~after:loc ;
       fmt_index_op c ctx ~parens {txt= index_op; loc} s [i] ~set:e
   | Pexp_apply (e0, [(Nolabel, e1)]) when is_prefix e0 ->
diff --git a/test/passing/index_op.ml b/test/passing/index_op.ml
index 221bfbb8e9..4d20be95b2 100644
--- a/test/passing/index_op.ml
+++ b/test/passing/index_op.ml
@@ -75,3 +75,9 @@ let _ = Bigarray.Genarray.set s [|1; 2|] 1
 let _ = Bigarray.Genarray.set s [|1; 2; 3|] 1
 
 let _ = s.{1, 2, 3, 4} <- 1
+
+let () =
+  let m = Mat.zeros 5 5 in
+  m.Mat.${[[2]; [5]]} |> ignore ;
+  let open Mat in
+  m.${[[2]; [5]]} |> ignore
diff --git a/test/passing/js_source.ml.ref b/test/passing/js_source.ml.ref
index 273c564e67..04d4051f84 100644
--- a/test/passing/js_source.ml.ref
+++ b/test/passing/js_source.ml.ref
@@ -9734,9 +9734,9 @@ module Indexop = struct
 
   ;;
   let h = Hashtbl.create 17 in
-  Def.( .%[]<- ) h "one" 1;
-  Def.( .%()<- ) h "two" 2;
-  Def.( .%{}<- ) h "three" 3
+  h.Def.%["one"] <- 1;
+  h.Def.%("two") <- 2;
+  h.Def.%{"three"} <- 3
 
   let x, y, z = Def.(h.%["one"], h.%("two"), h.%{"three"})
 end
diff --git a/test/passing/source.ml.ref b/test/passing/source.ml.ref
index 77e725a57c..005ccc2f80 100644
--- a/test/passing/source.ml.ref
+++ b/test/passing/source.ml.ref
@@ -9252,9 +9252,7 @@ module Indexop = struct
 
   ;;
   let h = Hashtbl.create 17 in
-  Def.( .%[]<- ) h "one" 1 ;
-  Def.( .%()<- ) h "two" 2 ;
-  Def.( .%{}<- ) h "three" 3
+  h.Def.%["one"] <- 1 ; h.Def.%("two") <- 2 ; h.Def.%{"three"} <- 3
 
   let x, y, z = Def.(h.%["one"], h.%("two"), h.%{"three"})
 end