feat: #717 - Update forgetNodesCreatedInBlock to allow returning a …

…value.

docs/manipulation/performance.md

@@ -180,6 +180,15 @@ project.forgetNodesCreatedInBlock(remember => {
 namespaceDeclaration.getText(); // ok, child was marked to remember
 interfaceDeclaration.getText(); // ok, was explicitly marked to remember
 classDeclaration.getText();     // throws, was forgotten
+
+// alternatively, return the node to remember it
+const node = project.forgetNodesCreatedInBlock(() => {
+    const classDec = sourceFile.getClassOrThrow("MyClass");
+    // ...do a lot of stuff...
+    return classDec;
+});
+
+node.getText(); // ok
 ```
 
 Also, do not be concerned about nesting forget blocks. That is perfectly fine to do:

lib/ts-morph.d.ts

@@ -594,16 +594,16 @@ export declare class Project {
      * Forgets the nodes created in the scope of the passed in block.
      *
      * This is an advanced method that can be used to easily "forget" all the nodes created within the scope of the block.
-     * @param block - Block of code to run.
+     * @param block - Block of code to run. Use the `remember` callback or return a node to remember it.
      */
-    forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => void): void;
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => T): T;
     /**
      * Forgets the nodes created in the scope of the passed in block asynchronously.
      *
      * This is an advanced method that can be used to easily "forget" all the nodes created within the scope of the block.
-     * @param block - Block of code to run.
+     * @param block - Block of code to run. Use the `remember` callback or return a node to remember it.
      */
-    forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => Promise<void>): void;
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => Promise<T>): Promise<T>;
     /**
      * Formats an array of diagnostics with their color and context into a string.
      * @param diagnostics - Diagnostics to get a string of.

src/Project.ts

@@ -579,16 +579,16 @@ export class Project {
      * Forgets the nodes created in the scope of the passed in block.
      *
      * This is an advanced method that can be used to easily "forget" all the nodes created within the scope of the block.
-     * @param block - Block of code to run.
+     * @param block - Block of code to run. Use the `remember` callback or return a node to remember it.
      */
-    forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => void): void;
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => T): T;
     /**
      * Forgets the nodes created in the scope of the passed in block asynchronously.
      *
      * This is an advanced method that can be used to easily "forget" all the nodes created within the scope of the block.
-     * @param block - Block of code to run.
+     * @param block - Block of code to run. Use the `remember` callback or return a node to remember it.
      */
-    forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => Promise<void>): void;
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => Promise<T>): Promise<T>;
     forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => (void | Promise<void>)) {
         return this._context.compilerFactory.forgetNodesCreatedInBlock(block);
     }

src/compiler/ast/common/Node.ts

@@ -914,33 +914,34 @@ export class Node<NodeType extends ts.Node = ts.Node> {
      * Gets the first source file text position that is not whitespace taking into account comment nodes and a previous node's trailing trivia.
      */
     getNonWhitespaceStart(): number {
-        // todo: use forgetNodesCreatedInBlock here
-        const parent = this.getParent() as Node | undefined;
-        const pos = this.getPos();
-        const parentTakesPrecedence = parent != null
-            && !TypeGuards.isSourceFile(parent)
-            && parent.getPos() === pos;
-
-        if (parentTakesPrecedence)
-            return this.getStart(true);
-
-        let startSearchPos: number;
-        const sourceFileFullText = this._sourceFile.getFullText();
-        const previousSibling = this.getPreviousSibling();
-
-        if (previousSibling != null && TypeGuards.isCommentNode(previousSibling))
-            startSearchPos = previousSibling.getEnd();
-        else if (previousSibling != null) {
-            if (hasNewLineInRange(sourceFileFullText, [pos, this.getStart(true)]))
-                startSearchPos = previousSibling.getTrailingTriviaEnd();
-            else
-                startSearchPos = pos;
-        }
-        else {
-            startSearchPos = this.getPos();
-        }
+        return this._context.compilerFactory.forgetNodesCreatedInBlock(() => {
+            const parent = this.getParent() as Node | undefined;
+            const pos = this.getPos();
+            const parentTakesPrecedence = parent != null
+                && !TypeGuards.isSourceFile(parent)
+                && parent.getPos() === pos;
+
+            if (parentTakesPrecedence)
+                return this.getStart(true);
+
+            let startSearchPos: number;
+            const sourceFileFullText = this._sourceFile.getFullText();
+            const previousSibling = this.getPreviousSibling();
+
+            if (previousSibling != null && TypeGuards.isCommentNode(previousSibling))
+                startSearchPos = previousSibling.getEnd();
+            else if (previousSibling != null) {
+                if (hasNewLineInRange(sourceFileFullText, [pos, this.getStart(true)]))
+                    startSearchPos = previousSibling.getTrailingTriviaEnd();
+                else
+                    startSearchPos = pos;
+            }
+            else {
+                startSearchPos = this.getPos();
+            }
 
-        return getNextNonWhitespacePos(sourceFileFullText, startSearchPos);
+            return getNextNonWhitespacePos(sourceFileFullText, startSearchPos);
+        });
     }
 
     /**

src/factories/CompilerFactory.ts

@@ -10,7 +10,7 @@ import { SourceFileStructure, OptionalKind } from "../structures";
 import { WriterFunction } from "../types";
 import { SyntaxKind, ts, TypeFlags, ScriptKind } from "../typescript";
 import { replaceSourceFileForCacheUpdate } from "../manipulation";
-import { EventContainer, FileUtils, KeyValueCache, WeakCache, StringUtils, getTextFromStringOrWriter } from "../utils";
+import { EventContainer, FileUtils, KeyValueCache, WeakCache, StringUtils, getTextFromStringOrWriter, TypeGuards } from "../utils";
 import { DirectoryCache } from "./DirectoryCache";
 import { ForgetfulNodeCache } from "./ForgetfulNodeCache";
 import { kindToWrapperMappings } from "./kindToWrapperMappings";
@@ -615,24 +615,44 @@ export class CompilerFactory {
      * Forgets the nodes created in the block.
      * @param block - Block of code to run.
      */
-    forgetNodesCreatedInBlock(block: (remember: (...node: Node[]) => void) => (void | Promise<void>)): Promise<void> {
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => T): T;
+    /**
+     * Asynchronously forgets the nodes created in the block.
+     * @param block - Block of code to run.
+     */
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => Promise<T>): Promise<T>;
+    forgetNodesCreatedInBlock<T = void>(block: (remember: (...node: Node[]) => void) => (T | Promise<T>)): Promise<T> | T {
         // can't use the async keyword here because exceptions that happen when doing this synchronously need to be thrown
         this.nodeCache.setForgetPoint();
         let wasPromise = false;
+        let result: T | Promise<T>;
         try {
-            const result = block((...nodes) => {
+            result = block((...nodes) => {
                 for (const node of nodes)
                     this.nodeCache.rememberNode(node);
             });
 
-            if (result != null && typeof result.then === "function") {
+            if (TypeGuards.isNode(result))
+                this.nodeCache.rememberNode(result);
+
+            if (isPromise(result)) {
                 wasPromise = true;
-                return result.then(() => this.nodeCache.forgetLastPoint());
+                return result.then(value => {
+                    if (TypeGuards.isNode(value))
+                        this.nodeCache.rememberNode(value);
+
+                    this.nodeCache.forgetLastPoint();
+                    return value;
+                });
             }
         } finally {
             if (!wasPromise)
                 this.nodeCache.forgetLastPoint();
         }
-        return Promise.resolve();
+        return result;
+
+        function isPromise<TValue>(value: unknown): value is Promise<TValue> {
+            return value != null && typeof (value as any).then === "function";
+        }
     }
 }

src/tests/projectTests.ts

@@ -10,6 +10,7 @@ import { SourceFileStructure, StructureKind } from "../structures";
 import { CompilerOptions, ScriptTarget, SyntaxKind, ts, ScriptKind } from "../typescript";
 import { OptionalKindAndTrivia } from "./compiler/testHelpers";
 import * as testHelpers from "./testHelpers";
+import { assert, IsExact } from "conditional-type-checks";
 
 console.log("");
 console.log("TypeScript version: " + ts.version);
@@ -1209,7 +1210,7 @@ describe(nameof(Project), () => {
             let interfaceNode3: Node;
             let interfaceNode4: Node;
             let interfaceNode5: Node;
-            project.forgetNodesCreatedInBlock(remember => {
+            const returnedNode = project.forgetNodesCreatedInBlock(remember => {
                 sourceFile = project.createSourceFile("test.ts", "class MyClass {} namespace MyNamespace { interface Interface1 {} interface Interface2 {} "
                     + "interface Interface3 {} interface Interface4 {} }");
                 sourceFileNotNavigated = project.createSourceFile("test2.ts", "class MyClass {}");
@@ -1227,6 +1228,8 @@ describe(nameof(Project), () => {
                 namespaceKeywordNode = namespaceNode.getFirstChildByKindOrThrow(SyntaxKind.NamespaceKeyword);
                 interfaceNode1 = namespaceNode.getInterfaceOrThrow("Interface1");
                 remember(interfaceNode1);
+
+                return namespaceNode.addInterface({ name: "Interface6" });
             });
 
             it("should not have forgotten the source file", () => {
@@ -1269,6 +1272,10 @@ describe(nameof(Project), () => {
                 expect(interfaceNode5.wasForgotten()).to.be.true;
             });
 
+            it("should not have forgotten the returned node", () => {
+                expect(returnedNode.wasForgotten()).to.be.false;
+            });
+
             it("should not throw if removing a created node in a block", () => {
                 const newSourceFile = project.createSourceFile("file3.ts", "class MyClass {}");
                 project.forgetNodesCreatedInBlock(remember => {
@@ -1291,6 +1298,12 @@ describe(nameof(Project), () => {
                     throw new Error("");
                 })).to.throw();
             });
+
+            it("should get the return value", () => {
+                const result = project.forgetNodesCreatedInBlock(() => 5);
+                assert<IsExact<typeof result, number>>(true);
+                expect(result).to.equal(5);
+            });
         });
 
         describe("asynchronous", () => {
@@ -1299,17 +1312,27 @@ describe(nameof(Project), () => {
                 const sourceFile = project.createSourceFile("file.ts");
                 let interfaceDec: InterfaceDeclaration;
                 let classDec: ClassDeclaration;
-                await project.forgetNodesCreatedInBlock(async remember => {
+                const returnedNode = await project.forgetNodesCreatedInBlock(async remember => {
                     // do something to cause this code to be added to the end of the execution queue
                     await new Promise((resolve, reject) => resolve());
 
                     classDec = sourceFile.addClass({ name: "Class" });
                     interfaceDec = sourceFile.addInterface({ name: "Interface" });
                     remember(interfaceDec);
+                    return sourceFile.addInterface({ name: "ReturnedInterface" });
                 });
 
                 expect(classDec!.wasForgotten()).to.be.true;
                 expect(interfaceDec!.wasForgotten()).to.be.false;
+                expect(returnedNode.wasForgotten()).to.be.false;
+            });
+
+            it("should get the return value", async () => {
+                const project = new Project({ useVirtualFileSystem: true });
+                const resultPromise = project.forgetNodesCreatedInBlock(() => Promise.resolve(5));
+                assert<IsExact<typeof resultPromise, Promise<number>>>(true);
+                const result = await resultPromise;
+                expect(result).to.equal(5);
             });
         });
     });