2 files changed, 184 insertions, 140 deletions

diff --git a/lib/command.ml b/lib/command.ml
index a55fb9e..9ab36f6 100644
--- a/lib/command.ml
+++ b/lib/command.ml
@@ -70,7 +70,7 @@ let to_scope = function
 let is_simple_movement k = List.mem ~equal:Poly.equal simple_movements k
 
 let instant_operation =
-  [ Key 'C'; Key 'D'; Key 'J'; Key 'P'; Key 'X'; Key 'p'; Key 'x' ]
+  [ Key 'C'; Key 'D'; Key 'Y'; Key 'J'; Key 'P'; Key 'X'; Key 'p'; Key 'x' ]
 
 let chord_operation = [ Key 'c'; Key 'd'; Key 'y' ]
 
diff --git a/lib/editor.ml b/lib/editor.ml
index 3368966..043c00e 100644
--- a/lib/editor.ml
+++ b/lib/editor.ml
@@ -95,12 +95,12 @@ module Action = struct
   include Applicative.Make (struct
     type 'a t = 'a action
 
-    let return a s = (a, s)
+    let return a e = (a, e)
 
     let apply f a e =
-      let f, e' = f e in
-      let x, e'' = a e' in
-      (f x, e'')
+      let g, e = f e in
+      let x, e = a e in
+      (g x, e)
 
     let map = `Define_using_apply
   end)
@@ -108,7 +108,7 @@ module Action = struct
   include Monad.Make (struct
     type 'a t = 'a action
 
-    let return x s = (x, s)
+    let return x e = (x, e)
 
     let bind a ~f x =
       let y, a' = a x in
@@ -117,6 +117,8 @@ module Action = struct
     let map = `Define_using_bind
   end)
 
+  let ( let* ) e f = bind e ~f
+  let ( and* ) = both
   let get e = (e, e)
   let put e _ = ((), e)
   let modify ~f e = ((), f e)
@@ -136,30 +138,30 @@ module Action = struct
   let get_focused_buffer e = (e.buffer, e)
   let set_focused_buffer b e = ((), { e with buffer = b })
 
+  let on_focused_buffer f =
+    let* b = get_focused_buffer in
+    let out, b = f b in
+    let* () = set_focused_buffer b in
+    let* () = update_cursor in
+    return out
+
+  let simulate f =
+    let* b = get_focused_buffer in
+    return (f b |> fst)
+
   let get_mode e = (e.mode, e)
   let set_mode m e = ((), { e with mode = m })
   let get_terminal_size e = (e.term.size, e)
 
-  let get_register r e =
+  let get_register ?(r = '"') e =
     assert (Char.('!' <= r && r <= '~'));
     (e.registers.(Char.to_int r), e)
 
-  let set_register r z e =
+  let set_register ?(r = '"') s e =
     assert (Char.('!' <= r && r <= '~'));
-    e.registers.(Char.to_int '"') <- z;
-    e.registers.(Char.to_int r) <- z;
+    e.registers.(Char.to_int r) <- s;
     ((), e)
 
-  let on_focused_buffer f =
-    let f e = { e with buffer = Option.map ~f e.buffer } in
-    modify ~f *> update_cursor
-
-  let on_focused_buffer_or_new f =
-    (get_focused_buffer
-    >>| Option.value ~default:Buffer.empty
-    >>| f >>= set_focused_buffer)
-    *> update_cursor
-
   let render =
     let aux e =
       let x, y = e.offset
@@ -228,30 +230,29 @@ let handle_insert_command =
   | Simple Arrow_right -> Buffer.Action.move_right |> on_focused_buffer
   | Simple Arrow_up -> Buffer.Action.move_up |> on_focused_buffer
   | Simple Backspace ->
-      Buffer.Action.delete_before ~cross_lines:true ~n:1 |> on_focused_buffer
+      Buffer.Action.delete_before ~cross_lines:true ~n:1
+      |> on_focused_buffer |> ignore_m
   | Simple (Ctrl 'Q') -> quit 0
   | Simple Delete ->
-      Buffer.Action.delete_after ~cross_lines:true ~n:1 |> on_focused_buffer
+      Buffer.Action.delete_after ~cross_lines:true ~n:1
+      |> on_focused_buffer |> ignore_m
   | Simple Enter -> Buffer.Action.newline |> on_focused_buffer
   | Simple Esc ->
-      (Buffer.Action.move_left |> on_focused_buffer) *> set_mode Normal
+      let* () = Buffer.Action.move_left |> on_focused_buffer in
+      set_mode Normal
   | Simple Page_down | Simple (Ctrl 'F') ->
