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

100% Statements 68/68
96.55% Branches 28/29
100% Functions 3/3
100% Lines 65/65
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 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 2311x 1x   1x                         1x 1x   1x 1x 1x 1x 1x 1x     1x       1x                 1x         1x         4041x         4041x         4041x                                                         4041x   4041x 4041x 4041x               8144x   8144x 5053x   3091x     3091x 3091x     3091x 40x     3091x             4041x   115113x         8145x                       8145x   8145x 1x     8144x 8144x   8144x 8144x   8144x 1028x     7116x   7116x 7116x 7116x   7116x               8144x   8144x 2536x 2526x     2536x   2536x     8144x               142437x                       8144x   142437x 3093x     139344x 125076x     14268x 2788x     11480x     11480x                
import { injectable, inject } from 'inversify';
import { ServiceIdentifiers } from '../../container/ServiceIdentifiers';
 
import * as estraverse from 'estraverse';
import * as ESTree from 'estree';
 
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 { ICustomNode } from '../../interfaces/custom-nodes/ICustomNode';
import { IOptions } from '../../interfaces/options/IOptions';
import { IStorage } from '../../interfaces/storages/IStorage';
import { IVisitor } from '../../interfaces/IVisitor';
 
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 { ControlFlowReplacers } from '../../enums/container/ControlFlowReplacers';
import { NodeUtils } from '../../node/NodeUtils';
import { RandomGeneratorUtils } from '../../utils/RandomGeneratorUtils';
 
