Improved pasting with a count

autoload/whitespaste.vim

@@ -187,19 +187,19 @@ function! s:Paste(command)
   let original_value   = getreg(default_register, 1)
   let original_type    = getregtype(default_register)
 
-  call setreg(default_register, getreg(v:register, 1), getregtype(v:register))
-
   if v:count >= 1
     " we've been given a count, repeat v:count times
-    let repeat = v:count
+    let times = v:count
   else
     " no count, just run once
-    let repeat = 1
+    let times = 1
   endif
 
-  for i in range(repeat)
-    exe a:command
-  endfor
+  let value = getreg(v:register, 1)
+  let value = repeat(value, times)
+  call setreg(default_register, value, getregtype(v:register))
+
+  exe a:command
 
   call setreg(default_register, original_value, original_type)
 endfunction

spec/plugin/compatibility_spec.rb

@@ -4,18 +4,41 @@
   let(:filename) { 'test.txt' }
   let(:vim) { @vim }
 
-  it "respects pasting with a count" do
+  it "respects pasting after with a count" do
     set_file_contents <<-EOF
       one
+      two
     EOF
 
-    vim.command("call feedkeys('yy2p')")
+    vim.command("call feedkeys('yj2p')")
     vim.write
 
     assert_file_contents <<-EOF
       one
       one
+      two
       one
+      two
+      two
+    EOF
+  end
+
+  it "respects pasting before with a count" do
+    set_file_contents <<-EOF
+      one
+      two
+    EOF
+
+    vim.command("call feedkeys('yj2P')")
+    vim.write
+
+    assert_file_contents <<-EOF
+      one
+      two
+      one
+      two
+      one
+      two
     EOF
   end