Improve keycode merging (#221)

autoload/which_key.vim

@@ -9,13 +9,36 @@ let s:TYPE = {
       \ 'string':  type(''),
       \ 'funcref': type(function('call'))
       \ }
+" For the purpose of mapping a keypress to internal data-structures
+let s:KEYCODES = {
+      \ "\<BS>": '<BS>',
+      \ "\<Tab>": '<Tab>',
+      \ "\<CR>": '<CR>',
+      \ "\<Esc>": '<Esc>',
+      \ "\<Del>": '<Del>'
+      \ }
+" For the purposes of merging identical keycodes in internal data-structures
+let s:MERGE_INTO = {
+      \ '<Space>': ' ',
+      \ '<C-H>': '<BS>',
+      \ '<C-I>': '<Tab>',
+      \ '<C-M>': '<CR>',
+      \ '<Return>': '<CR>',
+      \ '<Enter>': '<CR>',
+      \ '<C-[>': '<Esc>',
+      \ '<lt>': '<',
+      \ '<Bslash>': '\',
+      \ '<Bar>': '|'
+      \ }
+let s:REQUIRES_REGEX_ESCAPE = ['$', '*', '~', '.']
+
 let g:which_key#TYPE = s:TYPE
 
 let s:should_note_winid = exists('*win_getid')
 
 function! which_key#register(prefix, dict, ...) abort
-  let key = a:prefix ==? '<Space>' ? ' ' :
-    \ (a:prefix ==? '<C-I>' ? '<Tab>' : a:prefix)
+  let key = has_key(s:MERGE_INTO, a:prefix) ?
+    \ s:MERGE_INTO[a:prefix] : a:prefix
   let val = a:dict
   if a:0 == 1
     call extend(s:desc[a:1], {key:val})
@@ -30,7 +53,10 @@ function! s:handle_char_on_start_is_ok(c) abort
   if which_key#char_handler#is_exit_code(a:c)
     return 1
   endif
-  let char = a:c == 9 ? '<Tab>' : nr2char(a:c)
+  let char = type(a:c) == s:TYPE.number ? nr2char(a:c) : a:c
+  if has_key(s:KEYCODES, char)
+    let char = s:KEYCODES[char]
+  endif
   let s:which_key_trigger .= ' '.char
   let next_level = get(s:runtime, char)
   let ty = type(next_level)
@@ -52,6 +78,7 @@ endfunction
 function! which_key#start(vis, bang, prefix) " {{{
   let s:vis = a:vis ? 'gv' : ''
   let mode = a:vis ? 'v' : 'n'
+  let prefix = a:prefix
   let s:count = v:count != 0 ? v:count : ''
   let s:which_key_trigger = ''
 
@@ -60,14 +87,19 @@ function! which_key#start(vis, bang, prefix) " {{{
   endif
 
   if a:bang
-    for kv in keys(a:prefix)
+    for kv in keys(prefix)
       call s:cache_key(mode, kv)
     endfor
-    let s:runtime = deepcopy(a:prefix)
+    let s:runtime = deepcopy(prefix)
     call s:merge(s:runtime, s:cache[mode])
   else
-    let key = a:prefix ==? '<Space>' ? ' ' :
-      \ (a:prefix ==? '<C-I>' ? '<Tab>' : a:prefix)
+    if has_key(s:KEYCODES, prefix)
+      let prefix = s:KEYCODES[prefix]
+    endif
+    if has_key(s:MERGE_INTO, prefix)
+      let prefix = s:MERGE_INTO[prefix]
+    endif
+    let key = prefix
     let s:which_key_trigger = key ==# ' ' ? '<space>' : key
     call s:cache_key(mode, key)
 
@@ -127,23 +159,14 @@ endfunction
 function! s:merge(target, native) " {{{
   let target = a:target
   let native = a:native
-  " <C-І> is merged into <Tab>, '<Space>' is merged into ' '
-  if has_key(target, '<C-I>')
-    if has_key(target, '<Tab>')
-      call extend(target['<Tab>'], target['<C-I>'], 'keep')
-    else
-      let target['<Tab>'] = target['<C-I>']
-    endif
-    call remove(target, '<C-I>')
-  endif
-  if has_key(target, '<Space>')
-    if has_key(target, ' ')
-      call extend(target[' '], target['<Space>'], 'keep')
-    else
-      let target[' '] = target['<Space>']
-    endif
-    call remove(target, '<Space>')
-  endif
+  " e.g. <C-І> is merged into <Tab>, '<Space>' is merged into ' '
+  call map(target, {k,v ->
+  \ has_key(s:MERGE_INTO, k) ?
+  \   (has_key(target, s:MERGE_INTO[k]) ?
+  \     extend(target[s:MERGE_INTO[k]], target[k], 'keep') :
+  \     extend(target, {s:MERGE_INTO[k]: target[k]})) :
+  \   v})
+  call filter(target, {k,_ -> !has_key(s:MERGE_INTO, k)})
   for [k, V] in items(target)
     " Support a `Dictionary-function` for on-the-fly mappings
     while type(target[k]) == s:TYPE.funcref
@@ -191,15 +214,10 @@ function! s:echo_prompt() abort
 endfunction
 
 function! s:has_children(input) abort
-  " TODO: escape properly, E114: Missing quote: "^\"
-  if a:input ==# '\'
-    let group = map(keys(s:runtime), 'v:val =~# "^\'.a:input.'"')
-  elseif a:input ==# '"'
-    let group = map(keys(s:runtime), "v:val =~# '^".a:input."'")
-  elseif a:input == '.'
-    let group = map(keys(s:runtime), "v:val =~# '^\\.'")
+  if index(s:REQUIRES_REGEX_ESCAPE, a:input) != -1
+    let group = map(keys(s:runtime), {_,v -> v =~# '^\'.a:input})
   else
-    let group = map(keys(s:runtime), 'v:val =~# "^'.a:input.'"')
+    let group = map(keys(s:runtime), {_,v -> v =~# '^'.a:input})
   endif
   return len(filter(group, 'v:val == 1')) > 1
 endfunction
@@ -241,7 +259,7 @@ function! s:getchar() abort
   endif
 
   " Allow <BS> to go back to the upper level.
-  if c ==# "\<BS>"
+  if c == "\<BS>"
     call s:show_upper_level_mappings()
     return ''
   endif
@@ -296,12 +314,12 @@ endfunction
 function! s:handle_input(input) " {{{
   let ty = type(a:input)
 
-  if ty ==? s:TYPE.dict
+  if ty == s:TYPE.dict
     call s:show_next_level_mappings(a:input)
     return
   endif
 
-  if ty ==? s:TYPE.list
+  if ty == s:TYPE.list
     call which_key#window#close()
     call s:execute(a:input[0])
   elseif g:which_key_fallback_to_native_key

autoload/which_key/char_handler.vim

@@ -57,8 +57,7 @@ endfor
 
 function! which_key#char_handler#parse_raw(raw_char)
   if type(a:raw_char) == g:which_key#TYPE.number
-    " <Tab>, <C-I> = 9
-    return a:raw_char == 9 ? '<Tab>' : nr2char(a:raw_char)
+    return nr2char(a:raw_char)
   elseif has_key(s:special_keys, a:raw_char)
     " Special characters
     return s:special_keys[a:raw_char]
@@ -70,10 +69,13 @@ endfunction
 function! s:initialize_exit_code() abort
   if exists('g:which_key_exit')
     let ty = type(g:which_key_exit)
-    if ty == s:TYPE.number || ty == s:TYPE.string
+    if ty == s:TYPE.number
+      let s:exit_code = [nr2char(g:which_key_exit)]
+    elseif ty == s:TYPE.string
       let s:exit_code = [g:which_key_exit]
     elseif ty == s:TYPE.list
-      let s:exit_code = g:which_key_exit
+      let s:exit_code = map(g:which_key_exit, {e ->
+        \ type(e) == s:TYPE.number ? nr2char(e) : e})
     else
       echohl ErrorMsg
       echon '[which-key] '.g:which_key_exit.' is invalid for option g:which_key_exit'
@@ -82,41 +84,31 @@ function! s:initialize_exit_code() abort
     endif
   else
     " <Esc>, <C-[>: 27
-    let s:exit_code = [27]
+    let s:exit_code = ["\<Esc>"]
   endif
 endfunction
 
 if !exists('s:exit_code')
   call s:initialize_exit_code()
 endif
 
-" Argument: number
+" Argument: v:t_number or v:t_string as returned by getchar()
 function! which_key#char_handler#is_exit_code(raw_char) abort
-  for e in s:exit_code
-    let ty = type(e)
-    if ty == s:TYPE.number && e == a:raw_char
-      return 1
-    elseif ty == s:TYPE.string && e == nr2char(a:raw_char)
-      return 1
-    endif
-  endfor
-
-  return 0
+  return -1 != index(s:exit_code,
+    \ type(a:raw_char) == s:TYPE.number ? nr2char(a:raw_char) : a:raw_char)
 endfunction
 
 " Returns true if timed out
 function! s:wait_with_timeout(timeout) abort
   let timeout = a:timeout
-  while timeout >= 0
+  while timeout > 0
     if getchar(1)
       return 0
     endif
-    if timeout > 0
-      sleep 20m
-    endif
+    sleep 20m
     let timeout -= 20
   endwhile
-  return 1
+  return !getchar(1)
 endfunction
 
 " Wait timtout to see if there are more input chars.