diff --git a/autoload/EasyMotion.vim b/autoload/EasyMotion.vim
index df39bf2..479be65 100644
--- a/autoload/EasyMotion.vim
+++ b/autoload/EasyMotion.vim
@@ -739,14 +739,6 @@ function! s:GetVisualStartPosition(c_pos, v_start, v_end, search_direction) "{{{
 endfunction "}}}
 
 " -- Others ------------------------------
-function! s:is_cmdwin() "{{{
-  return bufname('%') ==# '[Command Line]'
-endfunction "}}}
-function! s:should_use_wundo() "{{{
-    " wundu cannot use in command-line window and
-    " unless undolist is not empty
-    return ! s:is_cmdwin() && undotree().seq_last != 0
-endfunction "}}}
 function! s:handleEmpty(input, visualmode) "{{{
     " if empty, reselect and return 1
     if empty(a:input)
@@ -1104,12 +1096,8 @@ function! s:PromptUser(groups) "{{{
     " }}}
 
     " -- Put labels on targets & Get User Input & Restore all {{{
-    " Save undo tree {{{
-    let s:undo_file = tempname()
-    if s:should_use_wundo()
-        execute "wundo" s:undo_file
-    endif
-    "}}}
+    " Save undo tree
+    let undo_lock = EasyMotion#undo#save()
     try
         " Set lines with markers {{{
         call s:SetLines(lines_items, 'marker')
@@ -1156,21 +1144,8 @@ function! s:PromptUser(groups) "{{{
             \ )
         " }}}
 
-        " Restore undo tree {{{
-        if s:should_use_wundo() && filereadable(s:undo_file)
-            silent execute "rundo" s:undo_file
-            call delete(s:undo_file)
-            unlet s:undo_file
-        else
-            " Break undo history (undobreak)
-            let old_undolevels = &undolevels
-            set undolevels=-1
-            keepjumps call setline('.', getline('.'))
-            let &undolevels = old_undolevels
-            unlet old_undolevels
-            " FIXME: Error occur by GundoToggle for undo number 2 is empty
-            keepjumps call setline('.', getline('.'))
-        endif "}}}
+        " Restore undo tree
+        call undo_lock.restore()
 
         redraw
     endtry "}}}
diff --git a/autoload/EasyMotion/undo.vim b/autoload/EasyMotion/undo.vim
new file mode 100644
index 0000000..a14ef88
--- /dev/null
+++ b/autoload/EasyMotion/undo.vim
@@ -0,0 +1,50 @@
+let s:Buffer = vital#easymotion#import('Vim.Buffer')
+
+function! EasyMotion#undo#save() abort
+  return s:undo_lock.save()
+endfunction
+
+let s:undo_lock = {}
+
+function! s:undo_lock.save() abort
+  let undo = deepcopy(self)
+  call undo._save()
+  return undo
+endfunction
+
+function! s:undo_lock._save() abort
+  if undotree().seq_last == 0
+    " if there are no undo history, disable undo feature by setting
+    " 'undolevels' to -1 and restore it.
+    let self.save_undolevels = &l:undolevels
+    let &l:undolevels = -1
+  elseif !s:Buffer.is_cmdwin()
+    " command line window doesn't support :wundo.
+    let self.undofile = tempname()
+    execute 'wundo!' self.undofile
+  else
+    let self.is_cmdwin = s:TRUE
+  endif
+endfunction
+
+function! s:undo_lock.restore() abort
+  if has_key(self, 'save_undolevels')
+    let &l:undolevels = self.save_undolevels
+  endif
+  if has_key(self, 'undofile') && filereadable(self.undofile)
+    silent execute 'rundo' self.undofile
+    call delete(self.undofile)
+  endif
+  if has_key(self, 'is_cmdwin')
+    " XXX: it breaks undo history. AFAIK, there are no way to save and restore
+    " undo history in commandline window.
+    call self.undobreak()
+  endif
+endfunction
+
+function! s:undo_lock.undobreak() abort
+  let old_undolevels = &l:undolevels
+  setlocal undolevels=-1
+  keepjumps call setline('.', getline('.'))
+  let &l:undolevels = old_undolevels
+endfunction