fix: #192 - Fix forget block crashes when removing node.

dsherret · Dec 23, 2017 · 3f195ea · 3f195ea
1 parent 3cb455c
commit 3f195ea
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/src/factories/ForgetfulNodeCache.ts b/src/factories/ForgetfulNodeCache.ts
@@ -1,5 +1,6 @@
 /* barrel:ignore */
 import * as ts from "typescript";
+import * as errors from "./../errors";
 import {Node} from "./../compiler";
 import {KeyValueCache, createHashSet, HashSet} from "./../utils";
 
@@ -29,6 +30,9 @@ export class ForgetfulNodeCache extends KeyValueCache<ts.Node, Node> {
     }
 
     rememberNode(node: Node) {
+        if (node.wasForgotten())
+            throw new errors.InvalidOperationError("Cannot remember a node that was removed or forgotten.");
+
         let wasInForgetStack = false;
         for (const stackItem of this.forgetStack) {
             if (stackItem.delete(node)) {
@@ -51,7 +55,7 @@ export class ForgetfulNodeCache extends KeyValueCache<ts.Node, Node> {
 
     private forgetNodes(nodes: IterableIterator<Node>) {
         for (const node of nodes) {
-            if (node.getKind() === ts.SyntaxKind.SourceFile)
+            if (node.wasForgotten() || node.getKind() === ts.SyntaxKind.SourceFile)
                 continue;
             node.forgetOnlyThis();
         }

diff --git a/src/tests/tsSimpleAstTests.ts b/src/tests/tsSimpleAstTests.ts
@@ -460,6 +460,7 @@ describe(nameof(TsSimpleAst), () => {
         let interfaceNode2: Node;
         let interfaceNode3: Node;
         let interfaceNode4: Node;
+        let interfaceNode5: Node;
         ast.forgetNodesCreatedInBlock(remember => {
             sourceFile = ast.createSourceFile("test.ts", "class MyClass {} namespace MyNamespace { interface Interface1 {} interface Interface2 {} " +
                 "interface Interface3 {} interface Interface4 {} }");
@@ -471,6 +472,7 @@ describe(nameof(TsSimpleAst), () => {
                 interfaceNode2 = namespaceNode.getInterfaceOrThrow("Interface2");
                 interfaceNode3 = namespaceNode.getInterfaceOrThrow("Interface3");
                 interfaceNode4 = namespaceNode.getInterfaceOrThrow("Interface4");
+                interfaceNode5 = namespaceNode.addInterface({ name: "Interface5" });
                 remember2(interfaceNode3, interfaceNode4);
             });
 
@@ -511,9 +513,30 @@ describe(nameof(TsSimpleAst), () => {
             expect(interfaceNode3.wasForgotten()).to.be.false;
         });
 
-        it("should not have forgotten the third interface because it was remembered", () => {
+        it("should not have forgotten the fourth interface because it was remembered", () => {
             expect(interfaceNode4.wasForgotten()).to.be.false;
         });
+
+        it("should have forgotten the created fifth interface because it was not remembered", () => {
+            expect(interfaceNode5.wasForgotten()).to.be.true;
+        });
+
+        it("should not throw if removing a created node in a block", () => {
+            const sourceFile = ast.createSourceFile("file3.ts", "class MyClass {}");
+            ast.forgetNodesCreatedInBlock(remember => {
+                const classDec = sourceFile.getClassOrThrow("MyClass");
+                classDec.remove();
+            });
+        });
+
+        it("should throw if attempting to remember a node that was forgotten", () => {
+            const sourceFile = ast.createSourceFile("file4.ts");
+            ast.forgetNodesCreatedInBlock(remember => {
+                const classDec = sourceFile.addClass({ name: "Class" });
+                classDec.forget();
+                expect(() => remember(classDec)).to.throw(errors.InvalidOperationError);
+            });
+        });
     });
 
     describe("manipulating then getting something from the type checker", () => {