Instrument: Move creating tracker [refactor]

lib/instrument/visitors/class.js

@@ -17,7 +17,7 @@ const t = require('@babel/types');
 // Imports
 const {
 		createFunction, createAndEnterFunctionOrClassNameBlock, visitFunctionParams, visitFunctionBody,
-		removeUnnecessaryUseStrictDirectives, insertTrackerComment, instrumentFunctionOrClassConstructor
+		removeUnnecessaryUseStrictDirectives, insertTrackerComment, createTrackerCall
 	} = require('./function.js'),
 	{visitMethod, getMethodName} = require('./method.js'),
 	Expression = require('./expression.js'),
@@ -28,6 +28,7 @@ const {
 		createBlockId, createBlock, createBlockWithId, createAndEnterBlock, createBindingWithoutNameCheck,
 		createThisBinding, createArgumentsBinding, createNewTargetBinding
 	} = require('../blocks.js'),
+	{insertTrackerCodeIntoFunction} = require('../tracking.js'),
 	{visitKey, visitKeyMaybe, visitKeyContainer} = require('../visit.js'),
 	{createTempVarNode} = require('../internalVars.js'),
 	{FN_TYPE_CLASS} = require('../../shared/constants.js');
@@ -447,8 +448,9 @@ function instrumentClass(
 	parent[key] = classNode;
 
 	// Add tracker code + block vars to constructor
-	instrumentFunctionOrClassConstructor(
-		constructorNode, fn, constructorParamsBlock, constructorBodyBlock, state
+	const trackerNode = createTrackerCall(fn, state);
+	insertTrackerCodeIntoFunction(
+		fn, constructorNode, constructorParamsBlock, constructorBodyBlock, trackerNode, state
 	);
 
 	// Insert tracking comment
@@ -460,7 +462,7 @@ function instrumentClass(
 
 /**
  * Create empty class constructor node.
- * It will be instrumented by `instrumentFunctionOrClassConstructor()`, same as a real constructor.
+ * It will be instrumented by `instrumentClass()`, same as a real constructor.
  * @param {Object} fn - Function object for class
  * @param {Object} state - State object
  * @returns {Object} - Class constructor AST node

lib/instrument/visitors/function.js

@@ -20,7 +20,7 @@ module.exports = {
 	escapeFilename,
 	withStrictModeState,
 	hoistSloppyFunctionDeclarations,
-	instrumentFunctionOrClassConstructor,
+	createTrackerCall,
 	insertTrackerComment,
 	createFunctionInfoFunction
 };
@@ -613,29 +613,12 @@ function instrumentFunction(node, fn, parent, key, paramsBlock, bodyBlock, state
 	// Restore to original place in AST
 	parent[key] = node;
 
-	instrumentFunctionOrClassConstructor(node, fn, paramsBlock, bodyBlock, state);
-
-	if (node.type === 'ArrowFunctionExpression') insertArrowFunctionTrackerComment(fn, node, state);
-}
-
-/**
- * Instrument function, method or class constructor.
- *   - Insert `livepack_tracker()` call in function body.
- *   - Insert block vars (`livepack_scopeId`, `livepack_temp`) in function body or params.
- *   - 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) {
 	// Insert tracker call and block vars (`livepack_scopeId` and `livepack_temp`) into function.
 	// NB: `bodyBlock` here may actually be params block if is an arrow function with no body.
 	const trackerNode = createTrackerCall(fn, state);
 	insertTrackerCodeIntoFunction(fn, node, paramsBlock, bodyBlock, trackerNode, state);
+
+	if (node.type === 'ArrowFunctionExpression') insertArrowFunctionTrackerComment(fn, node, state);
 }
 
 /**