From 4e7d573c14ea7bc7bd87448272c1cd54ddfdbfa5 Mon Sep 17 00:00:00 2001
From: Tim Pope <code@tpope.net>
Date: Tue, 29 May 2018 17:37:25 -0400
Subject: [PATCH] Make ys maps atomic

Closes https://github.com/tpope/vim-surround/issues/253
---
 plugin/surround.vim | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/plugin/surround.vim b/plugin/surround.vim
index 46f3252..a91bc30 100644
--- a/plugin/surround.vim
+++ b/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>