-      fun e ->
-        (* TODO consider using the buffer window size (i.e., subtract status_size) *)
-        let (n, _), e = get_terminal_size e in
-        on_focused_buffer (Buffer.Action.move_down ~n) e
+      (* TODO consider using the buffer window size (i.e., subtract status_size) *)
+      let* n, _ = get_terminal_size in
+      Buffer.Action.move_down ~n |> on_focused_buffer
   | Simple Page_up | Simple (Ctrl 'B') ->
-      fun e ->
-        let (n, _), e = get_terminal_size e in
-        on_focused_buffer (Buffer.Action.move_up ~n) e
+      let* n, _ = get_terminal_size in
+      Buffer.Action.move_up ~n |> on_focused_buffer
   | Simple (Ctrl 'D') ->
-      fun e ->
-        let (r, _), e = get_terminal_size e in
-        on_focused_buffer (Buffer.Action.move_down ~n:(r / 2)) e
+      let* r, _ = get_terminal_size in
+      Buffer.Action.move_down ~n:(r / 2) |> on_focused_buffer
   | Simple (Ctrl 'U') ->
-      fun e ->
-        let (r, _), e = get_terminal_size e in
-        on_focused_buffer (Buffer.Action.move_up ~n:(r / 2)) e
+      let* r, _ = get_terminal_size in
+      Buffer.Action.move_up ~n:(r / 2) |> on_focused_buffer
   | Simple Home -> Buffer.Action.bol |> on_focused_buffer
   | Simple End -> Buffer.Action.eol |> on_focused_buffer
   | Simple Tab -> Buffer.Action.insert '\t' |> on_focused_buffer
@@ -277,128 +278,172 @@ let handle_normal_command c =
     | Chord (_, n, Noop, _, Up) -> Buffer.Action.move_up ?n |> on_focused_buffer
     | Chord (_, n, Noop, _, To_bol) ->
         let n = Option.value ~default:1 n - 1 in
-        Buffer.Action.(move_down ~n &> bol) |> on_focused_buffer
+        Buffer.Action.bol ~n |> on_focused_buffer
     | Chord (_, n, Noop, _, To_eol) ->
         let n = Option.value ~default:1 n - 1 in
-        Buffer.Action.(move_down ~n &> eol) |> on_focused_buffer
+        Buffer.Action.eol ~n |> on_focused_buffer
     | Simple Page_down | Simple (Ctrl 'F') ->
-        fun e ->
-          let (n, _), e = get_terminal_size e in
-          on_focused_buffer (Buffer.Action.move_down ~n) e
+        let* n, _ = get_terminal_size in
+        Buffer.Action.move_down ~n |> on_focused_buffer
     | Simple Page_up | Simple (Ctrl 'B') ->
-        fun e ->
-          let (n, _), e = get_terminal_size e in
-          on_focused_buffer (Buffer.Action.move_up ~n) e
+        let* n, _ = get_terminal_size in
+        Buffer.Action.move_up ~n |> on_focused_buffer
     | Simple (Ctrl 'D') ->
-        fun e ->
-          let (r, _), e = get_terminal_size e in
-          on_focused_buffer (Buffer.Action.move_down ~n:(r / 2)) e
+        let* r, _ = get_terminal_size in
+        Buffer.Action.move_down ~n:(r / 2) |> on_focused_buffer
     | Simple (Ctrl 'U') ->
-        fun e ->
-          let (r, _), e = get_terminal_size e in
-          on_focused_buffer (Buffer.Action.move_up ~n:(r / 2)) e
-    (* Change *)
-    | Shortcut (_, n, Change) ->
+        let* r, _ = get_terminal_size in
+        Buffer.Action.move_up ~n:(r / 2) |> on_focused_buffer
+    (* Yank *)
+    | Shortcut (r, n, Yank) ->
         let n = Option.value ~default:1 n - 1 in
-        (Buffer.Action.(
-           delete_to_eol &> move_down &> delete_lines ~n &> move_up &> eol)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, Line) ->
-        let n = Option.((value ~default:1 n1 * value ~default:1 n2) - 1) in
-        (Buffer.Action.(delete_lines ~n &> bol &> delete_to_eol)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, Down) ->
+        let* out = Buffer.Action.delete_to_eol ~n |> simulate in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Yank, n2, Line) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        (Buffer.Action.(delete_lines ~n &> bol &> delete_to_eol)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, Left) ->
+        let* out = Buffer.Action.delete_lines ~n |> simulate in
+        set_register ?r (Linewise out)
+    | Chord (r, n1, Yank, n2, Down) ->
+        let n = Option.((value ~default:1 n1 * value ~default:1 n2) + 1) in
+        let* out = Buffer.Action.delete_lines ~n |> simulate in
+        set_register ?r (Linewise out)
+    | Chord (r, n1, Yank, n2, Left) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        (Buffer.Action.delete_before ~n |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, Right) ->
+        let* out = Buffer.Action.delete_before ~n |> simulate in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Yank, n2, Right) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        (Buffer.Action.delete_after ~n |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, Up) ->
+        let* out = Buffer.Action.delete_after ~n |> simulate in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Yank, n2, Up) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        (Buffer.Action.(bol &> delete_to_eol &> delete_lines_before ~n)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, To_bol) ->
+        let* out =
+          Buffer.Action.(move_up ~n *> delete_lines ~n:(n + 1)) |> simulate
+        in
+        set_register ?r (Linewise out)
+    | Chord (r, _, Yank, _, To_bol) ->
+        let* out = Buffer.Action.delete_to_bol |> simulate in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Yank, n2, To_eol) ->
         let n = Option.((value ~default:1 n1 * value ~default:1 n2) - 1) in
