Fix formatting of custom indexing operators (ocaml-ppx#975)

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 = (".", '[', ']')
 

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

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 ->

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

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

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