Refactor Hooks

packages/engine/src/executor.ts

@@ -1,6 +1,6 @@
 import { autorun, runInAction } from "mobx";
 import { TypeCellContext } from "./context";
-import { installHooks } from "./hookDisposables";
+import { executeWithHooks } from "./hooks";
 import { Module } from "./modules";
 import { isStored } from "./storage/stored";
 import { isView } from "./view";
@@ -96,7 +96,7 @@ export async function runModule(
     );
   }
 
-  const execute = async () => {
+  async function execute(this: any) {
     try {
       if (wouldLoopOnAutorun) {
         detectedLoop = true;
@@ -114,25 +114,22 @@ export async function runModule(
       disposeEveryRun.length = 0; // clear existing array in this way, because we've passed the reference to resolveDependencyArray and want to keep it intact
 
       beforeExecuting();
-      const hooks = installHooks();
-      disposeEveryRun.push(hooks.disposeAll);
-      let executionPromise: Promise<any>;
+
+      const { executeModel, disposeHooks } = executeWithHooks(
+        (hookContext) =>
+          mod.factoryFunction.apply(hookContext, argsToCallFunctionWith) // TODO: what happens with disposers if a rerun of this function is slow / delayed?
+      );
+
+      disposeEveryRun.push(disposeHooks);
+
       try {
-        executionPromise = mod.factoryFunction.apply(
-          undefined,
-          argsToCallFunctionWith
-        ); // TODO: what happens with disposers if a rerun of this function is slow / delayed?
+        await executeModel();
       } finally {
-        // Hooks are only installed for sync code. Ideally, we'd want to run it for all code, but then we have the chance hooks will affect other parts of the TypeCell (non-user) code
-        // (we ran into this that notebooks wouldn't be saved (_.debounce), and also that setTimeout of Monaco blink cursor would be hooked)
-        hooks.unHookAll();
         if (previousVariableDisposer) {
           previousVariableDisposer(exports);
         }
       }
 
-      await executionPromise;
-
       // Running the assignments to `context` in action should be a performance improvement to prevent triggering observers one-by-one
       wouldLoopOnAutorun = true;
       runInAction(() => {
@@ -211,7 +208,7 @@ export async function runModule(
       //reject(e);
       resolve();
     }
-  };
+  }
 
   const autorunDisposer = autorun(() => execute());
 

packages/engine/src/hooks.ts

@@ -0,0 +1,71 @@
+// These will be injected in the compiled function and link to hookContext
+export const overrideFunctions = [
+  "setTimeout",
+  "setInterval",
+  "console",
+  "EventTarget",
+] as const;
+
+function applyDisposer<T, Y>(
+  original: (...args: T[]) => Y,
+  disposes: Array<() => void>,
+  disposer: (ret: Y, args: T[]) => void
+) {
+  return function newFunction(this: any): Y {
+    const callerArguments = arguments;
+    const ret = original.apply(this, callerArguments as any); // TODO: fix any?
+    const ctx = this;
+    disposes.push(() => disposer.call(ctx, ret, callerArguments as any));
+    return ret;
+  };
+}
+
+export function executeWithHooks<T>(
+  modelFunction: (hookContext: any) => Promise<T>
+) {
+  const disposers: Array<() => void> = [];
+
+  const executeModel = async function (this: any) {
+    const hookContext: { [K in typeof overrideFunctions[number]]: any } = {
+      setTimeout: applyDisposer(setTimeout, disposers, (ret) => {
+        clearTimeout(ret);
+      }),
+      setInterval: applyDisposer(setInterval, disposers, (ret) => {
+        clearInterval(ret);
+      }),
+      console: {
+        ...console,
+        log: (...args: any) => {
+          // TODO: broadcast output to console view
+          console.log(...args);
+        },
+      },
+      EventTarget: undefined,
+    };
+
+    if (typeof EventTarget !== "undefined") {
+      (hookContext.EventTarget as any) = {
+        EventTarget,
+        prototype: {
+          ...EventTarget.prototype,
+          addEventListener: applyDisposer(
+            EventTarget.prototype.addEventListener as any,
+            disposers,
+            function (this: any, _ret, args) {
+              this.removeEventListener(args[0], args[1]);
+            }
+          ),
+        },
+      };
+    }
+
+    return modelFunction(hookContext);
+  };
+
+  return {
+    executeModel,
+    disposeHooks: () => {
+      disposers.forEach((d) => d());
+    },
+  };
+}

packages/engine/src/modules.ts

@@ -1,5 +1,6 @@
 import { TypeCellContext } from "./context";
 import { observable, untracked, computed, autorun } from "mobx";
+import { overrideFunctions } from "./hooks";
 // import { stored } from "./storage/stored";
 // import { view } from "./view";
 
@@ -112,5 +113,16 @@ export function getModulesFromTypeCellCode(compiledCode: string, scope: any) {
     "$1async function"
   ); // TODO: remove await?
 
+  // Adds overridden functions
+  // TODO: improve injection method
+  totalCode = totalCode.replace(
+    /("use strict";)/,
+    `"use strict";
+// Override functinos
+${overrideFunctions
+  .map((hook: string) => `let ${hook} = this.${hook};`)
+  .join("\n")}\n`
+  );
+
   return getModulesFromCode(totalCode, scope);
 }