-        (Buffer.Action.(
-           delete_to_bol &> move_down &> delete_lines ~n &> move_up &> bol)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
-    | Chord (_, n1, Change, n2, To_eol) ->
+        let* out = Buffer.Action.delete_to_eol ~n |> simulate in
+        set_register ?r (Glyphwise out)
+    (* Change *)
+    | Shortcut (r, n, Change) ->
+        let n = Option.value ~default:1 n - 1 in
+        let* out = Buffer.Action.delete_to_eol ~n |> on_focused_buffer in
+        let* () = set_register ?r (Glyphwise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, Line) ->
+        let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
+        let act =
+          let open Buffer.Action in
+          let* out = delete_lines ~n
+          and* () = insert_line ~before:true *> move_up in
+          return out
+        in
+        let* out = act |> on_focused_buffer in
+        let* () = set_register ?r (Linewise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, Down) ->
+        let n = Option.(value ~default:1 n1 * value ~default:1 n2) + 1 in
+        let act =
+          let open Buffer.Action in
+          let* out = delete_lines ~n
+          and* () = insert_line ~before:true *> move_up in
+          return out
+        in
+        let* out = act |> on_focused_buffer in
+        let* () = set_register ?r (Linewise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, Left) ->
+        let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
+        let* out = Buffer.Action.delete_before ~n |> on_focused_buffer in
+        let* () = set_register ?r (Glyphwise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, Right) ->
+        let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
+        let* out = Buffer.Action.delete_after ~n |> on_focused_buffer in
+        let* () = set_register ?r (Glyphwise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, Up) ->
+        let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
+        let act =
+          let open Buffer.Action in
+          let* out = move_up ~n *> delete_lines ~n:(n + 1)
+          and* () = insert_line ~before:true *> move_up in
+          return out
+        in
+        let* out = act |> on_focused_buffer in
+        let* () = set_register ?r (Linewise out) in
+        set_mode Insert
+    | Chord (r, _, Change, _, To_bol) ->
+        let* out = Buffer.Action.delete_to_bol |> on_focused_buffer in
+        let* () = set_register ?r (Glyphwise out) in
+        set_mode Insert
+    | Chord (r, n1, Change, n2, To_eol) ->
         let n = Option.((value ~default:1 n1 * value ~default:1 n2) - 1) in
-        (Buffer.Action.(
-           delete_to_eol &> move_down &> delete_lines ~n &> move_up &> eol)
-        |> on_focused_buffer_or_new)
-        *> set_mode Insert
+        let* out = Buffer.Action.delete_to_eol ~n |> on_focused_buffer in
+        let* () = set_register ?r (Glyphwise out) in
+        set_mode Insert
     (* Delete *)
-    | Shortcut (_, n, Delete) ->
+    | Shortcut (r, n, Delete) ->
         let n = Option.value ~default:1 n - 1 in
-        Buffer.Action.(
-          delete_to_eol &> move_down &> delete_lines ~n &> move_up &> eol)
-        |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, Line) ->
+        let* out = Buffer.Action.delete_to_eol ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Delete, n2, Line) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        Buffer.Action.delete_lines ~n |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, Down) ->
+        let* out = Buffer.Action.delete_lines ~n |> on_focused_buffer in
+        set_register ?r (Linewise out)
+    | Chord (r, n1, Delete, n2, Down) ->
         let n = Option.((value ~default:1 n1 * value ~default:1 n2) + 1) in
-        Buffer.Action.delete_lines ~n |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, Left) ->
+        let* out = Buffer.Action.delete_lines ~n |> on_focused_buffer in
+        set_register ?r (Linewise out)
+    | Chord (r, n1, Delete, n2, Left) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        Buffer.Action.delete_before ~n |> on_focused_buffer_or_new
-    | Shortcut (_, n, Erase_before) ->
-        let n = Option.(value ~default:1 n) in
-        Buffer.Action.delete_before ~n |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, Right) ->
+        let* out = Buffer.Action.delete_before ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Delete, n2, Right) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        Buffer.Action.delete_after ~n |> on_focused_buffer_or_new
-    | Shortcut (_, n, Erase_after) ->
-        let n = Option.(value ~default:1 n) in
-        Buffer.Action.delete_after ~n |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, Up) ->
+        let* out = Buffer.Action.delete_after ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Delete, n2, Up) ->
         let n = Option.(value ~default:1 n1 * value ~default:1 n2) in
