Increased control flow transformation speed.

dist/index.js

@@ -5086,6 +5086,7 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
         var _this = _possibleConstructorReturn(this, (FunctionControlFlowTransformer.__proto__ || Object.getPrototypeOf(FunctionControlFlowTransformer)).call(this, options));
 
         _this.controlFlowData = new Map();
+        _this.visitedFunctionNodes = new Set();
         _this.hostNodesWithControlFlowNode = [];
         _this.controlFlowStorageFactory = controlFlowStorageFactory;
         _this.controlFlowReplacerFactory = controlFlowReplacerFactory;
@@ -5100,7 +5101,7 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
 
             return {
                 leave: function leave(node, parentNode) {
-                    if (Node_1.Node.isFunctionDeclarationNode(node) || Node_1.Node.isFunctionExpressionNode(node)) {
+                    if (Node_1.Node.isFunctionDeclarationNode(node) || Node_1.Node.isFunctionExpressionNode(node) || Node_1.Node.isArrowFunctionExpressionNode(node)) {
                         return _this2.transformNode(node, parentNode);
                     }
                 }
@@ -5109,10 +5110,11 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
     }, {
         key: "transformNode",
         value: function transformNode(functionNode, parentNode) {
-            if (Node_1.Node.isArrowFunctionExpressionNode(functionNode)) {
+            this.visitedFunctionNodes.add(functionNode);
+            if (!Node_1.Node.isBlockStatementNode(functionNode.body)) {
                 return functionNode;
             }
-            var hostNode = FunctionControlFlowTransformer_1.getHostNode(functionNode);
+            var hostNode = FunctionControlFlowTransformer_1.getHostNode(functionNode.body);
             var controlFlowStorage = this.getControlFlowStorage(hostNode);
             this.controlFlowData.set(hostNode, controlFlowStorage);
             this.transformFunctionBody(functionNode.body, controlFlowStorage);
@@ -5138,6 +5140,11 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
             }
             return controlFlowStorage;
         }
+    }, {
+        key: "isVisitedFunctionNode",
+        value: function isVisitedFunctionNode(node) {
+            return (Node_1.Node.isFunctionDeclarationNode(node) || Node_1.Node.isFunctionExpressionNode(node) || Node_1.Node.isArrowFunctionExpressionNode(node)) && this.visitedFunctionNodes.has(node);
+        }
     }, {
         key: "transformFunctionBody",
         value: function transformFunctionBody(functionNodeBody, controlFlowStorage) {
@@ -5145,6 +5152,9 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
 
             estraverse.replace(functionNodeBody, {
                 enter: function enter(node, parentNode) {
+                    if (_this3.isVisitedFunctionNode(node)) {
+                        return estraverse.VisitorOption.Skip;
+                    }
                     if (!FunctionControlFlowTransformer_1.controlFlowReplacersMap.has(node.type)) {
                         return node;
                     }
@@ -5158,10 +5168,10 @@ var FunctionControlFlowTransformer = FunctionControlFlowTransformer_1 = function
         }
     }], [{
         key: "getHostNode",
-        value: function getHostNode(functionNode) {
-            var blockScopesOfNode = NodeUtils_1.NodeUtils.getBlockScopesOfNode(functionNode);
+        value: function getHostNode(functionNodeBody) {
+            var blockScopesOfNode = NodeUtils_1.NodeUtils.getBlockScopesOfNode(functionNodeBody);
             if (blockScopesOfNode.length === 1) {
-                return functionNode.body;
+                return functionNodeBody;
             } else {
                 blockScopesOfNode.pop();
             }

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

@@ -48,7 +48,12 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     /**
      * @type {Map<ESTree.Node, IStorage<ICustomNode>>}
      */
-    private controlFlowData: Map <ESTree.Node, IStorage<ICustomNode>> = new Map();
+    private readonly controlFlowData: Map <ESTree.Node, IStorage<ICustomNode>> = new Map();
+
+    /**
+     * @type {Set<ESTree.Function>}
+     */
+    private readonly visitedFunctionNodes: Set<ESTree.Function> = new Set();
 
     /**
      * @type {TNodeWithBlockStatement[]}
@@ -90,14 +95,14 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     }
 
     /**
-     * @param functionNode
+     * @param functionNodeBody
      * @returns {TNodeWithBlockStatement}
      */
-    private static getHostNode (functionNode: ESTree.FunctionDeclaration | ESTree.FunctionExpression): TNodeWithBlockStatement {
-        const blockScopesOfNode: TNodeWithBlockStatement[] = NodeUtils.getBlockScopesOfNode(functionNode);
+    private static getHostNode (functionNodeBody: ESTree.BlockStatement): TNodeWithBlockStatement {
+        const blockScopesOfNode: TNodeWithBlockStatement[] = NodeUtils.getBlockScopesOfNode(functionNodeBody);
 
         if (blockScopesOfNode.length === 1) {
-            return functionNode.body;
+            return functionNodeBody;
         } else {
             blockScopesOfNode.pop();
         }
@@ -119,7 +124,11 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     public getVisitor (): IVisitor {
         return {
             leave: (node: ESTree.Node, parentNode: ESTree.Node) => {
-                if (Node.isFunctionDeclarationNode(node) || Node.isFunctionExpressionNode(node)) {
+                if (
+                    Node.isFunctionDeclarationNode(node) ||
+                    Node.isFunctionExpressionNode(node) ||
+                    Node.isArrowFunctionExpressionNode(node)
+                ) {
                     return this.transformNode(node, parentNode);
                 }
             }
@@ -129,14 +138,16 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     /**
      * @param functionNode
      * @param parentNode
-     * @returns {ESTree.Node}
+     * @returns {ESTree.Function}
      */
-    public transformNode (functionNode: ESTree.Function, parentNode: ESTree.Node): ESTree.Node {
-        if (Node.isArrowFunctionExpressionNode(functionNode)) {
+    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);
+        const hostNode: TNodeWithBlockStatement = FunctionControlFlowTransformer.getHostNode(functionNode.body);
         const controlFlowStorage: IStorage<ICustomNode> = this.getControlFlowStorage(hostNode);
 
         this.controlFlowData.set(hostNode, controlFlowStorage);
@@ -175,6 +186,18 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
         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
@@ -182,6 +205,10 @@ export class FunctionControlFlowTransformer extends AbstractNodeTransformer {
     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;
                 }

test/functional-tests/node-transformers/control-flow-transformers/function-control-flow-transformer/FunctionControlFlowTransformer.spec.ts

@@ -178,5 +178,41 @@ describe('FunctionControlFlowTransformer', () => {
                 assert.notMatch(obfuscatedCode, controlFlowStorageMatch);
             });
         });
+
+        describe('arrow function expression', () => {
+            describe('variant #1 - arrow function expression with body', () => {
+                const obfuscationResult: IObfuscationResult = JavaScriptObfuscator.obfuscate(
+                    readFileAsString(__dirname + '/fixtures/arrow-function-expression-with-body.js'),
+                    {
+                        ...NO_CUSTOM_NODES_PRESET,
+                        controlFlowFlattening: true,
+                        controlFlowFlatteningThreshold: 1
+                    }
+                );
+                const obfuscatedCode: string = obfuscationResult.getObfuscatedCode();
+                const regexp: RegExp = new RegExp(rootControlFlowStorageNodeMatch);
+
+                it('should add `control flow storage` node to the obfuscated code', () => {
+                    assert.match(obfuscatedCode, regexp);
+                });
+            });
+
+            describe('variant #2 - arrow function expression without body', () => {
+                const obfuscationResult: IObfuscationResult = JavaScriptObfuscator.obfuscate(
+                    readFileAsString(__dirname + '/fixtures/arrow-function-expression-without-body.js'),
+                    {
+                        ...NO_CUSTOM_NODES_PRESET,
+                        controlFlowFlattening: true,
+                        controlFlowFlatteningThreshold: 1
+                    }
+                );
+                const obfuscatedCode: string = obfuscationResult.getObfuscatedCode();
+                const regexp: RegExp = new RegExp(`var *${variableMatch} *= *\\(\\) *=> *0x1 *\\+ *0x2;`);
+
+                it('shouldn\'t add `control flow storage` node to the obfuscated code', () => {
+                    assert.match(obfuscatedCode, regexp);
+                });
+            });
+        });
     });
 });

test/functional-tests/node-transformers/control-flow-transformers/function-control-flow-transformer/fixtures/arrow-function-expression-with-body.js

@@ -0,0 +1,5 @@
+(function () {
+    var foo = () => {
+        var variable = 1 + 2;
+    }
+})();

test/functional-tests/node-transformers/control-flow-transformers/function-control-flow-transformer/fixtures/arrow-function-expression-without-body.js

@@ -0,0 +1,3 @@
+(function () {
+    var foo = () => 1 + 2;
+})();