From 81408b8e1b608298f99fe4626e6ddd7a08ea3a9c Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 5 Jan 2024 22:54:08 -0500
Subject: [PATCH] add `last` to Pvec

---
 src/pvec/containers_pvec.ml  | 16 ++++++++++++++++
 src/pvec/containers_pvec.mli | 10 ++++++++++
 tests/pvec/t_pvec.ml         | 14 +++++++++++++-
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/pvec/containers_pvec.ml b/src/pvec/containers_pvec.ml
index 99dee905..ba3c504e 100644
--- a/src/pvec/containers_pvec.ml
+++ b/src/pvec/containers_pvec.ml
@@ -213,6 +213,22 @@ let pop_opt (self : 'a t) : ('a * 'a t) option =
   else
     Some (pop self)
 
+let[@inline] last self =
+  if self.size = 0 then invalid_arg "pvec.last";
+  A.get self.tail (A.length self.tail - 1)
+
+let last_opt self =
+  if self.size = 0 then
+    None
+  else
+    Some (A.get self.tail (A.length self.tail - 1))
+
+let drop_last self =
+  if self.size = 0 then
+    self
+  else
+    snd (pop self)
+
 let rec iteri_rec_ f idx (self : _ tree) =
   match self with
   | Empty -> ()
diff --git a/src/pvec/containers_pvec.mli b/src/pvec/containers_pvec.mli
index 2d3a815c..dd3cd71a 100644
--- a/src/pvec/containers_pvec.mli
+++ b/src/pvec/containers_pvec.mli
@@ -44,6 +44,12 @@ val get : 'a t -> int -> 'a
 
 val get_opt : 'a t -> int -> 'a option
 
+val last : 'a t -> 'a
+(** Last element.
+      @raise Invalid_argument if the vec is empty *)
+
+val last_opt : 'a t -> 'a option
+
 val pop : 'a t -> 'a * 'a t
 (** Pop last element.
    @raise Invalid_argument in case the vec is empty. *)
@@ -51,6 +57,10 @@ val pop : 'a t -> 'a * 'a t
 val pop_opt : 'a t -> ('a * 'a t) option
 (** Pop last element. *)
 
+val drop_last : 'a t -> 'a t
+(** Like {!pop_opt} but doesn't return the last element.
+  Returns the same vector if it's empty. *)
+
 val iter : ('a -> unit) -> 'a t -> unit
 
 val iteri : (int -> 'a -> unit) -> 'a t -> unit
diff --git a/tests/pvec/t_pvec.ml b/tests/pvec/t_pvec.ml
index 25c7391a..5f8ba83a 100644
--- a/tests/pvec/t_pvec.ml
+++ b/tests/pvec/t_pvec.ml
@@ -89,6 +89,12 @@ module Ref_impl = struct
     | x :: tl -> x, List.rev tl
     | [] -> invalid_arg "empty"
 
+  let last_opt l =
+    if l = [] then
+      None
+    else
+      Some (List.nth l (List.length l - 1))
+
   let is_empty l = l = []
 
   let choose l =
@@ -109,6 +115,7 @@ module Op = struct
     | Check_len
     | Check_to_list
     | Check_to_gen
+    | Check_last
 
   let well_formed ops : bool =
     let rec loop size = function
@@ -121,6 +128,7 @@ module Op = struct
       | Check_is_empty :: tl
       | Check_len :: tl
       | Check_to_list :: tl
+      | Check_last :: tl
       | Check_to_gen :: tl ->
         loop size tl
     in
@@ -137,6 +145,7 @@ module Op = struct
     | Check_len -> "check_len"
     | Check_to_list -> "check_to_list"
     | Check_to_gen -> "check_to_gen"
+    | Check_last -> "check_last"
 
   let shrink shrink_x (op : _ t) : _ Q.Iter.t =
     let open Q.Shrink in
@@ -146,7 +155,7 @@ module Op = struct
     | Pop -> empty
     | Add_list l -> list ~shrink:shrink_x l >|= fun x -> Add_list x
     | Check_get _ | Check_choose | Check_is_empty | Check_len | Check_to_list
-    | Check_to_gen ->
+    | Check_to_gen | Check_last ->
       empty
 
   let shrink_l shrink_x : _ t list Q.Shrink.t =
@@ -171,6 +180,7 @@ module Op = struct
                    1, return (Check_is_empty, size);
                    1, return (Check_to_list, size);
                    1, return (Check_to_gen, size);
+                   1, return (Check_last, size);
                  ];
                  (if size > 0 then
                    [
@@ -231,6 +241,8 @@ let check_ops ~show_x (ops : 'a Op.t list) : unit =
         if to_list !cur <> Ref_impl.to_list !cur_ref then fail ()
       | Op.Check_choose ->
         if Option.is_some (choose !cur) <> Ref_impl.choose !cur_ref then fail ()
+      | Op.Check_last ->
+        if last_opt !cur <> Ref_impl.last_opt !cur_ref then fail ()
       | Op.Check_to_gen ->
         if
           to_seq !cur |> CCSeq.to_list