PR review fixes

Author: mdevils

lib/js-file.js

@@ -144,22 +144,25 @@ JsFile.prototype = {
     setWhitespaceBefore: function(token, whitespace) {
         var whitespaceToken = this.getPrevToken(token, {includeWhitespace: true});
         if (whitespaceToken && whitespaceToken.type === 'Whitespace') {
+
+            // Modifying already existing token.
             if (whitespace === '') {
                 this.removeToken(whitespaceToken);
             } else {
                 whitespaceToken.value = whitespace;
             }
-        } else {
-            if (whitespace === '') {
-                return;
-            }
-            var tokenIndex = this._tokens.indexOf(token);
+        } else if (whitespace !== '') {
+            var tokenIndex = token._tokenIndex;
+
+            // Adding a token before specified one.
             this._tokens.splice(tokenIndex, 0, {
                 type: 'Whitespace',
                 value: whitespace,
                 isWhitespace: true
             });
-            for (var i = tokenIndex, l = this._tokens.length; i < l; i++) {
+
+            // Quickly updating modified token order
+            for (var i = tokenIndex; i < this._tokens.length; i++) {
                 this._tokens[i]._tokenIndex = i;
             }
         }
@@ -362,7 +365,7 @@ JsFile.prototype = {
      * @returns {Object}
      */
     getFirstToken: function(options) {
-        return this._getNextTokenFromIndex(0, options);
+        return this._getTokenFromIndex(0, 1, options);
     },
 
     /**
@@ -374,70 +377,36 @@ JsFile.prototype = {
      * @returns {Object}
      */
     getLastToken: function(options) {
-        return this._getPrevTokenFromIndex(this._tokens.length - 1, options);
+        return this._getTokenFromIndex(this._tokens.length - 1, -1, options);
     },
 
     /**
-     * Returns the first token before the given index.
+     * Returns the first token after the given using direction and specified conditions.
      *
      * @param {Number} index
+     * @param {Number} direction `1` - forward or `-1` - backwards
      * @param {Object} [options]
      * @param {Boolean} [options.includeComments=false]
      * @param {Boolean} [options.includeWhitespace=false]
      * @returns {Object|null}
      */
-    _getPrevTokenFromIndex: function(index, options) {
-        if (index < 0) {
-            return null;
-        }
+    _getTokenFromIndex: function(index, direction, options) {
+        while (true) {
+            var followingToken = this._tokens[index];
 
-        do {
-            var prevToken = this._tokens[index];
-
-            if (!prevToken) {
+            if (!followingToken) {
                 return null;
             }
 
             if (
-                (!prevToken.isComment || (options && options.includeComments)) &&
-                (!prevToken.isWhitespace || (options && options.includeWhitespace))
+                (!followingToken.isComment || (options && options.includeComments)) &&
+                (!followingToken.isWhitespace || (options && options.includeWhitespace))
             ) {
-                return prevToken;
+                return followingToken;
             }
-        } while (--index >= 0);
 
-        return null;
-    },
-
-    /**
-     * Returns the first token after the given.
-     *
-     * @param {Number} index
-     * @param {Object} [options]
-     * @param {Boolean} [options.includeComments=false]
-     * @param {Boolean} [options.includeWhitespace=false]
-     * @returns {Object|null}
-     */
-    _getNextTokenFromIndex: function(index, options) {
-        if (index >= this._tokens.length) {
-            return null;
+            index += direction;
         }
-
-        do {
-            var nextToken = this._tokens[index];
-            if (!nextToken) {
-                return null;
-            }
-
-            if (
-                (!nextToken.isComment || (options && options.includeComments)) &&
-                (!nextToken.isWhitespace || (options && options.includeWhitespace))
-            ) {
-                return nextToken;
-            }
-        } while (++index < this._tokens.length);
-
-        return null;
     },
 
     /**
@@ -450,7 +419,7 @@ JsFile.prototype = {
      * @returns {Object|null}
      */
     getPrevToken: function(token, options) {
-        return this._getPrevTokenFromIndex(token._tokenIndex - 1, options);
+        return this._getTokenFromIndex(token._tokenIndex - 1, -1, options);
     },
 
     /**
@@ -463,7 +432,7 @@ JsFile.prototype = {
      * @returns {Object|null}
      */
     getNextToken: function(token, options) {
-        return this._getNextTokenFromIndex(token._tokenIndex + 1, options);
+        return this._getTokenFromIndex(token._tokenIndex + 1, 1, options);
     },
 
     /**

test/specs/js-file.js

@@ -1490,4 +1490,55 @@ describe('js-file', function() {
             expect(spy.getCall(1).args[0].value).to.equal('2');
         });
     });
+
+    describe('getWhitespaceBefore', function() {
+        it('should return whitespace before the first token', function() {
+            var file = createJsFile('  \nx;');
+            expect(file.getWhitespaceBefore(file.getFirstToken())).to.equal('  \n');
+        });
+
+        it('should return whitespace before the EOF', function() {
+            var file = createJsFile('\n  ');
+            expect(file.getWhitespaceBefore(file.getLastToken())).to.equal('\n  ');
+        });
+
+        it('should return empty for tokens without whitespace before', function() {
+            var file = createJsFile('x');
+            expect(file.getWhitespaceBefore(file.getFirstToken())).to.equal('');
+        });
+    });
+
+    describe('setWhitespaceBefore', function() {
+        it('should alter existing whitespace token', function() {
+            var file = createJsFile('  x');
+            expect(file.getTokens().length).to.equal(3);
+            file.setWhitespaceBefore(file.getFirstToken(), '\n');
+            expect(file.getTokens().length).to.equal(3);
+            expect(file.getTokens()[0].value).to.equal('\n');
+        });
+
+        it('should insert new whitespace token', function() {
+            var file = createJsFile('x');
+            expect(file.getTokens().length).to.equal(2);
+            file.setWhitespaceBefore(file.getFirstToken(), '\n');
+            expect(file.getTokens().length).to.equal(3);
+            expect(file.getTokens()[0].value).to.equal('\n');
+        });
+
+        it('should drop whitespace token', function() {
+            var file = createJsFile('  x');
+            expect(file.getTokens().length).to.equal(3);
+            file.setWhitespaceBefore(file.getFirstToken(), '');
+            expect(file.getTokens().length).to.equal(2);
+            expect(file.getTokens()[0].type).to.equal('Identifier');
+        });
+
+        it('should ignore already absent whitespace token', function() {
+            var file = createJsFile('x');
+            expect(file.getTokens().length).to.equal(2);
+            file.setWhitespaceBefore(file.getFirstToken(), '');
+            expect(file.getTokens().length).to.equal(2);
+            expect(file.getTokens()[0].type).to.equal('Identifier');
+        });
+    });
 });