@injectable()
export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
    /**
     * @type {Map <string, ControlFlowReplacers>}
     */
    private static readonly controlFlowReplacersMap: Map <string, ControlFlowReplacers> = new Map([
        [NodeType.BinaryExpression, ControlFlowReplacers.BinaryExpressionControlFlowReplacer],
        [NodeType.CallExpression, ControlFlowReplacers.CallExpressionControlFlowReplacer],
        [NodeType.LogicalExpression, ControlFlowReplacers.LogicalExpressionControlFlowReplacer]
    ]);
 
    /**
     * @type {number}
     */
    private static readonly hostNodeSearchMinDepth: number = 0;
 
    /**
     * @type {number}
     */
    private static readonly hostNodeSearchMaxDepth: number = 2;
 
    /**
     * @type {Map<ESTree.Node, IStorage<ICustomNode>>}
     */
    private readonly controlFlowData: Map <ESTree.Node, IStorage<ICustomNode>> = new Map();
 
    /**
     * @type {Set<ESTree.Function>}
     */
    private readonly visitedFunctionNodes: Set<ESTree.Function> = new Set();
 
    /**
     * @type {Set<TNodeWithBlockStatement>}
     */
    private readonly hostNodesWithControlFlowNode: Set<TNodeWithBlockStatement> = new Set();
 
    /**
     * @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__TControlFlowStorage) 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 functionNodeBody
     * @returns {TNodeWithBlockStatement}
     */
    private static getHostNode (functionNodeBody: ESTree.BlockStatement): TNodeWithBlockStatement {
        const blockScopesOfNode: TNodeWithBlockStatement[] = NodeUtils.getBlockScopesOfNode(functionNodeBody);
 
        if (blockScopesOfNode.length === 1) {
            return functionNodeBody;
        } else {
            blockScopesOfNode.pop();
        }
 
        Eif (blockScopesOfNode.length > FunctionControlFlowTransformer.hostNodeSearchMinDepth) {
            blockScopesOfNode.splice(0, FunctionControlFlowTransformer.hostNodeSearchMinDepth);
        }
 
        if (blockScopesOfNode.length > FunctionControlFlowTransformer.hostNodeSearchMaxDepth) {
            blockScopesOfNode.length = FunctionControlFlowTransformer.hostNodeSearchMaxDepth;
        }
 
        return RandomGeneratorUtils.getRandomGenerator().pickone(blockScopesOfNode);
    }
 
    /**
     * @return {IVisitor}
     */
    public getVisitor (): IVisitor {
        return {
            leave: (node: ESTree.Node, parentNode: ESTree.Node) => {
                if (
                    Node.isFunctionDeclarationNode(node) ||
                    Node.isFunctionExpressionNode(node) ||
                    Node.isArrowFunctionExpressionNode(node)
                ) {
                    return this.transformNode(node, parentNode);
                }
            }
        };
    }
 
    /**
     * @param functionNode
     * @param parentNode
     * @returns {ESTree.Function}
     */
    public transformNode (functionNode: ESTree.Function, parentNode: ESTree.Node): ESTree.Function {
        this.visitedFunctionNodes.add(functionNode);
 
        if (!Node.isBlockStatementNode(functionNode.body)) {
            return functionNode;
        }
 
        const hostNode: TNodeWithBlockStatement = FunctionControlFlowTransformer.getHostNode(functionNode.body);
        const controlFlowStorage: IStorage<ICustomNode> = this.getControlFlowStorage(hostNode);
 
        this.controlFlowData.set(hostNode, controlFlowStorage);
        this.transformFunctionBody(functionNode.body, controlFlowStorage);
 
        if (!controlFlowStorage.getLength()) {
            return functionNode;
        }
 
        const controlFlowStorageCustomNode: ICustomNode = this.customNodeFactory(CustomNodes.ControlFlowStorageNode);
 
        controlFlowStorageCustomNode.initialize(controlFlowStorage);
        NodeAppender.prependNode(hostNode, controlFlowStorageCustomNode.getNode());
        this.hostNodesWithControlFlowNode.add(hostNode);
 
        return functionNode;
    }
 
    /**
     * @param hostNode
     * @return {IStorage<ICustomNode>}
     */
    private getControlFlowStorage (hostNode: TNodeWithBlockStatement): IStorage<ICustomNode> {
        const controlFlowStorage: IStorage <ICustomNode> = this.controlFlowStorageFactory();
 
        if (this.controlFlowData.has(hostNode)) {
            if (this.hostNodesWithControlFlowNode.has(hostNode)) {
                hostNode.body.shift();
            }
 
            const hostControlFlowStorage: IStorage<ICustomNode> = <IStorage<ICustomNode>>this.controlFlowData.get(hostNode);
 
            controlFlowStorage.mergeWith(hostControlFlowStorage, true);
        }
 
        return controlFlowStorage;
    }
 
    /**
     * @param node
     * @return {boolean}
     */
    private isVisitedFunctionNode (node: ESTree.Node): boolean {
        return (
            Node.isFunctionDeclarationNode(node) ||
            Node.isFunctionExpressionNode(node) ||
            Node.isArrowFunctionExpressionNode(node)
        ) && this.visitedFunctionNodes.has(node);
    }
 
    /**
     * @param functionNodeBody
     * @param controlFlowStorage
     */
    private transformFunctionBody (functionNodeBody: ESTree.BlockStatement, controlFlowStorage: IStorage<ICustomNode>): void {
        estraverse.replace(functionNodeBody, {
            enter: (node: ESTree.Node, parentNode: ESTree.Node): any => {
                if (this.isVisitedFunctionNode(node)) {
                    return estraverse.VisitorOption.Skip;
                }
 
                if (!FunctionControlFlowTransformer.controlFlowReplacersMap.has(node.type)) {
                    return node;
                }
 
                if (RandomGeneratorUtils.getMathRandom() > this.options.controlFlowFlatteningThreshold) {
                    return node;
                }
 
                const controlFlowReplacerName: ControlFlowReplacers = <ControlFlowReplacers>FunctionControlFlowTransformer
                    .controlFlowReplacersMap.get(node.type);
 
                return {
                    ...this.controlFlowReplacerFactory(controlFlowReplacerName).replace(node, parentNode, controlFlowStorage),
                    parentNode
                };
            }
        });
    }
}