overlookmotel / livepack / 6923877558

/* --------------------

 * livepack module

 * Code instrumentation visitor for functions

 * ------------------*/

'use strict';

// Export

module.exports = {

        ArrowFunctionExpression,

        FunctionDeclaration,

        FunctionExpression,

        createAndEnterFunctionOrClassNameBlock,

        visitFunctionOrMethod,

        visitFunctionParamsAndBody,

        removeUnnecessaryUseStrictDirectives,

        createFunction,

        getFunctionType,

        escapeFilename,

        withStrictModeState,

        hoistSloppyFunctionDeclarations,

        instrumentFunctionOrClassConstructor,

        insertTrackerComment,

        createFunctionInfoFunction

};

// Modules

const mapValues = require('lodash/mapValues'),

        t = require('@babel/types');

// Imports

const Statement = require('./statement.js'),

        Expression = require('./expression.js'),

        {IdentifierLet, AssigneeLet} = require('./assignee.js'),

        {visitKey, visitKeyContainer} = require('../visit.js'),

        {

                createBlock, createAndEnterBlock, createBinding, createThisBinding, createArgumentsBinding,

                createNewTargetBinding, getOrCreateExternalVar

        } = require('../blocks.js'),

        {insertTrackerCodeIntoFunctionBody, insertTrackerCodeIntoFunctionParams} = require('../tracking.js'),

        {createFnInfoVarNode} = require('../internalVars.js'),

        {insertComment, hasUseStrictDirective, stringLiteralWithSingleQuotes} = require('../utils.js'),

        {combineArraysWithDedup} = require('../../shared/functions.js'),

        {

                FN_TYPE_FUNCTION, FN_TYPE_ASYNC_FUNCTION, FN_TYPE_GENERATOR_FUNCTION,

                FN_TYPE_ASYNC_GENERATOR_FUNCTION, TRACKER_COMMENT_PREFIX

        } = require('../../shared/constants.js');

// Exports

/**

 * Visitor for arrow function.

 * @param {Object} node - Arrow function AST node

 * @param {Object} state - State object

 * @param {Object|Array} parent - Parent AST node/container

 * @param {string|number} key - Node's key on parent AST node/container

 * @returns {undefined}

 */

function ArrowFunctionExpression(node, state, parent, key) {

        // Tracker comment is inserted later to avoid it getting relocated if function has complex params

        // which are later relocated to inside function body

        visitFunction(node, parent, key, undefined, false, node.body.type === 'BlockStatement', state);

}

/**

 * Visitor for function declaration.

 * @param {Object} node - Function declaration AST node

 * @param {Object} state - State object

 * @param {Object|Array} parent - Parent AST node/container

 * @param {string|number} key - Node's key on parent AST node/container

 * @returns {undefined}

 */

function FunctionDeclaration(node, state, parent, key) {

        // Visit function

        const fnName = node.id.name;

        const fn = visitFunction(node, parent, key, fnName, true, true, state);

        // Create binding for function name

        // TODO: Whether the value of the function is hoisted depends on whether is a top level statement

        // (including in a labeled statement)

        // `() => { console.log(typeof x); function x() {} }` -> Logs 'function'

        // `() => { console.log(typeof x); q: function x() {} }` -> Logs 'function'

        // `() => { console.log(typeof x); if (true) function x() {} }` -> Logs 'undefined'

        const {currentBlock: block, currentHoistBlock: hoistBlock} = state;

        let binding = block.bindings[fnName];

                binding = createBinding(block, fnName, {isVar: isTopLevel, isFunction: true}, state);

                // If is a sloppy-mode function declaration which may need to be hoisted, record that.

                // Determine whether to hoist after 1st pass is complete, once all other bindings created.

                if (!isTopLevel && !state.isStrict && !node.async && !node.generator) {

                        state.sloppyFunctionDeclarations.push({

                                varName: fnName,

                                binding,

                                block,

                                hoistBlock,

                                parentFn: state.currentFunction

                        });

                }

        } else if (!binding.isFunction) {

                // Existing binding is a `var` declaration.

                // Flag binding as a function so `var` declarations' identifiers don't get renamed.

                binding.isFunction = true;

                binding.trails.length = 0;

        }

        // Insert tracker comment

        insertFunctionDeclarationOrExpressionTrackerComment(fn, node, state);

}

