Instantiate temporary variable if not available

__snapshots__/test.js.snap

@@ -7,7 +7,8 @@ const a = () => {
 
 	// profile
 	if (true) {
-		_babel_temp_alias = Math.random(12);
+		let _babel_temp_alias = Math.random(12);
+
 		console.profileEnd();
 
 		return _babel_temp_alias;
@@ -158,20 +159,46 @@ const b = () => {
 
 	const a = () => {
 		console.profile(\\"a\\");
-		_babel_temp_alias = Math.sqrt(1000000000000);
+
+		let _babel_temp_alias = Math.sqrt(1000000000000);
+
 		console.profileEnd();
 
 		// profile
 		return _babel_temp_alias;
 	};
 
 	// profile
-	_babel_temp_alias = a();
+
+	let _babel_temp_alias = a();
+
 	console.profileEnd();
 	return _babel_temp_alias;
 };"
 `;
 
+exports[`usage inside classes works inside function properties 1`] = `
+"
+class A {
+	changeVisualState(id) {
+		console.profile(\\"changeVisualState\\");
+
+		let _babel_temp_alias = item => {
+			this.setState({
+				newState: {}
+			});
+			callSomeMethod();
+			console.profileEnd();
+		};
+
+		console.profileEnd();
+
+		// profile
+		return _babel_temp_alias;
+	}
+}"
+`;
+
 exports[`wraps a function in profile block 1`] = `
 "
 const b = () => {
@@ -184,7 +211,9 @@ const b = () => {
 
 const a = () => {
 	console.profile(\\"a\\");
-	_babel_temp_alias = b();
+
+	let _babel_temp_alias = b();
+
 	console.profileEnd();
 
 	// profile

index.js

@@ -48,14 +48,27 @@ const generateProfileEnd = () =>
 
 const generateIdentifier = () => t.identifier('_babel_temp_alias');
 
-const assignArgument = (identifier, argument) =>
-	t.expressionStatement(
-		t.assignmentExpression(
-			"=",
-			t.identifier('_babel_temp_alias'),
-			argument
-		)
-	);
+const assignArgument = (identifier, argument, opts) => {
+	if (opts.blockDepth === 1) {
+		return t.variableDeclaration(
+			'let',
+			[
+				t.variableDeclarator(
+					t.identifier('_babel_temp_alias'),
+					argument
+				)
+			]
+		);
+	} else {
+		return t.expressionStatement(
+			t.assignmentExpression(
+				"=",
+				t.identifier('_babel_temp_alias'),
+				argument
+			)
+		);
+	}
+}
 
 
 const SkipVisitor = function (path, args) {
@@ -68,6 +81,7 @@ const ReturnVisitor = {
 		const argument = path.node.argument;
 
 		args.state.gotReturn = true;
+		args.state.blockDepth++;
 
 		if (args.state.gotProfileComment) {
 			if (!argument || argument.type === 'Literal' || argument.type === 'NumericLiteral' || argument.type === 'StringLiteral' || argument.type === 'Identifier') {
@@ -81,7 +95,9 @@ const ReturnVisitor = {
 				// set the argument of return to the variable identifier
 
 				const identifier = generateIdentifier();
-				const assignment = assignArgument(identifier, argument);
+				const assignment = assignArgument(identifier, argument, {
+					blockDepth: args.state.blockDepth
+				});
 				path.insertBefore(assignment);
 				path.insertBefore(generateProfileEnd());
 				path.node.argument = identifier;
@@ -139,7 +155,8 @@ const ConsoleProfileVisitor = function (babel) {
 			BlockStatement: function (path, args) {
 				let state = {
 					gotProfileComment: false,
-					gotReturn: false
+					gotReturn: false,
+					blockDepth: 0
 				};
 
 				BlockVisitor(path, { state });

test.js

@@ -170,3 +170,25 @@ it('maintains indentation when for if blocks', () => {
 	expect(code).toMatchSnapshot();
 });
 
+
+describe('usage inside classes', () => {
+	it('works inside function properties', () => {
+		const example = `
+		class A {
+			changeVisualState (id) {
+				// profile
+				return item => {
+					this.setState({
+						newState: {}
+					});
+					callSomeMethod();
+				};
+			}
+		}
+		`;
+
+		const code = babel.transform(example, { plugins: [ plugin ] }).code;
+		expect(code).toMatchSnapshot();
+	});
+});
+