Fix control flow flattening + optional chaining compatibility (#1334)

src/custom-nodes/control-flow-flattening-nodes/CallExpressionFunctionNode.ts

@@ -25,6 +25,13 @@ export class CallExpressionFunctionNode extends AbstractCustomNode {
     @initializable()
     private expressionArguments!: (ESTree.Expression | ESTree.SpreadElement)[];
 
+    /**
+     * @type {boolean}
+     * @private
+     */
+    @initializable()
+    private isChainExpressionParent!: boolean;
+
     /**
      * @param {TIdentifierNamesGeneratorFactory} identifierNamesGeneratorFactory
      * @param {ICustomCodeHelperFormatter} customCodeHelperFormatter
@@ -48,9 +55,11 @@ export class CallExpressionFunctionNode extends AbstractCustomNode {
 
     /**
      * @param {(Expression | SpreadElement)[]} expressionArguments
+     * @param {boolean} isChainExpressionParent
      */
-    public initialize (expressionArguments: (ESTree.Expression | ESTree.SpreadElement)[]): void {
+    public initialize (expressionArguments: (ESTree.Expression | ESTree.SpreadElement)[], isChainExpressionParent: boolean): void {
         this.expressionArguments = expressionArguments;
+        this.isChainExpressionParent = isChainExpressionParent;
     }
 
     /**
@@ -82,6 +91,12 @@ export class CallExpressionFunctionNode extends AbstractCustomNode {
             }
         }
 
+        const callExpression = NodeFactory.callExpressionNode(
+          calleeIdentifier,
+          callArguments,
+          this.isChainExpressionParent
+        );
+        
         const structure: TStatement = NodeFactory.expressionStatementNode(
             NodeFactory.functionExpressionNode(
                 [
@@ -90,10 +105,9 @@ export class CallExpressionFunctionNode extends AbstractCustomNode {
                 ],
                 NodeFactory.blockStatementNode([
                     NodeFactory.returnStatementNode(
-                        NodeFactory.callExpressionNode(
-                            calleeIdentifier,
-                            callArguments
-                        )
+                        this.isChainExpressionParent
+                          ? NodeFactory.chainExpressionNode(<ESTree.ChainElement>callExpression)
+                          : callExpression
                     )
                 ])
             )

src/node-transformers/control-flow-transformers/control-flow-replacers/CallExpressionControlFlowReplacer.ts

@@ -66,12 +66,14 @@ export class CallExpressionControlFlowReplacer extends AbstractControlFlowReplac
             return callExpressionNode;
         }
 
+        const isChainExpressionParent = NodeGuards.isChainExpressionNode(parentNode);
+
         const replacerId: number = callExpressionNode.arguments.length;
         const callExpressionFunctionCustomNode: ICustomNode<TInitialData<CallExpressionFunctionNode>> =
             this.controlFlowCustomNodeFactory(ControlFlowCustomNode.CallExpressionFunctionNode);
         const expressionArguments: (ESTree.Expression | ESTree.SpreadElement)[] = callExpressionNode.arguments;
 
-        callExpressionFunctionCustomNode.initialize(expressionArguments);
+        callExpressionFunctionCustomNode.initialize(expressionArguments, isChainExpressionParent);
 
         const storageKey: string = this.insertCustomNodeToControlFlowStorage(
             callExpressionFunctionCustomNode,
@@ -84,7 +86,7 @@ export class CallExpressionControlFlowReplacer extends AbstractControlFlowReplac
             controlFlowStorage.getStorageId(),
             storageKey,
             callee,
-            expressionArguments
+            expressionArguments,
         );
     }
 

src/node/NodeFactory.ts

@@ -118,6 +118,20 @@ export class NodeFactory {
         };
     }
 
+    /**
+     * @param {ChainElement} expression
+     * @return {ChainExpression}
+     */
+    public static chainExpressionNode (
+        expression: ESTree.ChainElement,
+    ): ESTree.ChainExpression {
+        return {
+            type: NodeType.ChainExpression,
+            expression,
+            metadata: { ignoredNode: false }
+        };
+    }
+
     /**
      * @param {ESTree.Expression} test
      * @param {ESTree.Expression} consequent

src/node/NodeGuards.ts

@@ -91,6 +91,14 @@ export class NodeGuards {
         return node.type === NodeType.CallExpression;
     }
 
+  /**
+   * @param {Node} node
+   * @returns {boolean}
+   */
+  public static isChainExpressionNode (node: ESTree.Node): node is ESTree.ChainExpression {
+    return node.type === NodeType.ChainExpression;
+  }
+
     /**
      * @param {Node} node
      * @returns {boolean}