From 9271fc3f8f847326810c7e664c74cbf6ad6ee2a5 Mon Sep 17 00:00:00 2001
From: Junegunn Choi <junegunn.c@gmail.com>
Date: Thu, 5 Dec 2013 23:16:33 +0900
Subject: [PATCH] Make visual-operator repeatable (#24)

---
 README.md              |  5 +++-
 plugin/easy_align.vim  | 52 ++++++++++++++++++++++++++--------
 test/interactive.vader | 64 ++++++++++++++++++++++++++++++++++--------
 3 files changed, 97 insertions(+), 24 deletions(-)

diff --git a/README.md b/README.md
index c58c82f..8317fdd 100644
--- a/README.md
+++ b/README.md
@@ -80,6 +80,9 @@ try these commands:
     - Start EasyAlign command (`<Leader>a`) for `i`nner `p`aragraph
     - Align around `=`
 
+Notice that the commands are repeatable with `.` key if you have installed
+[repeat.vim](https://github.com/tpope/vim-repeat).
+
 Usage
 -----
 
@@ -217,7 +220,7 @@ nmap <leader>a <Plug>(EasyAlign)
 
 Now without going into visual mode, you can align the lines in the paragraph
 with `<Leader>aip=`, `<Leader>aip*|`, or `<Leader>aip:`. And if you have
-installed [vim-repeat](https://github.com/tpope/vim-repeat) by Tim Pope, the
+installed [repeat.vim](https://github.com/tpope/vim-repeat) by Tim Pope, the
 exact alignment can be repeated with `.` key.
 
 ### Live interactive mode
diff --git a/plugin/easy_align.vim b/plugin/easy_align.vim
index 986f0fe..7a360d9 100644
--- a/plugin/easy_align.vim
+++ b/plugin/easy_align.vim
@@ -31,12 +31,41 @@ command! -nargs=* -range -bang LiveEasyAlign <line1>,<line2>call easy_align#alig
 
 let s:last_command = 'EasyAlign'
 
+function! s:remember_visual(mode)
+  let s:last_visual = [a:mode, abs(line("'>") - line("'<")), abs(col("'>") - col("'<"))]
+endfunction
+
+function! s:repeat_visual()
+  let [mode, ldiff, cdiff] = s:last_visual
+  let cmd = 'normal! '.mode
+  if ldiff > 0
+    let cmd .= ldiff . 'j'
+  endif
+
+  let ve_save = &virtualedit
+  try
+    if mode == "\<C-V>"
+      if cdiff > 0
+        let cmd .= cdiff . 'l'
+      endif
+      set virtualedit+=block
+    endif
+    execute cmd.":\<C-r>=g:easy_align_last_command\<Enter>\<Enter>"
+    silent! call repeat#set("\<Plug>(EasyAlignRepeat)")
+  finally
+    if ve_save != &virtualedit
+      let &virtualedit = ve_save
+    endif
+  endtry
+endfunction
+
 function! s:generic_easy_align_op(type, vmode, live)
   let sel_save = &selection
   let &selection = "inclusive"
 
   if a:vmode
     let vmode = a:type
+    call s:remember_visual(vmode)
   else
     let tail = "\<C-c>"
     if a:type == 'line'
@@ -47,17 +76,14 @@ function! s:generic_easy_align_op(type, vmode, live)
       silent execute "normal! `[v`]".tail
     endif
     let vmode = ''
+    unlet! s:last_visual
   endif
 
   try
     if get(g:, 'easy_align_need_repeat', 0)
-      execute "'<,'>". s:last_command
+      execute "'<,'>". g:easy_align_last_command
     else
       '<,'>call easy_align#align('<bang>' == '!', a:live, vmode, '')
-      " Currently visual-mode operator is not repeatable
-      if !a:vmode
-        let s:last_command = g:easy_align_last_command
-      endif
     end
     silent! call repeat#set("\<Plug>(EasyAlignRepeat)")
   finally
@@ -74,12 +100,16 @@ function! s:live_easy_align_op(type, ...)
 endfunction
 
 function! s:easy_align_repeat()
-  try
-    let g:easy_align_need_repeat = 1
-    normal! .
-  finally
-    unlet! g:easy_align_need_repeat
-  endtry
+  if exists('s:last_visual')
+    call s:repeat_visual()
+  else
+    try
+      let g:easy_align_need_repeat = 1
+      normal! .
+    finally
+      unlet! g:easy_align_need_repeat
+    endtry
+  endif
 endfunction
 
 nnoremap <silent> <Plug>(EasyAlign) :set opfunc=<SID>easy_align_op<Enter>g@
diff --git a/test/interactive.vader b/test/interactive.vader
index d665b97..3c08a26 100644
--- a/test/interactive.vader
+++ b/test/interactive.vader
@@ -1528,12 +1528,30 @@ Given (Two paragraphs (requires vim-repeat)):
   cc = 2
   bbb = 3
   aaaa = 4
+  _____ = 5
 
 Do (Align and repeat):
   \<Space>Aip\<Enter>=
   6G
   .
 
+Expect:
+     a = 1
+    bb = 2
+   ccc = 3
+  dddd = 4
+
+      d = 1
+     cc = 2
+    bbb = 3
+   aaaa = 4
+  _____ = 5
+
+Do (Visual-mode operator is also repeatable):
+  vip\<Enter>\<Enter>=
+  6G
+  .
+
 Expect:
      a = 1
     bb = 2
@@ -1544,24 +1562,46 @@ Expect:
     cc = 2
    bbb = 3
   aaaa = 4
+  _____ = 5
 
-Do (Visual-mode operator is not repeatable and shouldn't affect normal-mode repeat):
-  \<Space>Aj=
-  3G
-  Vj\<Enter>\<Enter>=
+Given:
+  :: a : 1
+  :: bb : 2
+  :: ccc : 3
+  :: dd : 4
+  :: e : 5
+
+  :: :: a:1
+  :: :: b :2
+  :: :: ccc : 3
+  :: :: dd : 4
+  :: :: e : 5
+  :: :: f : 6
+
+Do (Blockwise-visual-operator is also repeatable):
+  fa
+  \<C-V>
+  f1
+  4j
+  \<Enter>:
   7G
+  fa
   .
 
 Expect:
-  a  = 1
-  bb = 2
-   ccc = 3
-  dddd = 4
+  :: a:   1
+  :: bb:  2
+  :: ccc:  3
+  :: dd:  4
+  :: e:   5
+
+  :: :: a:   1
+  :: :: b:   2
+  :: :: ccc:  3
+  :: :: dd:  4
+  :: :: e:   5
+  :: :: f : 6
 
-  d = 1
-  cc  = 2
-  bbb = 3
-  aaaa = 4
 ###########################################################
 Execute:
   Restore