From d6b1e36d4908ca30f9f5d3a527c9c5741aefd424 Mon Sep 17 00:00:00 2001
From: Logan Smyth <loganfsmyth@gmail.com>
Date: Sat, 2 Jul 2016 15:31:30 -0700
Subject: [PATCH] Make the 'catchUp' call implicit to source location updates.

---
 packages/babel-generator/src/buffer.js             | 14 ++++++--------
 packages/babel-generator/src/printer.js            |  6 +-----
 .../expected.js                                    |  4 +++-
 .../return-with-retainlines-option/expected.js     |  7 ++++---
 4 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/packages/babel-generator/src/buffer.js b/packages/babel-generator/src/buffer.js
index 801042dfeb..a5fc0e432c 100644
--- a/packages/babel-generator/src/buffer.js
+++ b/packages/babel-generator/src/buffer.js
@@ -37,14 +37,10 @@ export default class Buffer {
   buf: string;
   last: string;
 
-  /**
-   * Description
-   */
-
-  catchUp(node: Object) {
+  _catchUp(){
     // catch up to this nodes newline if we're behind
-    if (node.loc && this.format.retainLines && this.buf) {
-      while (this.getCurrentLine() < node.loc.start.line) {
+    if (this.format.retainLines && this._sourcePosition.line !== null) {
+      while (this.getCurrentLine() < this._sourcePosition.line) {
         this.push("\n");
       }
     }
@@ -276,6 +272,8 @@ export default class Buffer {
     this._sourcePosition.line = pos ? pos.line : null;
     this._sourcePosition.column = pos ? pos.column : null;
     this._sourcePosition.filename = loc && loc.filename || null;
+
+    this._catchUp();
   }
 
   /**
@@ -283,7 +281,7 @@ export default class Buffer {
    */
 
   withSource(prop: string, loc: Location, cb: () => void) {
-    if (!this.opts.sourceMaps) return cb();
+    if (!this.opts.sourceMaps && !this.format.retainLines) return cb();
 
     // Use the call stack to manage a stack of "source location" data.
     let originalLine = this._sourcePosition.line;
diff --git a/packages/babel-generator/src/printer.js b/packages/babel-generator/src/printer.js
index 79fcb230fa..f14ca8f99e 100644
--- a/packages/babel-generator/src/printer.js
+++ b/packages/babel-generator/src/printer.js
@@ -45,8 +45,6 @@ export default class Printer extends Buffer {
 
     this.printLeadingComments(node, parent);
 
-    this.catchUp(node);
-
     this._printNewline(true, node, parent, opts);
 
     if (opts.before) opts.before();
@@ -264,9 +262,7 @@ export default class Printer extends Buffer {
     }
 
     // Exclude comments from source mappings since they will only clutter things.
-    this.withSource(null, null, () => {
-      this.catchUp(comment);
-
+    this.withSource("start", comment.loc, () => {
       // whitespace before
       this.newline(this.whitespace.getNewlinesBefore(comment));
 
diff --git a/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-and-compact-option/expected.js b/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-and-compact-option/expected.js
index 7441245e8e..d13817e5aa 100644
--- a/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-and-compact-option/expected.js
+++ b/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-and-compact-option/expected.js
@@ -1,4 +1,6 @@
 function foo(l){
 return(
 
-l);}
+l);
+
+}
diff --git a/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-option/expected.js b/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-option/expected.js
index cc911760dc..fd1cba7b14 100644
--- a/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-option/expected.js
+++ b/packages/babel-generator/test/fixtures/edgecase/return-with-retainlines-option/expected.js
@@ -1,10 +1,11 @@
 function foo(l) {
   return (
-    l);}
-
+    l);
 
+}
 
 function foo() {
   return (
     1 && 2 || 
-    3);}
+    3);
+}