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

95.45% Statements 63/66
83.33% Branches 15/18
100% Functions 2/2
95.16% Lines 59/62
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 2021x 1x   1x   1x                       1x 1x   1x 1x 1x 1x 1x 1x     1x       1x             1x         1x         202x         202x                                                         202x   202x 202x 202x               222x   222x 215x   7x     7x       7x       7x                     21x 21x   21x 20x     1x     5x             222x             222x       222x 222x   222x 216x   6x     6x   6x   6x     222x   5154x 327x     4827x     4827x 4346x     481x               222x 9x     213x   213x   213x   213x 213x      
import { injectable, inject } from 'inversify';
import { ServiceIdentifiers } from '../../container/ServiceIdentifiers';
 
import * as estraverse from 'estraverse';
import * as ESTree from 'estree';
import * as _ from 'lodash';
 
import { TControlFlowReplacerFactory } from '../../types/container/TControlFlowReplacerFactory';
import { TControlFlowStorageFactory } from '../../types/container/TControlFlowStorageFactory';
import { TCustomNodeFactory } from '../../types/container/TCustomNodeFactory';
import { TNodeWithBlockStatement } from '../../types/node/TNodeWithBlockStatement';
import { TStatement } from '../../types/node/TStatement';
 
import { ICustomNode } from '../../interfaces/custom-nodes/ICustomNode';
import { IOptions } from '../../interfaces/options/IOptions';
import { IStorage } from '../../interfaces/storages/IStorage';
 
import { CustomNodes } from '../../enums/container/CustomNodes';
import { NodeType } from '../../enums/NodeType';
 
import { AbstractNodeTransformer } from '../AbstractNodeTransformer';
import { Node } from '../../node/Node';
import { NodeAppender } from '../../node/NodeAppender';
import { NodeControlFlowReplacers } from '../../enums/container/NodeControlFlowReplacers';
import { NodeUtils } from '../../node/NodeUtils';
import { RandomGeneratorUtils } from '../../utils/RandomGeneratorUtils';
 
