fix: cleanup JUMP_TAB in generator to support Generator.prototype.return

`Generator.prototype.return` allows to return a result at the location of the last yield expression. We should cleanup the instrumentation variable `JUMP_TAB` in finally blocks to record resolve events.
getappmap · Mar 10, 2023 · 361afc3 · 361afc3
1 parent 0e390ed
commit 361afc3
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 6 deletions.
diff --git a/components/instrumentation/default/visit.mjs b/components/instrumentation/default/visit.mjs
@@ -253,7 +253,10 @@ const makeReturnStatement = (argument) => ({
 const isJumpClosureNode = (node) => {
   if (node.type === "Program") {
     return node.sourceType === "module";
-  } else if (node.type === "FunctionExpression" || node.type === "FunctionDeclaration") {
+  } else if (
+    node.type === "FunctionExpression" ||
+    node.type === "FunctionDeclaration"
+  ) {
     return node.async || node.generator;
   } else if (node.type === "ArrowFunctionExpression") {
     return node.async;
@@ -500,6 +503,38 @@ const instrumentClosure = (node, parent, grand_parent, closure, context) => {
             ]),
           ),
           makeBlockStatement([
+            ...(isJumpClosureNode(node)
+              ? [
+                  makeIfStatement(
+                    makeBinaryExpression(
+                      "!==",
+                      makeIdentifier(`${context.apply}_JUMP_TAB`),
+                      makeLiteral(null),
+                    ),
+                    makeBlockStatement([
+                      makeStatement(
+                        makeCallExpression(
+                          makeRegularMemberExpression(
+                            context.apply,
+                            "recordResolve",
+                          ),
+                          [
+                            makeIdentifier(`${context.apply}_JUMP_TAB`),
+                            makeRegularMemberExpression(context.apply, "empty"),
+                          ],
+                        ),
+                      ),
+                      makeStatement(
+                        makeAssignmentExpression(
+                          makeIdentifier(`${context.apply}_JUMP_TAB`),
+                          makeLiteral(null),
+                        ),
+                      ),
+                    ]),
+                    null,
+                  ),
+                ]
+              : []),
             makeIfStatement(
               makeUnaryExpression("!", makeIdentifier(`${context.apply}_DONE`)),
               makeBlockStatement([
@@ -707,9 +742,36 @@ const instrumenters = {
                 ]),
           ]),
         ),
-        node.finalizer === null
-          ? null
-          : visitNode(node.finalizer, node, parent, closure, context),
+        makeBlockStatement([
+          makeIfStatement(
+            makeBinaryExpression(
+              "!==",
+              makeIdentifier(`${context.apply}_JUMP_TAB`),
+              makeLiteral(null),
+            ),
+            makeBlockStatement([
+              makeStatement(
+                makeCallExpression(
+                  makeRegularMemberExpression(context.apply, "recordResolve"),
+                  [
+                    makeIdentifier(`${context.apply}_JUMP_TAB`),
+                    makeRegularMemberExpression(context.apply, "empty"),
+                  ],
+                ),
+              ),
+              makeStatement(
+                makeAssignmentExpression(
+                  makeIdentifier(`${context.apply}_JUMP_TAB`),
+                  makeLiteral(null),
+                ),
+              ),
+            ]),
+            null,
+          ),
+          ...(node.finalizer === null
+            ? []
+            : [visitNode(node.finalizer, node, parent, closure, context)]),
+        ]),
       );
     } else {
       return null;

diff --git a/components/instrumentation/default/visit.test.mjs b/components/instrumentation/default/visit.test.mjs
@@ -263,6 +263,10 @@ const makeCodeLocation = (source, line, column) =>
             );
             throw APPLY_ERROR;
           } finally {
+            if (APPLY_JUMP_TAB !== null) {
+              APPLY.recordResolve(APPLY_JUMP_TAB, APPLY.empty);
+              APPLY_JUMP_TAB = null;
+            }
             if (!APPLY_DONE) {
               APPLY.recordReturn(
                 APPLY_BUNDLE_TAB,
@@ -336,7 +340,13 @@ const makeCodeLocation = (source, line, column) =>
                   456;
                 }
               } finally {
-                789;
+                if (APPLY_JUMP_TAB !== null) {
+                  APPLY.recordResolve(APPLY_JUMP_TAB, APPLY.empty);
+                  APPLY_JUMP_TAB = null;
+                }
+                {
+                  789;
+                }
               }
             }
           } catch (APPLY_ERROR) {
@@ -352,6 +362,10 @@ const makeCodeLocation = (source, line, column) =>
             );
             throw APPLY_ERROR;
           } finally {
+            if (APPLY_JUMP_TAB !== null) {
+              APPLY.recordResolve(APPLY_JUMP_TAB, APPLY.empty);
+              APPLY_JUMP_TAB = null;
+            }
             if (!APPLY_DONE) {
               APPLY.recordReturn(
                 APPLY_BUNDLE_TAB,
@@ -421,6 +435,11 @@ assertEqual(
         {
           456;
         }
+      } finally {
+        if (APPLY_JUMP_TAB !== null) {
+          APPLY.recordResolve(APPLY_JUMP_TAB, APPLY.empty);
+          APPLY_JUMP_TAB = null;
+        }
       }
       export default 789;
     `,
@@ -451,7 +470,13 @@ assertEqual(
           APPLY_JUMP_TAB = null;
         }
       } finally {
-        456;
+        if (APPLY_JUMP_TAB !== null) {
+          APPLY.recordResolve(APPLY_JUMP_TAB, APPLY.empty);
+          APPLY_JUMP_TAB = null;
+        }
+        {
+          456;
+        }
       }
       export default 789;
     `,