refactor: make zippers return 'pop'ped elements

1 files changed, 7 insertions, 2 deletions

diff --git a/lib/zipper.ml b/lib/zipper.ml
index 413f365..5f3b37a 100644
--- a/lib/zipper.ml
+++ b/lib/zipper.ml
@@ -49,7 +49,9 @@ let goto n z =
   let step = if n < 0 then left else right in
   Fn.apply_n_times ~n:(abs n) step z
 
-let pop ?(n = 1) z = { z with after = Sequence.drop_eagerly z.after n }
+let pop ?(n = 1) z =
+  let a, b = Sequence.split_n z.after n in
+  (Sequence.of_list a, { z with after = b })
 
 let pop_before ?(n = 1) =
   let rec aux m z =
@@ -58,7 +60,10 @@ let pop_before ?(n = 1) =
   aux 0
 
 let pop_after ?(n = 1) z =
-  if right_length z < 2 then z else right z |> pop ~n |> left
+  if right_length z < 2 then (Sequence.empty, z)
+  else
+    let a, b = right z |> pop ~n in
+    (a, left b)
 
 let push x z = { z with after = Sequence.shift_right z.after x }
 let push_seq s z = { z with after = Sequence.append s z.after }