diff --git a/src/node-transformers/control-flow-transformers/FunctionControlFlowTransformer.ts b/src/node-transformers/control-flow-transformers/FunctionControlFlowTransformer.ts index d3acec00..cb029901 100644 --- a/src/node-transformers/control-flow-transformers/FunctionControlFlowTransformer.ts +++ b/src/node-transformers/control-flow-transformers/FunctionControlFlowTransformer.ts @@ -62,9 +62,9 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer { private readonly visitedFunctionNodes: Set = new Set(); /** - * @type {Set} + * @type {WeakMap} */ - private readonly hostNodesWithControlFlowNode: Set = new Set(); + private readonly hostNodesWithControlFlowNode: WeakMap = 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 = 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> = + 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} diff --git a/src/node/NodeAppender.ts b/src/node/NodeAppender.ts index 9fb46434..e16fd3f3 100644 --- a/src/node/NodeAppender.ts +++ b/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 diff --git a/test/unit-tests/node/node-appender/NodeAppender.spec.ts b/test/unit-tests/node/node-appender/NodeAppender.spec.ts index 5fa1d240..c0891eeb 100644 --- a/test/unit-tests/node/node-appender/NodeAppender.spec.ts +++ b/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, 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, astTree.body[5]); + }); + + it('should keep `BlockStatement` as is', () => { + assert.deepEqual(astTree, expectedAstTree); + }); + }); + }); }); diff --git a/test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index-expected.js b/test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index-expected.js new file mode 100644 index 00000000..fd46b072 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index.js b/test/unit-tests/node/node-appender/fixtures/remove-node/invalid-index.js new file mode 100644 index 00000000..fd46b072 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/test/unit-tests/node/node-appender/fixtures/remove-node/valid-index-expected.js b/test/unit-tests/node/node-appender/fixtures/remove-node/valid-index-expected.js new file mode 100644 index 00000000..d4a41646 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/test/unit-tests/node/node-appender/fixtures/remove-node/valid-index.js b/test/unit-tests/node/node-appender/fixtures/remove-node/valid-index.js new file mode 100644 index 00000000..fd46b072 --- /dev/null +++ b/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; \ No newline at end of file