/**

 * Visitor for function expression.

 * @param {Object} node - Function expression AST node

 * @param {Object} state - State object

 * @param {Object|Array} parent - Parent AST node/container

function FunctionExpression(node, state, parent, key) {

        // If not anonymous, create and enter name block

        let nameBlock, fnName;

        const idNode = node.id;

                fnName = idNode.name;

                nameBlock = createAndEnterFunctionOrClassNameBlock(fnName, true, state);

        }

        // Visit function

        const fn = visitFunction(node, parent, key, fnName, true, true, state);

        if (nameBlock) {

                // Exit name block

                state.currentBlock = nameBlock.parent;

                // If contains `eval()`, change function ID to make name block be considered internal to function.

                // An external var representing the function is pointless as its name will be frozen anyway

                // due to being in scope of the `eval()`.

                if (fn.containsEval) fn.id = nameBlock.id;

        }

        // Insert tracker comment

        insertFunctionDeclarationOrExpressionTrackerComment(fn, node, state);

}

/**

 * Create block for function/class name accessed within the function/class.

 * @param {string} fnName - Function name

 * @param {boolean} isSilentConst - `true` if assignment to function name silently fails in sloppy mode

 * @returns {Object} - Block object

 */

function createAndEnterFunctionOrClassNameBlock(fnName, isSilentConst, state) {

        createBinding(block, fnName, {isConst: true, isSilentConst, isFunction: true}, state);

        return block;

}

/**

 * Visit function declaration, function expression or arrow function.

 * @param {boolean} hasBodyBlock - `true` if has body block (only arrow functions can not)

 * @param {Object} state - State object

 * @returns {Object} - Function object

 */

function visitFunction(node, parent, key, fnName, isFullFunction, hasBodyBlock, state) {

        return withStrictModeState(

                node, hasBodyBlock, state,

                (isStrict, isEnteringStrict) => visitFunctionOrMethod(

                        node, parent, key, fnName, isStrict, isEnteringStrict, isFullFunction, hasBodyBlock, state

                )

        );

}

/**

 * Visit function or method.

 * @param {Object} fnNode - Function declaration, function expression, arrow function,

 *   class method, class private method or object method AST node

 * @param {Object|Array} parent - Parent AST node/container

 * @param {string|number} key - Node's key on parent AST node/container

 * @param {string} [fnName] - Function name (`undefined` if unnamed)

 * @param {boolean} isStrict - `true` if function is strict mode

 * @param {boolean} isEnteringStrict - `true` if function is strict mode and parent is sloppy mode

 * @param {boolean} isFullFunction - `true` if is not arrow function

 * @param {boolean} hasBodyBlock - `true` if has body block (only arrow functions can not)

 * @param {Object} state - State object

 * @returns {Object} - Function object

 */