@injectable()
export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
    /**
     * @type {Map <string, NodeControlFlowReplacers>}
     */
    private static readonly controlFlowReplacersMap: Map <string, NodeControlFlowReplacers> = new Map([
        [NodeType.BinaryExpression, NodeControlFlowReplacers.BinaryExpressionControlFlowReplacer]
    ]);
 
    /**
     * @type {number}
     */
    private static readonly hostNodeSearchMinDepth: number = 2;
 
    /**
     * @type {number}
     */
    private static readonly hostNodeSearchMaxDepth: number = 10;
 
    /**
     * @type {Map<ESTree.Node, IStorage<ICustomNode>>}
     */
    private controlFlowData: Map <ESTree.Node, IStorage<ICustomNode>> = new Map();
 
    /**
     * @type {TStatement[][]}
     */
    private readonly controlFlowNodesList: TStatement[][] = [];
 
    /**
     * @type {TControlFlowReplacerFactory}
     */
    private readonly controlFlowReplacerFactory: TControlFlowReplacerFactory;
 
    /**
     * @type {TControlFlowStorageFactory}
     */
    private readonly controlFlowStorageFactory: TControlFlowStorageFactory;
 
    /**
     * @type {TCustomNodeFactory}
     */
    private readonly customNodeFactory: TCustomNodeFactory;
 
    /**
     * @param controlFlowStorageFactory
     * @param controlFlowReplacerFactory
     * @param customNodeFactory
     * @param options
     */
    constructor (
        @inject(ServiceIdentifiers['Factory<IStorage<ICustomNode>>']) controlFlowStorageFactory: TControlFlowStorageFactory,
        @inject(ServiceIdentifiers['Factory<IControlFlowReplacer>']) controlFlowReplacerFactory: TControlFlowReplacerFactory,
        @inject(ServiceIdentifiers['Factory<ICustomNode>']) customNodeFactory: TCustomNodeFactory,
        @inject(ServiceIdentifiers.IOptions) options: IOptions
    ) {
        super(options);
 
        this.controlFlowStorageFactory = controlFlowStorageFactory;
        this.controlFlowReplacerFactory = controlFlowReplacerFactory;
        this.customNodeFactory = customNodeFactory;
    }
 
    /**
     * @param functionNode
     * @returns {TNodeWithBlockStatement}
     */
    private static getHostNode (functionNode: ESTree.FunctionDeclaration | ESTree.FunctionExpression): TNodeWithBlockStatement {
        const blockScopesOfNode: TNodeWithBlockStatement[] = NodeUtils.getBlockScopesOfNode(functionNode);
 
        if (blockScopesOfNode.length === 1) {
            return functionNode.body;
        } else {
            blockScopesOfNode.pop();
        }
 
        Iif (blockScopesOfNode.length > FunctionControlFlowTransformer.hostNodeSearchMinDepth) {
            blockScopesOfNode.splice(0, FunctionControlFlowTransformer.hostNodeSearchMinDepth);
        }
 
        Iif (blockScopesOfNode.length > FunctionControlFlowTransformer.hostNodeSearchMaxDepth) {
            blockScopesOfNode.length = FunctionControlFlowTransformer.hostNodeSearchMaxDepth;
        }
 
        return RandomGeneratorUtils.getRandomGenerator().pickone(blockScopesOfNode);
    }
 
    /**
     * @param hostNodeBody
     * @param controlFlowNodesList
     */
    private static removeOldControlFlowNodeFromHostNodeBody (
        hostNodeBody: TStatement[],
        controlFlowNodesList: TStatement[][]
    ): TStatement[] {
        for (let controlFlowNode of controlFlowNodesList) {
            const firstIndexOfNode: number = hostNodeBody.indexOf(controlFlowNode[0]);
 
            if (firstIndexOfNode === -1) {
                continue;
            }
 
            return _.difference(hostNodeBody, controlFlowNode);
        }
 
        return hostNodeBody;
    }
 
    /**
     * @param functionNode
     */
    public transformNode (functionNode: ESTree.Function): void {
        this.changeFunctionBodyControlFlow(functionNode);
    }
 
    /**
     * @param functionNode
     */
    private changeFunctionBodyControlFlow (functionNode: ESTree.Function): void {
        Iif (Node.isArrowFunctionExpressionNode(functionNode)) {
            return;
        }
 
        const controlFlowStorage: IStorage <ICustomNode> = this.controlFlowStorageFactory();
        const hostNode: TNodeWithBlockStatement = FunctionControlFlowTransformer.getHostNode(functionNode);
 
        if (!this.controlFlowData.has(hostNode)) {
            this.controlFlowData.set(hostNode, controlFlowStorage);
        } else {
            hostNode.body = <ESTree.Statement[]>FunctionControlFlowTransformer
                .removeOldControlFlowNodeFromHostNodeBody(hostNode.body, this.controlFlowNodesList);
 
            const hostControlFlowStorage: IStorage<ICustomNode> = <IStorage<ICustomNode>>this.controlFlowData.get(hostNode);
 
            controlFlowStorage.mergeWith(hostControlFlowStorage, true);
 
            this.controlFlowData.set(hostNode, controlFlowStorage);
        }
 
        estraverse.replace(functionNode.body, {
            enter: (node: ESTree.Node, parentNode: ESTree.Node): any => {
                if (RandomGeneratorUtils.getRandomFloat(0, 1) > this.options.controlFlowFlatteningThreshold) {
                    return;
                }
 
                const controlFlowReplacerName: NodeControlFlowReplacers | undefined = FunctionControlFlowTransformer
                    .controlFlowReplacersMap.get(node.type);
 
                if (controlFlowReplacerName === undefined) {
                    return;
                }
 
                return {
                    ...this.controlFlowReplacerFactory(controlFlowReplacerName)
                        .replace(node, parentNode, controlFlowStorage),
                    parentNode
                };
            }
        });
 
        if (!controlFlowStorage.getLength()) {
            return;
        }
 
        const controlFlowStorageCustomNode: ICustomNode = this.customNodeFactory(CustomNodes.ControlFlowStorageNode);
 
        controlFlowStorageCustomNode.initialize(controlFlowStorage);
 
        const controlFlowStorageNode: TStatement[] = controlFlowStorageCustomNode.getNode();
 
        this.controlFlowNodesList.push(controlFlowStorageNode);
        NodeAppender.prependNode(hostNode, controlFlowStorageNode);
    }
}