-        Buffer.Action.(delete_lines ~n:1 &> delete_lines_before ~n)
-        |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, To_bol) ->
-        let n = Option.((value ~default:1 n1 * value ~default:1 n2) - 1) in
-        Buffer.Action.(
-          delete_to_bol &> move_down &> delete_lines ~n &> move_up &> bol)
-        |> on_focused_buffer_or_new
-    | Chord (_, n1, Delete, n2, To_eol) ->
+        let* out =
+          Buffer.Action.(move_up ~n *> delete_lines ~n:(n + 1))
+          |> on_focused_buffer
+        in
+        set_register ?r (Linewise out)
+    | Shortcut (r, n, Erase_before) ->
+        let n = Option.value ~default:1 n in
+        let* out = Buffer.Action.delete_before ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Shortcut (r, n, Erase_after) ->
+        let n = Option.value ~default:1 n in
+        let* out = Buffer.Action.delete_after ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Chord (r, _, Delete, _, To_bol) ->
+        let* out = Buffer.Action.delete_to_bol |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
+    | Chord (r, n1, Delete, n2, To_eol) ->
         let n = Option.((value ~default:1 n1 * value ~default:1 n2) - 1) in
-        Buffer.Action.(
-          delete_to_eol &> move_down &> delete_lines ~n &> move_up &> eol)
-        |> on_focused_buffer_or_new
+        let* out = Buffer.Action.delete_to_eol ~n |> on_focused_buffer in
+        set_register ?r (Glyphwise out)
     (* Paste *)
-    | Shortcut (r, n, Paste_after) ->
-        let r = Option.value ~default:'"' r in
-        let paste = function
-          | Empty -> noop
-          | Glyphwise z -> Buffer.Action.paste ?n z |> on_focused_buffer_or_new
-          | Linewise z ->
-              Buffer.Action.paste ~linewise:true ?n z
-              |> on_focused_buffer_or_new
-        in
-        get_register r >>= paste
-    | Shortcut (r, n, Paste_before) ->
-        let r = Option.value ~default:'"' r in
-        let paste = function
-          | Empty -> noop
-          | Glyphwise z ->
-              Buffer.Action.paste ~before:true ?n z |> on_focused_buffer_or_new
-          | Linewise z ->
-              Buffer.Action.paste ~before:true ~linewise:true ?n z
-              |> on_focused_buffer_or_new
-        in
-        get_register r >>= paste
+    | Shortcut (r, n, Paste_after) -> (
+        get_register ?r >>= function
+        | Empty -> noop
+        | Glyphwise z -> Buffer.Action.paste ?n z |> on_focused_buffer
+        | Linewise z ->
+            Buffer.Action.paste ~linewise:true ?n z |> on_focused_buffer)
+    | Shortcut (r, n, Paste_before) -> (
+        get_register ?r >>= function
+        | Empty -> noop
+        | Glyphwise z ->
+            Buffer.Action.paste ~before:true ?n z |> on_focused_buffer
+        | Linewise z ->
+            Buffer.Action.paste ~before:true ~linewise:true ?n z
+            |> on_focused_buffer)
     (* Join *)
     | Shortcut (_, n, Join) ->
         let n = Option.(value ~default:2 n) in
@@ -407,13 +452,12 @@ let handle_normal_command c =
     | Simple (Ctrl 'Q') -> quit 0
     (* Misc *)
     | Simple (Key 'A') ->
-        (Buffer.Action.eol |> on_focused_buffer_or_new) *> set_mode Insert
+        (Buffer.Action.eol |> on_focused_buffer) *> set_mode Insert
     | Simple (Key 'a') ->
-        (Buffer.Action.move_right |> on_focused_buffer_or_new)
-        *> set_mode Insert
+        (Buffer.Action.move_right |> on_focused_buffer) *> set_mode Insert
     (* | Key 'G' -> Buffer.Action.eof |> on_focused_buffer_or_new *)
     (* | Key 'I' -> noop *)
-    | Simple (Key 'i') -> (Fn.id |> on_focused_buffer_or_new) *> set_mode Insert
+    | Simple (Key 'i') -> set_mode Insert
     (* | Key 's' -> *)
     (*     (Buffer.Action.delete_after |> on_focused_buffer_or_new) *)
     (* *> set_mode Insert *)