function visitFunctionOrMethod(

        fnNode, parent, key, fnName, isStrict, isEnteringStrict, isFullFunction, hasBodyBlock, state

) {

        // Create params block. If is full function, create bindings for `this` + `new.target`.

        const paramsBlock = createAndEnterBlock(fnName, false, state);

        let parentThisBlock;

        if (isFullFunction) {

                createThisBinding(paramsBlock);

                createNewTargetBinding(paramsBlock);

                parentThisBlock = state.currentThisBlock;

        // Create and enter function

        const fn = createFunction(paramsBlock.id, fnNode, isStrict, state);

        state.currentFunction = fn;

        const parentTrail = state.trail;

        state.trail = [];

        // Create body block or make params block a vars block

        let bodyBlock;

        if (hasBodyBlock) {

                bodyBlock = createBlock(fnName, true, state);

        const argNames = visitFunctionParamsAndBody(fn, fnNode, paramsBlock, bodyBlock, isStrict, state);

        // If is full function, create binding for `arguments` in params block.

        // Param called `arguments` takes precedence.

        if (isFullFunction) {

                if (!paramsBlock.bindings.arguments) createArgumentsBinding(paramsBlock, isStrict, argNames);

                state.currentThisBlock = parentThisBlock;

        }

        // Serialize AST

        fn.astJson = serializeFunctionAst(fnNode, hasBodyBlock, isStrict, isEnteringStrict);

        // Exit function

        state.currentBlock = paramsBlock.parent;

        state.currentFunction = fn.parent;

        state.trail = parentTrail;

        // Temporarily remove from AST so is not included in parent function's serialized AST

        parent[key] = null;

        // Queue instrumentation of function in 2nd pass

        state.secondPass(instrumentFunction, fnNode, fn, parent, key, paramsBlock, bodyBlock, state);

        return fn;

}

/**

 * Visit function params and body.

 * Caller must ensure `state.currentBlock` is params block,

 * and `state.currentThisBlock` is params block if is a full function.

 * Caller must set `state.currentBlock` and `state.currentThisBlock` back afterwards.

 * @param {Object} fn - Function object

 * @param {Object} fnNode - Function/method AST node

 * @param {boolean} isStrict - `true` if function is strict mode

 * @param {Object} state - State object

 * @returns {Array<string>} - Array of argument names which are linked to argument vars.

 *   Will be empty array if strict mode, or arguments are not all simple identifiers.

 */

function visitFunctionParamsAndBody(fn, fnNode, paramsBlock, bodyBlock, isStrict, state) {

        // Visit params and find first complex param (i.e. param which is not just an identifier).

        // Do not consider `...x` to be complex.

        // Also get array of params which are linked to `arguments`.

        // i.e. `((x) => { arguments[0] = 2; return x; })(1) === 2`

        // `arguments` is only linked in sloppy mode functions with no complex params.

        const argNames = [];

        let hasLinkedArguments = !isStrict,

                firstComplexParamIndex;

        visitKeyContainer(fnNode, 'params', (paramNode, _state, paramNodes, index) => {

                if (paramNode.type === 'Identifier') {

                        if (hasLinkedArguments) argNames.push(paramNode.name);

                        IdentifierLet(paramNode, state, paramNodes, index);

                } else {

                        if (firstComplexParamIndex === undefined) {

                                hasLinkedArguments = false;

                                argNames.length = 0;

                                if (paramNode.type !== 'RestElement' || paramNode.argument.type !== 'Identifier') {

                                        firstComplexParamIndex = index;

                                        // Make params block the vars block. Tracker call and scope ID var will go in params.

                                        if (bodyBlock) paramsBlock.varsBlock = bodyBlock.varsBlock = paramsBlock;

                                }

                        }

                        AssigneeLet(paramNode, state, paramNodes, index);

                }

        }, state);

        // Record index of first complex param

        fn.firstComplexParamIndex = firstComplexParamIndex;

        // Visit body

        if (bodyBlock) {

                state.currentBlock = bodyBlock;

                const parentHoistBlock = state.currentHoistBlock;

                state.currentHoistBlock = bodyBlock;

                visitKey(fnNode, 'body', FunctionBodyBlock, state);

                state.currentHoistBlock = parentHoistBlock;

        } else {

                visitKey(fnNode, 'body', Expression, state);

        }

        // Return array of params which are linked to `arguments`

        return argNames;

}

/**

 * Visitor for function body block.

 * @param {Object} node - Function body block AST node

 * @param {Object} state - State object

 * @returns {undefined}

 */

function FunctionBodyBlock(node, state) {

        visitKeyContainer(node, 'body', Statement, state);

}

/**

 * Serialize function AST to JSON.

 * Remove unnecessary 'use strict' directives if present.

 * Do not mutate the node passed in, to ensure these changes only affect the serialized AST,

 * and not the instrumented output.

 * @param {Object} fnNode - Function AST node

 * @param {boolean} hasBodyBlock - `true` if has body block (only arrow functions can not)

 * @param {boolean} isStrict - `true` if function is strict mode

 * @param {boolean} isEnteringStrict - `true` if function transitions into strict mode

 * @returns {string} - Function AST as JSON

 */

function serializeFunctionAst(fnNode, hasBodyBlock, isStrict, isEnteringStrict) {

        // Remove unnecessary 'use strict' directives

        if (hasBodyBlock && isStrict) {

                const alteredFnNode = removeUnnecessaryUseStrictDirectives(fnNode, !isEnteringStrict);

                if (alteredFnNode) fnNode = alteredFnNode;

        }

        // Stringify AST to JSON

        return JSON.stringify(fnNode);

}

/**

 * Remove unnecessary 'use strict' directives from function body.

 * Relocate any comments from deleted directives.

 * Do not mutate input - clone if needs alteration.

 * Return `undefined` if no alteration necessary.

 * @param {Object} fnNode - Function AST node

 * @param {boolean} needsNoDirective - `true` if function already strict mode from outer environment

function removeUnnecessaryUseStrictDirectives(fnNode, needsNoDirective) {

        let bodyNode = fnNode.body;

        const directiveNodes = bodyNode.directives;

        if (directiveNodes.length === 0) return undefined;

        // Remove unnecessary directives

        const newDirectiveNodes = [],

                commentNodes = [];

        for (let directiveNode of directiveNodes) {

                        if (needsNoDirective) {

                                commentNodes.push(

                                        ...(directiveNode.leadingComments || []),

                                        ...(directiveNode.trailingComments || [])

                                );

                                continue;

                        }

                        needsNoDirective = true;

                }

                // Add any comments removed from previous directives

                newDirectiveNodes.push(directiveNode);

        }

        // If no directives need to be removed, return `undefined`

        if (newDirectiveNodes.length === directiveNodes.length) return undefined;

        // Clone function node

        bodyNode = {...bodyNode, directives: newDirectiveNodes};

        fnNode = {...fnNode, body: bodyNode};

        // Add any remaining comments to first statement / last directive / function body

                        const firstStatementNode = statementNodes[0];

                        statementNodes[0] = {

                                ...firstStatementNode,

                                leadingComments: combineArraysWithDedup(commentNodes, firstStatementNode.leadingComments)

                } else if (newDirectiveNodes.length !== 0) {

        // Return clone of function node

        return fnNode;

}

/**

 * Create object representing function.

 * @param {number} id - Function ID

 * @param {Object} node - Function/method/class AST node

 * @param {boolean} isStrict - `true` if function is strict mode

 * @param {Object} state - State object

 * @returns {Object} - Function object

 */

function createFunction(id, node, isStrict, state) {

        const parentFunction = state.currentFunction;

        const fn = {

                id,

                node,

                astJson: undefined,

                isStrict,

                trail: parentFunction ? [...state.trail] : undefined,

                bindings: [],

                externalVars: new Map(), // Keyed by block

                scopes: undefined,

                globalVarNames: new Set(),

                amendments: [],

                superIsProto: false,

                containsImport: false,

                hasSuperClass: false,

                firstSuperStatementIndex: undefined,

                returnsSuper: false,

                hasThisBindingForSuper: false,

                firstComplexParamIndex: undefined

        };

        if (parentFunction) parentFunction.children.push(fn);

        state.functions.push(fn);

        state.fileContainsFunctionsOrEval = true;

        return fn;

}

/**

 * Get type code for function.

 * @param {Object} fnNode - Function AST node

 * @returns {string} - Function type code

 */

function getFunctionType(fnNode) {

        return fnNode.async

                ? fnNode.generator

                        ? FN_TYPE_ASYNC_GENERATOR_FUNCTION

                        : FN_TYPE_ASYNC_FUNCTION

                : fnNode.generator

                        ? FN_TYPE_GENERATOR_FUNCTION

                        : FN_TYPE_FUNCTION;

}

/**

 * Escape filename so it can be safely included in tracker comment.

 * @param {string} filename - File path

 * @returns {string} - Filename escaped

 */

function escapeFilename(filename) {

        // Encode filename as JSON to escape non-ascii chars.

        // Convert `*/` to `*\/` so does not terminate comment early.

        return JSON.stringify(filename).slice(1, -1).replace(/\*\//g, '*\\/');

}

function withStrictModeState(fnNode, hasBodyBlock, state, visit) {

        // Get if strict mode

        let {isStrict} = state,

                isEnteringStrict = false;

        if (!isStrict && hasBodyBlock && hasUseStrictDirective(fnNode.body)) {

                isStrict = isEnteringStrict = state.isStrict = true;

        }

        // Call visitor

        const res = visit(isStrict, isEnteringStrict);

        // If entered strict mode, exit it again

        if (isEnteringStrict) state.isStrict = false;

        // Return visitor's return value

        return res;

}

 * @returns {undefined}

 */

function hoistSloppyFunctionDeclarations(state) {

        for (const declaration of state.sloppyFunctionDeclarations) {

                hoistSloppyFunctionDeclaration(declaration);

        }

}

/**

 * Hoist function declaration, if it can be hoisted.

 * Should only be called for functions which:

 *   1. are not already in the top block in parent function / program

 *   2. are not async or generator functions

 *

 * @param {Object} declaration - Declaration object

 * @param {string} declaration.varName - Function name

 * @param {Object} declaration.binding - Binding object for binding in block where originally declared

 * @param {Object} declaration.hoistBlock - Block object for block where could be hoisted to

 * @param {Object} [declaration.parentFn] - Function object for parent function

 * @returns {undefined}

 */

function hoistSloppyFunctionDeclaration(declaration) {

        const {varName, hoistBlock} = declaration;

        // Do not hoist if existing const, let or class declaration binding in hoist block

        // If parent function's params include var with same name, do not hoist.

        // NB: `hoistBlock.parent` here is either parent function params block or file block.

        // In CommonJS code, file block is CommonJS wrapper function's params, and in any other context

        // file block contains no bindings except `this`. So it's safe to treat as a params block in all cases.

        // If any binding in intermediate blocks, do not hoist.

        // NB: This includes function declarations which will be themselves hoisted.

        let inBetweenBlock = declaration.block.parent;

                if (inBetweenBlock.bindings[varName]) return;

                inBetweenBlock = inBetweenBlock.parent;

        }

        // Can be hoisted

                hoistBlock.bindings[varName] = {...declaration.binding, isVar: true};

        } else if (!hoistBlockBinding.isFunction) {

                // Existing binding is a `var` declaration.

                // Flag binding as a function so `var` declarations' identifiers don't get renamed.

                hoistBlockBinding.isFunction = true;

                hoistBlockBinding.trails.length = 0;

        }

}

/**

 * Instrument function.

 * @param {Object} node - Function or method AST node

 * @param {Object} fn - Function object

 * @param {Object|Array} parent - Parent AST node/container

 * @param {string|number} key - Node's key on parent AST node/container

 * @param {Object} paramsBlock - Function's params block object

 * @param {Object} [bodyBlock] - Function's body block object (if it has one, arrow functions may not)

 * @param {Object} state - State object

 * @returns {undefined}

 */

function instrumentFunction(node, fn, parent, key, paramsBlock, bodyBlock, state) {

        // Restore to original place in AST

        instrumentFunctionOrClassConstructor(node, fn, paramsBlock, bodyBlock, state);

}

 * Instrument function, method or class constructor.

 *   - Insert `livepack_tracker()` call in function body.

 *   - Create function info function containing function AST, details of internal and external vars etc.

 *

 * @param {Object} node - Function, method, or class constructor AST node

 * @param {Object} fn - Function object

 * @param {Object} paramsBlock - Function's params block object

 * @param {Object} bodyBlock - Function's body block object

 * @param {Object} state - State object

 * @returns {undefined}

 */

function instrumentFunctionOrClassConstructor(node, fn, paramsBlock, bodyBlock, state) {

        // NB `bodyBlock` here may actually be params block if is an arrow function with no body.

        // Sort scopes in ascending order of block ID.

        const scopes = [...fn.externalVars].map(([block, vars]) => ({block, vars}))

                .sort((scope1, scope2) => (scope1.block.id > scope2.block.id ? 1 : -1));

        fn.scopes = scopes;

        // Add external vars to parent function where external to that function too

        const parentFunction = fn.parent;

        if (parentFunction) {

                // External vars

                const {id: parentFunctionId, externalVars: parentExternalVars} = parentFunction;

                for (const {block, vars} of scopes) {

                        if (block.id >= parentFunctionId) break;

                        for (const [varName, {binding}] of Object.entries(vars)) {

                                getOrCreateExternalVar(parentExternalVars, block, varName, binding);

                        }

                }

        }

        // If function has complex params, insert into params. Otherwise, into function body.

        const trackerNode = createTrackerCall(fn.id, scopes, state);

        const {firstComplexParamIndex} = fn;

        if (firstComplexParamIndex === undefined) {

                insertTrackerCodeIntoFunctionBody(node, paramsBlock, bodyBlock, trackerNode, state);

        } else {

                insertTrackerCodeIntoFunctionParams(

                        node, paramsBlock, bodyBlock, trackerNode, firstComplexParamIndex, state

                );

        }

}

/**

 * Create tracker call.

 * @param {number} fnId - Function ID

 * @param {Array<Object>} scopes - Array of scope objects in ascending order of block ID

 * @param {Object} state - State object

 * @returns {undefined}

 */

function createTrackerCall(fnId, scopes, state) {

                                scope.block.varsBlock.scopeIdVarNode,

                                ...Object.values(scope.vars).map(blockVar => blockVar.binding.varNode)

function insertFunctionDeclarationOrExpressionTrackerComment(fn, node, state) {

        // Insert tracker comment before identifier, or before 1st param or before function body

        const commentHolderNode = node.id || (node.params.length !== 0 ? node.params[0] : node.body);

        insertTrackerComment(fn.id, getFunctionType(node), commentHolderNode, 'leading', state);

}

function insertArrowFunctionTrackerComment(fn, node, state) {

        // Insert tracker comment before first param. If no params, before function body.

        const paramNodes = node.params;

        insertTrackerComment(

                fn.id, node.async ? FN_TYPE_ASYNC_FUNCTION : FN_TYPE_FUNCTION,

                paramNodes.length !== 0 ? paramNodes[0] : node.body, 'leading', state

        );

}

 * @param {string} commentType - 'leading' / 'inner' / 'trailing'

 * @param {Object} state - State object

 * @returns {undefined}

 */

function insertTrackerComment(fnId, fnType, commentHolderNode, commentType, state) {

        insertComment(

                commentHolderNode, commentType,

                `${TRACKER_COMMENT_PREFIX}${fnId};${fnType};${state.filenameEscaped}`

        );

}

/**

 * Create function info function containing function AST, details of internal and external vars etc.

 * @param {Object} fn - Function object

 * @param {Object} state - State object

 * @returns {Object} - AST node for function info function

 */

function createFunctionInfoFunction(fn, state) {

        // Compile internal vars

        const internalVars = Object.create(null);

        for (const {name: varName, trails, isFunction} of fn.bindings) {

                // but can be included if binding is initially a `var` statement e.g. `var x; function x() {}`.

                // So skip them here.

                if (isFunction) continue;

                const internalVar = internalVars[varName] || (internalVars[varName] = []);

                internalVar.push(...trails);

        }

        // Create JSON function info string

        const {children} = fn;

        let argNames;

        let json = JSON.stringify({

                scopes: fn.scopes.map(({block, vars}) => ({

                        blockId: block.id,

                        blockName: block.name,

                        vars: mapValues(vars, (varProps, varName) => {

                                if (varName === 'arguments') argNames = varProps.binding.argNames;

                                return {

                                        isReadFrom: varProps.isReadFrom || undefined,

                                        isAssignedTo: varProps.isAssignedTo || undefined,

                                        isFrozenInternalName: varProps.binding.isFunction || undefined,

                                        trails: varProps.trails

                                };

                        })

                })),

                isStrict: fn.isStrict || undefined,

                superIsProto: fn.superIsProto || undefined,

                containsEval: fn.containsEval || undefined,

                containsImport: fn.containsImport || undefined,

                argNames,

                internalVars,

                globalVarNames: fn.globalVarNames.size !== 0 ? [...fn.globalVarNames] : undefined,

                        ? fn.amendments.map(({type, blockId, trail}) => [type, blockId, ...trail]).reverse()

                        : undefined,

                hasSuperClass: fn.hasSuperClass || undefined,

                firstSuperStatementIndex: fn.firstSuperStatementIndex,

                returnsSuper: fn.returnsSuper || undefined,

                childFns: children.map(childFn => childFn.trail)

        });

        // Add AST JSON to JSON

        json = `${json.slice(0, -1)},"ast":${fn.astJson}}`;

        // Create function returning JSON string and references to function info function for child fns.

        // Make output as single-quoted string for shorter output (not escaping every `"`).

                        stringLiteralWithSingleQuotes(json),

                        t.arrayExpression(children.map(childFn => createFnInfoVarNode(childFn.id, state))),

                        state.getSourcesNode

                ]))

        ]));

}