From e8d783db80e2ab8decd37fc52b06a47d150a20cd Mon Sep 17 00:00:00 2001
From: "Haochen M. Kotoi-Xie" <haochenx@acm.org>
Date: Wed, 9 Aug 2023 05:25:42 +0900
Subject: [PATCH] add String.partition

---
 kxclib.ml                       | 16 ++++++++++++
 unit_test/baselib_unit_tests.ml | 45 +++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/kxclib.ml b/kxclib.ml
index 3dffaf8..e53d626 100644
--- a/kxclib.ml
+++ b/kxclib.ml
@@ -1154,6 +1154,22 @@ end
 module String = struct
   include String
 
+  (** [partition_opt n s] returns [None] if [n] is greater than the length of [s] or
+      [Some (s1, s2)] where [s = s1^s2 && n = length s1] otherwise *)
+  let partition_opt n str =
+    let len = length str in
+    if n > len then None
+    else Some (sub str 0 n, sub str n (len - n))
+
+  (** [partition n s] returns [Some (s1, s2)] where [s = s1^s2 && n = length s1].
+
+      {!Invalid_argument} will be thrown if [n] is greater than the length of [s] *)
+  let partition n str =
+    partition_opt n str
+    |> Option.v' (fun () ->
+           invalid_arg'
+             "Kxclib.String.partition - n (= %d) is greater than the length (=%d) of %S"
+             n (length str) str)
 
   (** [empty str] returns true when str is of zero length *)
   let empty str = length str = 0
diff --git a/unit_test/baselib_unit_tests.ml b/unit_test/baselib_unit_tests.ml
index 19bc729..fc88a43 100644
--- a/unit_test/baselib_unit_tests.ml
+++ b/unit_test/baselib_unit_tests.ml
@@ -66,6 +66,50 @@ let seq_drop_0 = seq_drop (list int) 2 [2; 3; 4] [4]
 let seq_drop_1 = seq_drop (list string) 0 ["A"; "B"; "C"] ["A"; "B"; "C"]
 let seq_drop_2 = seq_drop (list string) 3 ["A"; "B"; "C"] []
 
+let string_partition =
+  let counter = ref 0 in
+  let case str n (p1, p2) =
+    let id = get_and_decr counter |> (~-) in
+    test_case (sprintf "string_partition_%d" id) `Quick (fun () ->
+        check (option (pair string string))
+          (sprintf "string_partition_opt.case %S %d (%S, %S)"
+             str n p1 p2)
+          (Some (p1, p2))
+          (String.partition_opt n str);
+        check (pair string string)
+          (sprintf "string_partition_exn.case %S %d (%S, %S)"
+             str n p1 p2)
+          (p1, p2)
+          (String.partition n str);
+      ) in
+  let case_fail str n = (* failure cases *)
+    let id = get_and_decr counter |> (~-) in
+    test_case (sprintf "string_partition_%d" id) `Quick (fun () ->
+        check (option (pair string string))
+          (sprintf "string_partition_opt.should_fail %S %d"
+             str n)
+          (None)
+          (String.partition_opt n str);
+        try
+          let p1, p2 = String.partition n str in
+          failf "string_partition_exn.should_fail %S %d but got (%S, %S)"
+            str n p1 p2
+        with Invalid_argument _ -> ()
+      )
+  in [
+    "string_partition", [
+      case "" 0 ("", "");
+      case_fail "" 1;
+      case "a" 0 ("", "a");
+      case "a" 1 ("a", "");
+      case_fail "a" 2;
+      case "hello" 0 ("", "hello");
+      case "hello" 3 ("hel", "lo");
+      case "hello" 5 ("hello", "");
+      case_fail "hello" 7;
+    ];
+  ]
+
 
 [%%if ocaml_version < (4, 14, 0)]
 let stream_take tstbl n org_lst expected_lst () =
@@ -732,6 +776,7 @@ let () =
       test_case "seq_drop_2" `Quick seq_drop_2
     ];
   ] @ stream_suite
+    @ string_partition
     @ json_escaped_suite
     @ json_unparse @ json_show
     @ jvpath_unparse