Fix: getTokenAfter was wrong

lib/token-store/index.js

@@ -45,7 +45,7 @@ const PUBLIC_METHODS = Object.freeze([
  * Creates the map from locations to indices in `tokens`.
  *
  * The first/last location of tokens is mapped to the index of the token.
- * The first/last location of comments is mapped to the index of the previous token of each comment.
+ * The first/last location of comments is mapped to the index of the next token of each comment.
  *
  * @param {Token[]} tokens - The array of tokens.
  * @param {Comment[]} comments - The array of comments.

lib/token-store/utils.js

@@ -61,7 +61,15 @@ exports.getFirstIndex = function getFirstIndex(tokens, indexMap, startLoc) {
         return indexMap[startLoc];
     }
     if ((startLoc - 1) in indexMap) {
-        return indexMap[startLoc - 1] + 1;
+        const index = indexMap[startLoc - 1];
+        const token = (index >= 0 && index < tokens.length) ? tokens[index] : null;
+
+        // For the map of "comment's location -> token's index", it points the next token of a comment.
+        // In that case, +1 is unnecessary.
+        if (token && token.range[0] >= startLoc) {
+            return index;
+        }
+        return index + 1;
     }
     return 0;
 };
@@ -81,7 +89,15 @@ exports.getLastIndex = function getLastIndex(tokens, indexMap, endLoc) {
         return indexMap[endLoc] - 1;
     }
     if ((endLoc - 1) in indexMap) {
-        return indexMap[endLoc - 1];
+        const index = indexMap[endLoc - 1];
+        const token = (index >= 0 && index < tokens.length) ? tokens[index] : null;
+
+        // For the map of "comment's location -> token's index", it points the next token of a comment.
+        // In that case, -1 is necessary.
+        if (token && token.range[1] > endLoc) {
+            return index - 1;
+        }
+        return index;
     }
     return tokens.length - 1;
 };

tests/lib/token-store.js

@@ -271,6 +271,33 @@ describe("TokenStore", () => {
             );
         });
 
