From e743d070ea2d58f9ac98aa917ff2be223dac598a Mon Sep 17 00:00:00 2001
From: Oleksandr T <oleksandr.tarasiuk@outlook.com>
Date: Sat, 23 Sep 2023 02:14:06 +0300
Subject: [PATCH] fix(55404): Remove braces from arrow function" generates
 invalid code in JavaScript (#55429)

Co-authored-by: Nathan Shively-Sanders <293473+sandersn@users.noreply.github.com>
---
 .../addOrRemoveBracesToArrowFunction.ts       |  5 ++++-
 ...actorAddOrRemoveBracesToArrowFunction36.ts | 20 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction36.ts

diff --git a/src/services/refactors/addOrRemoveBracesToArrowFunction.ts b/src/services/refactors/addOrRemoveBracesToArrowFunction.ts
index 0d7489554cf..6cd5bb9f05d 100644
--- a/src/services/refactors/addOrRemoveBracesToArrowFunction.ts
+++ b/src/services/refactors/addOrRemoveBracesToArrowFunction.ts
@@ -12,11 +12,13 @@ import {
     factory,
     first,
     getContainingFunction,
+    getLeftmostExpression,
     getLocaleSpecificMessage,
     getTokenAtPosition,
     isArrowFunction,
     isBlock,
     isExpression,
+    isObjectLiteralExpression,
     isReturnStatement,
     needsParentheses,
     rangeContainsRange,
@@ -147,7 +149,8 @@ function getConvertibleArrowFunctionAtPosition(file: SourceFile, startPosition:
     else if (refactorKindBeginsWith(removeBracesAction.kind, kind) && isBlock(func.body) && func.body.statements.length === 1) {
         const firstStatement = first(func.body.statements);
         if (isReturnStatement(firstStatement)) {
-            return { func, addBraces: false, expression: firstStatement.expression, returnStatement: firstStatement };
+            const expression = firstStatement.expression && isObjectLiteralExpression(getLeftmostExpression(firstStatement.expression, /*stopAtCallExpressions*/ false)) ? factory.createParenthesizedExpression(firstStatement.expression) : firstStatement.expression;
+            return { func, addBraces: false, expression, returnStatement: firstStatement };
         }
     }
     return undefined;
diff --git a/tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction36.ts b/tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction36.ts
new file mode 100644
index 00000000000..42195b9f770
--- /dev/null
+++ b/tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction36.ts
@@ -0,0 +1,20 @@
+/// <reference path='fourslash.ts' />
+
+////const f = /*a*/(param) => {
+////    return {
+////        "a": 1,
+////        "b": 2,
+////    }[param];
+////}/*b*/;
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent:
+`const f = (param) => ({
+    "a": 1,
+    "b": 2,
+}[param]);`,
+});
\ No newline at end of file