Refactored control flow storage logic related to removing existing control flow storage from the AST tree

src/node-transformers/control-flow-transformers/FunctionControlFlowTransformer.ts

@@ -62,9 +62,9 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     private readonly visitedFunctionNodes: Set<ESTree.Function> = new Set();
 
     /**
-     * @type {Set<TNodeWithStatements>}
+     * @type {WeakMap<TNodeWithStatements, VariableDeclaration>}
      */
-    private readonly hostNodesWithControlFlowNode: Set<TNodeWithStatements> = new Set();
+    private readonly hostNodesWithControlFlowNode: WeakMap<TNodeWithStatements, ESTree.VariableDeclaration> = new WeakMap();
 
     /**
      * @type {TControlFlowReplacerFactory}
@@ -157,8 +157,11 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
             this.controlFlowCustomNodeFactory(ControlFlowCustomNode.ControlFlowStorageNode);
 
         controlFlowStorageCustomNode.initialize(controlFlowStorage);
-        NodeAppender.prepend(hostNode, controlFlowStorageCustomNode.getNode());
-        this.hostNodesWithControlFlowNode.add(hostNode);
+
+        const controlFlowStorageNode: ESTree.VariableDeclaration = this.getControlFlowStorageNode(controlFlowStorage);
+
+        NodeAppender.prepend(hostNode, [controlFlowStorageNode]);
+        this.hostNodesWithControlFlowNode.set(hostNode, controlFlowStorageNode);
 
         NodeUtils.parentizeAst(functionNode);
 
@@ -173,12 +176,11 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
         const controlFlowStorage: TControlFlowStorage = this.controlFlowStorageFactory();
 
         if (this.controlFlowData.has(hostNode)) {
-            if (this.hostNodesWithControlFlowNode.has(hostNode)) {
-                if (NodeGuards.isSwitchCaseNode(hostNode)) {
-                    hostNode.consequent.shift();
-                } else {
-                    hostNode.body.shift();
-                }
+            const existingControlFlowStorageNode: ESTree.VariableDeclaration | null =
+                this.hostNodesWithControlFlowNode.get(hostNode) ?? null;
+
+            if (existingControlFlowStorageNode) {
+                NodeAppender.remove(hostNode, existingControlFlowStorageNode);
             }
 
             const hostControlFlowStorage: TControlFlowStorage = <TControlFlowStorage>this.controlFlowData.get(hostNode);
@@ -189,6 +191,25 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
         return controlFlowStorage;
     }
 
+    /**
+     * @param {TControlFlowStorage} controlFlowStorage
+     * @returns {VariableDeclaration}
+     */
+    private getControlFlowStorageNode (controlFlowStorage: TControlFlowStorage): ESTree.VariableDeclaration {
+        const controlFlowStorageCustomNode: ICustomNode<TInitialData<ControlFlowStorageNode>> =
+            this.controlFlowCustomNodeFactory(ControlFlowCustomNode.ControlFlowStorageNode);
+
+        controlFlowStorageCustomNode.initialize(controlFlowStorage);
+
+        const controlFlowStorageNode: ESTree.Node = controlFlowStorageCustomNode.getNode()[0];
+
+        if (!NodeGuards.isVariableDeclarationNode(controlFlowStorageNode)) {
+            throw new Error('`controlFlowStorageNode` should contain `VariableDeclaration` node with control flow storage object');
+        }
+
+        return controlFlowStorageNode;
+    }
+
     /**
      * @param {BlockStatement} functionNodeBody
      * @returns {TNodeWithStatements}

src/node/NodeAppender.ts

@@ -160,6 +160,24 @@ export class NodeAppender {
         ]);
     }
 
+    /**
+     * @param {TNodeWithStatements} nodeWithStatements
+     * @param {Statement} statement
+     */
+    public static remove (nodeWithStatements: TNodeWithStatements, statement: ESTree.Statement): void {
+        const scopeStatements: TStatement[] = NodeAppender.getScopeStatements(nodeWithStatements);
+        const indexInScopeStatement: number = scopeStatements.indexOf(statement);
+
+        if (indexInScopeStatement === -1) {
+            return;
+        }
+
+        const updatedStatements: TStatement[] = [...scopeStatements];
+        updatedStatements.splice(indexInScopeStatement, 1);
+
+        NodeAppender.setScopeStatements(nodeWithStatements, updatedStatements);
+    }
+
     /**
      * @param {TNodeWithStatements} nodeWithStatements
      * @param {TStatement[]} statements

test/unit-tests/node/node-appender/NodeAppender.spec.ts

@@ -13,12 +13,12 @@ import { ICallsGraphAnalyzer } from '../../../../src/interfaces/analyzers/calls-
 import { ICallsGraphData } from '../../../../src/interfaces/analyzers/calls-graph-analyzer/ICallsGraphData';
 
 import { readFileAsString } from '../../../helpers/readFileAsString';
+import { removeRangesFromStructure } from '../../../helpers/removeRangesFromStructure';
 
 import { InversifyContainerFacade } from '../../../../src/container/InversifyContainerFacade';
 import { NodeAppender } from '../../../../src/node/NodeAppender';
 import { NodeFactory } from '../../../../src/node/NodeFactory';
 import { NodeUtils } from '../../../../src/node/NodeUtils';
-import { removeRangesFromStructure } from '../../../helpers/removeRangesFromStructure';
 
 /**
  * @param fixturePath
@@ -254,4 +254,44 @@ describe('NodeAppender', () => {
             assert.deepEqual(astTree, expectedAstTree);
         });
     });
+
+    describe('remove', () => {
+        describe('Variant #1: valid index', () => {
+            let astTree: ESTree.Program,
+                expectedAstTree: ESTree.Program;
+
+            before(() => {
+                astTree = convertCodeToAst('/fixtures/remove-node/valid-index.js');
+                expectedAstTree = convertCodeToAst('/fixtures/remove-node/valid-index-expected.js');
+
+                astTree = NodeUtils.parentizeAst(astTree);
+                expectedAstTree = NodeUtils.parentizeAst(expectedAstTree);
+
+                NodeAppender.remove(astTree, <ESTree.Statement>astTree.body[2]);
+            });
+
+            it('should remove given node from a `BlockStatement` node body', () => {
+                assert.deepEqual(astTree, expectedAstTree);
+            });
+        });
+
+        describe('Variant #2: invalid index', () => {
+            let astTree: ESTree.Program,
+                expectedAstTree: ESTree.Program;
+
+            before(() => {
+                astTree = convertCodeToAst('/fixtures/remove-node/invalid-index.js');
+                expectedAstTree = convertCodeToAst('/fixtures/remove-node/invalid-index-expected.js');
+
+                astTree = NodeUtils.parentizeAst(astTree);
+                expectedAstTree = NodeUtils.parentizeAst(expectedAstTree);
+
+                NodeAppender.remove(astTree, <ESTree.Statement>astTree.body[5]);
+            });
+
+            it('should keep `BlockStatement` as is', () => {
+                assert.deepEqual(astTree, expectedAstTree);
+            });
+        });
+    });
 });

test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index-expected.js

@@ -0,0 +1,4 @@
+var foo = 1;
+var bar = 2;
+var baz = 3;
+var bark = 4;

test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index.js

@@ -0,0 +1,4 @@
+var foo = 1;
+var bar = 2;
+var baz = 3;
+var bark = 4;

test/unit-tests/node/node-appender/fixtures/remove-node/valid-index-expected.js

@@ -0,0 +1,3 @@
+var foo = 1;
+var bar = 2;
+var bark = 4;

test/unit-tests/node/node-appender/fixtures/remove-node/valid-index.js

@@ -0,0 +1,4 @@
+var foo = 1;
+var bar = 2;
+var baz = 3;
+var bark = 4;