From 2a3cab5ccc3922eef6d7e635b8e4d294733dbf24 Mon Sep 17 00:00:00 2001
From: octachron <octa@polychoron.fr>
Date: Wed, 29 May 2024 20:00:53 +0200
Subject: [PATCH] special_function: fix indexing operator parsing

---
 src/longident.ml  | 43 +++++++++++++++++++++++++++++--------------
 test/base/test.ml | 28 +++++++++++++++++++++++++---
 2 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/src/longident.ml b/src/longident.ml
index 803ea2bc..acdfffff 100644
--- a/src/longident.ml
+++ b/src/longident.ml
@@ -46,26 +46,41 @@ let parse_simple s =
   | [] -> assert false
   | s :: l -> unflatten ~init:(Lident s) l
 
+(* find the first matching pair of parentheses *)
+let rec parentheses lpos opened pos len s =
+  if pos >= len then if opened > 1 then Error () else Ok None
+  else
+    match s.[pos] with
+    | '(' ->
+        let lpos = if opened = 0 then pos else lpos in
+        parentheses lpos (opened + 1) (pos + 1) len s
+    | ')' ->
+        let opened = opened - 1 in
+        if opened = 0 then Ok (Some (lpos, pos))
+        else if opened < 0 then Error ()
+        else parentheses lpos opened (pos + 1) len s
+    | _ -> parentheses lpos opened (pos + 1) len s
+
 (* handle ["A.B.(+.+)"] or ["Vec.(.%.()<-)"] *)
 let parse s =
-  let invalid () =
-    invalid_arg (Printf.sprintf "Ppxlib.Longident.parse: %S" s)
+  let invalid variant =
+    invalid_arg (Printf.sprintf "Ppxlib.Longident.parse(%s): %S" variant s)
   in
-  if String.length s < 1 then invalid ();
-  let open_par = String.index_opt s '(' in
-  let close_par = String.index_opt s ')' in
-  match (s.[0], open_par, close_par) with
-  | ('A' .. 'Z' | 'a' .. 'z' | '_'), None, None -> parse_simple s
-  | _, None, None -> Lident s (* This is a raw operator, no module path *)
-  | _, None, _ | _, _, None -> invalid ()
-  | _, Some l, Some r -> (
-      if Int.(r <> String.length s - 1) then invalid ();
+  if String.length s < 1 then invalid "empty string";
+  let par = parentheses (-1) 0 0 (String.length s) s in
+  match (s.[0], par) with
+  | ('A' .. 'Z' | 'a' .. 'z' | '_'), Ok None -> parse_simple s
+  | _, Ok None -> Lident s (* This is a raw operator, no module path *)
+  | _, Error _ -> invalid "unbalanced parenthesis"
+  | _, Ok (Some (l, r)) -> (
+      if Int.(r <> String.length s - 1) then
+        invalid "right parenthesis misplaced";
       let group =
-        if Int.(r = l + 1) then "()"
-        else String.trim (String.sub s ~pos:(l + 1) ~len:(r - l - 1))
+        let inside = String.trim (String.sub s ~pos:(l + 1) ~len:(r - l - 1)) in
+        if String.(inside = "") then "()" else inside
       in
       if Int.(l = 0) then Lident group
-      else if Char.(s.[l - 1] <> '.') then invalid ()
+      else if Char.(s.[l - 1] <> '.') then invalid "application in path"
       else
         let before = String.sub s ~pos:0 ~len:(l - 1) in
         match String.split_on_char before ~sep:'.' with
diff --git a/test/base/test.ml b/test/base/test.ml
index 3f9fef22..7b9f3714 100644
--- a/test/base/test.ml
+++ b/test/base/test.ml
@@ -101,17 +101,20 @@ let _ = convert_longident "Base.( land )"
 
 let _ = convert_longident "A(B)"
 [%%expect{|
-Exception: Invalid_argument "Ppxlib.Longident.parse: \"A(B)\"".
+Exception:
+Invalid_argument "Ppxlib.Longident.parse(application in path): \"A(B)\"".
 |}]
 
 let _ = convert_longident "A.B(C)"
 [%%expect{|
-Exception: Invalid_argument "Ppxlib.Longident.parse: \"A.B(C)\"".
+Exception:
+Invalid_argument "Ppxlib.Longident.parse(application in path): \"A.B(C)\"".
 |}]
 
 let _ = convert_longident ")"
 [%%expect{|
-Exception: Invalid_argument "Ppxlib.Longident.parse: \")\"".
+Exception:
+Invalid_argument "Ppxlib.Longident.parse(unbalanced parenthesis): \")\"".
 |}]
 
 let _ = convert_longident "+."
@@ -136,6 +139,25 @@ let _ = convert_longident "Foo.( *. )"
 ("Foo.( *. )", Ppxlib.Longident.Ldot (Ppxlib.Longident.Lident "Foo", "*."))
 |}]
 
+(* Indexing operators  *)
+let _ = convert_longident "(.!())"
+[%%expect{|
+- : string * longident = ("( .!() )", Ppxlib.Longident.Lident ".!()")
+|}]
+
+let _ = convert_longident "(.%(;..)<-)"
+[%%expect{|
+- : string * longident =
+("( .%(;..)<- )", Ppxlib.Longident.Lident ".%(;..)<-")
+|}]
+
+let _ = convert_longident "Vec.(.%(;..)<-)"
+[%%expect{|
+- : string * longident =
+("Vec.( .%(;..)<- )",
+ Ppxlib.Longident.Ldot (Ppxlib.Longident.Lident "Vec", ".%(;..)<-"))
+|}]
+
 let _ = Ppxlib.Code_path.(file_path @@ top_level ~file_path:"dir/main.ml")
 [%%expect{|
 - : string = "dir/main.ml"