From 1da7d97db32196bc8c2d79140eef313735c8ac46 Mon Sep 17 00:00:00 2001
From: Colin Caine <cmc@cs.man.ac.uk>
Date: Fri, 31 May 2019 15:22:25 +0100
Subject: [PATCH] Fix #1606

Introduces tests for this too.
---
 src/lib/keyseq.test.ts | 17 +++++++++++++++++
 src/lib/keyseq.ts      | 29 +++++++++++++++++++++--------
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/src/lib/keyseq.test.ts b/src/lib/keyseq.test.ts
index dbe4ec7d..224e4929 100644
--- a/src/lib/keyseq.test.ts
+++ b/src/lib/keyseq.test.ts
@@ -17,6 +17,7 @@ function mk(k, mod?: ks.KeyModifiers) {
         [mks(":"), "fillcmdline"],
         [mks("Av"), "whatever"],
         [mks("i<c-j>"), "testmods"],
+        [mks("0"), "panleft"],
     ])
     // Keymap for negative tests
     const keymap2 = new Map([
@@ -53,6 +54,22 @@ function mk(k, mod?: ks.KeyModifiers) {
             { value: "testmods", isMatch: true },
         ],
 
+        // Test count behaviour
+        // Zero isn't a prefix.
+        [[mks("0g"), keymap2], { keys: mks("g"), isMatch: true }],
+        [[mks("0gg"), keymap2], { value: "scrolltop", exstr: "scrolltop", isMatch: true, numericPrefix: undefined }],
+        // If zero is a map, then it should still work
+        [[mks("0"), keymap], { value: "panleft", exstr: "panleft", isMatch: true, numericPrefix: undefined }],
+
+        // Do match numbers starting with a non-zero
+        [[mks("2gg"), keymap2], { value: "scrolltop", exstr: "scrolltop 2", isMatch: true, numericPrefix: 2 }],
+        [[mks("20gg"), keymap2], { value: "scrolltop", exstr: "scrolltop 20", isMatch: true, numericPrefix: 20 }],
+        // If zero is a map, then you can still use zero in counts.
+        [[mks("20gg"), keymap], { value: "scrolltop", exstr: "scrolltop 20", isMatch: true, numericPrefix: 20 }],
+
+        // Don't match function keys as counts.
+        [[mks("<F2>gg"), keymap2], { value: "scrolltop", exstr: "scrolltop", isMatch: true }],
+
         // Test prefix problems
         [[mks("g"), keymap2], { keys: mks("g"), isMatch: true }],
         [[mks("go"), keymap2], { keys: mks("go"), isMatch: true }],
diff --git a/src/lib/keyseq.ts b/src/lib/keyseq.ts
index b7fbc844..1744edbe 100644
--- a/src/lib/keyseq.ts
+++ b/src/lib/keyseq.ts
@@ -21,7 +21,7 @@
 */
 
 /** */
-import { drop, filter, find, izip, take } from "@src/lib/itertools"
+import { filter, find, izip } from "@src/lib/itertools"
 import { Parser } from "@src/lib/nearley_utils"
 import * as bracketexpr_grammar from "@src/grammars/.bracketexpr.generated"
 const bracketexpr_parser = new Parser(bracketexpr_grammar)
@@ -114,6 +114,23 @@ export interface ParserResponse {
     numericPrefix?: number
 }
 
+function splitNumericPrefix(keyseq: KeyEventLike[]): [KeyEventLike[], KeyEventLike[]] {
+    // If the first key is in 1:9, partition all numbers until you reach a non-number.
+    if ([1, 2, 3, 4, 5, 6, 7, 8, 9].includes(Number(keyseq[0].key))) {
+        const prefix = [keyseq[0]]
+        for (const ke of keyseq.slice(1)) {
+            if ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9].includes(Number(ke.key)))
+                prefix.push(ke)
+            else
+                break
+        }
+        const rest = keyseq.slice(prefix.length)
+        return [prefix, rest]
+    } else {
+        return [[], keyseq]
+    }
+}
+
 export function parse(keyseq: KeyEventLike[], map: KeyMap): ParserResponse {
     // Remove bare modifiers
     keyseq = keyseq.filter(
@@ -122,8 +139,8 @@ export function parse(keyseq: KeyEventLike[], map: KeyMap): ParserResponse {
     )
 
     // Split into numeric prefix and non-numeric suffix
-    let numericPrefix = Array.from(take(keyseq, isNumeric))
-    keyseq = Array.from(drop(keyseq, isNumeric))
+    let numericPrefix: KeyEventLike[]
+    [numericPrefix, keyseq] = splitNumericPrefix(keyseq)
 
     // If keyseq is a prefix of a key in map, proceed, else try dropping keys
     // from keyseq until it is empty or is a prefix.
@@ -146,6 +163,7 @@ export function parse(keyseq: KeyEventLike[], map: KeyMap): ParserResponse {
                 value: perfect[1],
                 exstr: perfect[1] + numericPrefixToExstrSuffix(numericPrefix),
                 isMatch: true,
+                numericPrefix: numericPrefix.length ? Number(numericPrefix.map(ke => ke.key).join("")) : undefined,
             }
         } catch (e) {
             if (!(e instanceof RangeError)) throw e
@@ -344,11 +362,6 @@ export function isSimpleKey(keyEvent: KeyEventLike) {
     return !(keyEvent.key.length > 1 || hasNonShiftModifiers(keyEvent))
 }
 
-const NUMERIC_REG = /\d/
-export function isNumeric(keyEvent: KeyEventLike) {
-    return !hasModifiers(keyEvent) && keyEvent.key.match(NUMERIC_REG)
-}
-
 function numericPrefixToExstrSuffix(numericPrefix: KeyEventLike[]) {
     if (numericPrefix.length > 0) {
         return " " + numericPrefix.map(k => k.key).join("")