+        it("should retrieve the previous node if the comment at the end of source code is specified.", () => {
+            const code = "a + b /*comment*/";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenBefore(ast.comments[0]);
+
+            assert.strictEqual(token.value, "b");
+        });
+
+        it("should retrieve the previous comment if the first token is specified.", () => {
+            const code = "/*comment*/ a + b";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenBefore(ast.tokens[0], { includeComments: true });
+
+            assert.strictEqual(token.value, "comment");
+        });
+
+        it("should retrieve null if the first comment is specified.", () => {
+            const code = "/*comment*/ a + b";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenBefore(ast.comments[0], { includeComments: true });
+
+            assert.strictEqual(token, null);
+        });
+
     });
 
     describe("when calling getTokensAfter", () => {
@@ -417,6 +444,33 @@ describe("TokenStore", () => {
             );
         });
 
+        it("should retrieve the next node if the comment at the first of source code is specified.", () => {
+            const code = "/*comment*/ a + b";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenAfter(ast.comments[0]);
+
+            assert.strictEqual(token.value, "a");
+        });
+
+        it("should retrieve the next comment if the last token is specified.", () => {
+            const code = "a + b /*comment*/";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenAfter(ast.tokens[2], { includeComments: true });
+
+            assert.strictEqual(token.value, "comment");
+        });
+
+        it("should retrieve null if the last comment is specified.", () => {
+            const code = "a + b /*comment*/";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const token = tokenStore.getTokenAfter(ast.comments[0], { includeComments: true });
+
+            assert.strictEqual(token, null);
+        });
+
     });
 
     describe("when calling getFirstTokens", () => {
@@ -567,6 +621,35 @@ describe("TokenStore", () => {
             );
         });
 
+        it("should retrieve the first comment if the comment is at the last of nodes", () => {
+            const code = "a + b\n/*comment*/ c + d";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+
+            // Actually, the first of nodes is always tokens, not comments.
+            // But I think this test case is needed for completeness.
+            const token = tokenStore.getFirstToken(
+                { range: [ast.comments[0].range[0], ast.tokens[5].range[1]] },
+                { includeComments: true }
+            );
+
+            assert.strictEqual(token.value, "comment");
+        });
+
+        it("should retrieve the first token (without includeComments option) if the comment is at the last of nodes", () => {
+            const code = "a + b\n/*comment*/ c + d";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+
+            // Actually, the first of nodes is always tokens, not comments.
+            // But I think this test case is needed for completeness.
+            const token = tokenStore.getFirstToken(
+                { range: [ast.comments[0].range[0], ast.tokens[5].range[1]] }
+            );
+
+            assert.strictEqual(token.value, "c");
+        });
+
     });
 
     describe("when calling getLastTokens", () => {
@@ -721,6 +804,35 @@ describe("TokenStore", () => {
             );
         });
 
+        it("should retrieve the last comment if the comment is at the last of nodes", () => {
+            const code = "a + b /*comment*/\nc + d";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+
+            // Actually, the last of nodes is always tokens, not comments.
+            // But I think this test case is needed for completeness.
+            const token = tokenStore.getLastToken(
+                { range: [ast.tokens[0].range[0], ast.comments[0].range[1]] },
+                { includeComments: true }
+            );
+
+            assert.strictEqual(token.value, "comment");
+        });
+
+        it("should retrieve the last token (without includeComments option) if the comment is at the last of nodes", () => {
+            const code = "a + b /*comment*/\nc + d";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+
+            // Actually, the last of nodes is always tokens, not comments.
+            // But I think this test case is needed for completeness.
+            const token = tokenStore.getLastToken(
+                { range: [ast.tokens[0].range[0], ast.comments[0].range[1]] }
+            );
+
+            assert.strictEqual(token.value, "b");
+        });
+
     });
 
     describe("when calling getFirstTokensBetween", () => {
@@ -1066,4 +1178,79 @@ describe("TokenStore", () => {
 
     });
 
+    describe("when calling getFirstToken & getTokenAfter", () => {
+        it("should retrieve all tokens and comments in the node", () => {
+            const code = "(function(a, /*b,*/ c){})";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const tokens = [];
+            let token = tokenStore.getFirstToken(ast);
+
+            while (token) {
+                tokens.push(token);
+                token = tokenStore.getTokenAfter(token, { includeComments: true });
+            }
+
+            check(
+                tokens,
+                ["(", "function", "(", "a", ",", "b,", "c", ")", "{", "}", ")"]
+            );
+        });
+
+        it("should retrieve all tokens and comments in the node (no spaces)", () => {
+            const code = "(function(a,/*b,*/c){})";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const tokens = [];
+            let token = tokenStore.getFirstToken(ast);
+
+            while (token) {
+                tokens.push(token);
+                token = tokenStore.getTokenAfter(token, { includeComments: true });
+            }
+
+            check(
+                tokens,
+                ["(", "function", "(", "a", ",", "b,", "c", ")", "{", "}", ")"]
+            );
+        });
+    });
+
+    describe("when calling getLastToken & getTokenBefore", () => {
+        it("should retrieve all tokens and comments in the node", () => {
+            const code = "(function(a, /*b,*/ c){})";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const tokens = [];
+            let token = tokenStore.getLastToken(ast);
+
+            while (token) {
+                tokens.push(token);
+                token = tokenStore.getTokenBefore(token, { includeComments: true });
+            }
+
+            check(
+                tokens.reverse(),
+                ["(", "function", "(", "a", ",", "b,", "c", ")", "{", "}", ")"]
+            );
+        });
+
+        it("should retrieve all tokens and comments in the node (no spaces)", () => {
+            const code = "(function(a,/*b,*/c){})";
+            const ast = espree.parse(code, { loc: true, range: true, tokens: true, comment: true });
+            const tokenStore = new TokenStore(ast.tokens, ast.comments);
+            const tokens = [];
+            let token = tokenStore.getLastToken(ast);
+
+            while (token) {
+                tokens.push(token);
+                token = tokenStore.getTokenBefore(token, { includeComments: true });
+            }
+
+            check(
+                tokens.reverse(),
+                ["(", "function", "(", "a", ",", "b,", "c", ")", "{", "}", ")"]
+            );
+        });
+    });
 });