Make ys maps atomic

Closes #253

Author: tpope

plugin/surround.vim

@@ -478,7 +478,11 @@ function! s:changesurround(...) " {{{1
  call s:dosurround(a,b,a:0 && a:1)
 endfunction " }}}1
 
-function! s:opfunc(type,...) " {{{1
+function! s:opfunc(type, ...) abort " {{{1
+ if a:type ==# 'setup'
+ let &opfunc = matchstr(expand('<sfile>'), '<SNR>\w\+$')
+ return 'g@'
+ endif
  let char = s:inputreplacement()
  if char == ""
  return s:beep()
@@ -541,8 +545,12 @@ function! s:opfunc(type,...) " {{{1
  endif
 endfunction
 
-function! s:opfunc2(arg)
- call s:opfunc(a:arg,1)
+function! s:opfunc2(...) abort
+ if !a:0 || a:1 ==# 'setup'
+ let &opfunc = matchstr(expand('<sfile>'), '<SNR>\w\+$')
+ return 'g@'
+ endif
+ call s:opfunc(a:1, 1)
 endfunction " }}}1
 
 function! s:closematch(str) " {{{1
@@ -565,11 +573,10 @@ nnoremap <silent> <Plug>SurroundRepeat .
 nnoremap <silent> <Plug>Dsurround :<C-U>call <SID>dosurround(<SID>inputtarget())<CR>
 nnoremap <silent> <Plug>Csurround :<C-U>call <SID>changesurround()<CR>
 nnoremap <silent> <Plug>CSurround :<C-U>call <SID>changesurround(1)<CR>
-nnoremap <silent> <Plug>Yssurround :<C-U>call <SID>opfunc(v:count1)<CR>
-nnoremap <silent> <Plug>YSsurround :<C-U>call <SID>opfunc2(v:count1)<CR>
-" <C-U> discards the numerical argument but there's not much we can do with it
-nnoremap <silent> <Plug>Ysurround :<C-U>set opfunc=<SID>opfunc<CR>g@
-nnoremap <silent> <Plug>YSurround :<C-U>set opfunc=<SID>opfunc2<CR>g@
+nnoremap <expr> <Plug>Yssurround <SID>opfunc('setup').'_'
+nnoremap <expr> <Plug>YSsurround <SID>opfunc2('setup').'_'
+nnoremap <expr> <Plug>Ysurround <SID>opfunc('setup')
+nnoremap <expr> <Plug>YSurround <SID>opfunc2('setup')
 vnoremap <silent> <Plug>VSurround :<C-U>call <SID>opfunc(visualmode(),visualmode() ==# 'V' ? 1 : 0)<CR>
 vnoremap <silent> <Plug>VgSurround :<C-U>call <SID>opfunc(visualmode(),visualmode() ==# 'V' ? 0 : 1)<CR>
 inoremap <silent> <Plug>Isurround <C-R>=<SID>insert()<CR>