From 8a3124e8c8f03fc64716cd8a959986f08c51ee4c Mon Sep 17 00:00:00 2001
From: Kashyap <kashyap.kmbc@gmail.com>
Date: Fri, 23 Feb 2018 11:43:48 +0530
Subject: [PATCH] Fix bug where profileEnd() gets added even if no comment was
 found

This is due to a bug in logic. If we receive a `BlockStatement`, we were
adding `console.profileEnd()` even there was no `// profile` comment by
checking `!state.gotReturn` assuming this condition meant we reached the
end of the block but didn't get a return (`if` condition, for instance).
Since `gotReturn` is by default `false`, we were adding the closing
comment always. This introduces a perf issue in Safari because that
browser will try to run the profile by seeing this comment.
---
 __snapshots__/test.js.snap | 30 +++++++++++++++++++++++++++++-
 index.js                   |  5 ++++-
 package.json               | 13 ++++++++++---
 test.js                    | 28 ++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/__snapshots__/test.js.snap b/__snapshots__/test.js.snap
index cf195bf..158a203 100644
--- a/__snapshots__/test.js.snap
+++ b/__snapshots__/test.js.snap
@@ -201,7 +201,6 @@ class A {
 				newState: {}
 			});
 			callSomeMethod();
-			console.profileEnd();
 		};
 
 		console.profileEnd();
@@ -212,6 +211,35 @@ class A {
 }"
 `;
 
+exports[`usage inside if conditions 1`] = `
+"
+const b = () => {
+	if (true) {
+		console.profile('b');
+		console.profileEnd();
+
+		// profile
+
+		return 12;
+	}
+
+	switch (true) {
+		case 'case1':
+			return 'what';
+		case 'case 2':
+			{
+				return 'where';
+			}
+		case 'case 3':
+			return 'who';
+		default:
+			return 'default';
+	}
+
+	return 42;
+};"
+`;
+
 exports[`wraps a function in profile block 1`] = `
 "
 const b = () => {
diff --git a/index.js b/index.js
index a95440c..13590c4 100644
--- a/index.js
+++ b/index.js
@@ -162,6 +162,7 @@ const ConsoleProfileVisitor = function (babel) {
 			BlockStatement: function (path, args) {
 				let state = {
 					gotProfileComment: false,
+					addedProfileComment: false,
 					gotReturn: false,
 					blockDepth: 0
 				};
@@ -173,11 +174,13 @@ const ConsoleProfileVisitor = function (babel) {
 						'body',
 						generateProfileStart(getName(path))
 					);
+
+					state.addedProfileComment = true;
 				}
 
 				path.traverse(ReturnVisitor, { state });
 
-				if (!state.gotReturn) {
+				if (state.addedProfileComment && !state.gotReturn) {
 					path.pushContainer('body', generateProfileEnd());
 				}
 			}
diff --git a/package.json b/package.json
index bf050dd..0dee7cd 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "babel-plugin-console-perf",
-  "version": "0.0.2",
+  "version": "0.0.3",
   "description": "A plugin that makes it easy to profile functions using console.profile and console.time functions",
   "main": "index.js",
   "author": "Kashyap Kondamudi",
@@ -16,6 +16,13 @@
     "babel-plugin-transform-class-properties": "^6.24.1",
     "jest": "^21.2.1"
   },
-  "keywords": [ "babel-plugin", "console.profile", "console.time", "performance" ],
-  "files": [ "index.js" ]
+  "keywords": [
+    "babel-plugin",
+    "console.profile",
+    "console.time",
+    "performance"
+  ],
+  "files": [
+    "index.js"
+  ]
 }
diff --git a/test.js b/test.js
index 6e07e25..21ce626 100644
--- a/test.js
+++ b/test.js
@@ -187,6 +187,34 @@ it('maintains indentation when for if blocks', () => {
 	expect(code).toMatchSnapshot();
 });
 
+fit('usage inside if conditions', () => {
+	const example = `
+	const b = () => {
+		if (true) {
+			// profile
+		
+			return 12;
+		}
+
+		switch (true) {
+			case 'case1':
+				return 'what';
+			case 'case 2': {
+				return 'where';
+			}
+			case 'case 3':
+				return 'who';
+			default:
+				return 'default'
+		}
+
+		return 42;
+	};
+	`;
+
+	const code = babel.transform(example, { plugins: [ plugin ] }).code;
+	expect(code).toMatchSnapshot();
+});
 
 describe('usage inside classes', () => {
 	it('works inside function properties', () => {