feat: #176 - Support transformations using the compiler API (`Node#tr…

…ansform(...)`)

docs/_layouts/default.html

@@ -74,6 +74,7 @@ <h1 onclick="document.location.href = '{{ "/" | prepend: site.baseurl }}'" class
                                     <li class="{% if page.path == 'manipulation/order.md' %}active{% endif %}"><a href="{{ "/manipulation/order" | prepend: site.baseurl }}">Order</a></li>
                                     <li class="{% if page.path == 'manipulation/code-writer.md' %}active{% endif %}"><a href="{{ "/manipulation/code-writer" | prepend: site.baseurl }}">Code Writer</a></li>
                                     <li class="{% if page.path == 'manipulation/performance.md' %}active{% endif %}"><a href="{{ "/manipulation/performance" | prepend: site.baseurl }}">Performance</a></li>
+                                    <li class="{% if page.path == 'manipulation/transforms.md' %}active{% endif %}"><a href="{{ "/manipulation/transforms" | prepend: site.baseurl }}">Transforms</a></li>
                                 </ul>
                                 {% endif %}
                             </li>

docs/manipulation/transforms.md

@@ -0,0 +1,71 @@
+---
+title: Transforms
+---
+
+## Transforms
+
+It is possible to transform the AST using the compiler API, though this is not a typical scenario.
+
+For example:
+
+```ts
+import * as ts from "typescript";
+
+const project = new Project();
+const sourceFile = project.createSourceFile("Example.ts", "1; 2; 3;");
+
+// this can be done starting on any node and not just the root node
+sourceFile.transform(traversal => {
+    const node = traversal.visitChildren(); // return type is `ts.Node`
+
+    if (ts.isNumericLiteral(node)) {
+        const incrementedValue = parseInt(node.text, 10) + 1;
+        return ts.createNumericLiteral(incrementedValue.toString());
+    }
+
+    return node;
+});
+
+// outputs: 2; 3; 4;
+console.log(sourceFile.getFullText());
+```
+
+Doing this is more performant, but you won't have type checking, symbols, and you'll be dealing directly with the TypeScript compiler API nodes. Additionally, all previously wrapped descendant nodes of transformed nodes will be forgotten (using them will result in an error being thrown).
+
+### Conditionally visiting children
+
+```ts
+import * as ts from "typescript";
+
+const project = new Project();
+const sourceFile = project.createSourceFile("Example.ts", `
+class C1 {
+    myMethod() {
+        function nestedFunction() {
+        }
+    }
+}
+
+class C2 {
+    prop1: string;
+}
+
+function f1() {
+    console.log("1");
+
+    function nestedFunction() {
+    }
+}`);
+
+sourceFile.transform(traversal => {
+    // this will skip visiting the children of the classes
+    if (ts.isClassDeclaration(traversal.currentNode))
+        return traversal.currentNode;
+
+    const node = traversal.visitChildren();
+    if (ts.isFunctionDeclaration(node))
+        return ts.updateFunctionDeclaration(node, [], [], undefined, ts.createIdentifier("newName"),
+            [], [], undefined, ts.createBlock([]))
+    return node;
+});
+```

lib/ts-simple-ast.d.ts

@@ -4139,6 +4139,18 @@ export interface ForEachDescendantTraversalControl extends ForEachChildTraversal
     up(): void;
 }
 
+export interface TransformTraversalControl {
+    /**
+     * The node currently being transformed.
+     * @remarks Use the result of `.visitChildren()` instead before transforming if visiting the children.
+     */
+    currentNode: ts.Node;
+    /**
+     * Visits the children of the current node and returns a new node for the current node.
+     */
+    visitChildren(): ts.Node;
+}
+
 export declare type NodePropertyToWrappedType<NodeType extends ts.Node, KeyName extends keyof NodeType, NonNullableNodeType = NonNullable<NodeType[KeyName]>> = NodeType[KeyName] extends ts.NodeArray<infer ArrayNodeTypeForNullable> | undefined ? CompilerNodeToWrappedType<ArrayNodeTypeForNullable>[] | undefined : NodeType[KeyName] extends ts.NodeArray<infer ArrayNodeType> ? CompilerNodeToWrappedType<ArrayNodeType>[] : NodeType[KeyName] extends ts.Node ? CompilerNodeToWrappedType<NodeType[KeyName]> : NonNullableNodeType extends ts.Node ? CompilerNodeToWrappedType<NonNullableNodeType> | undefined : NodeType[KeyName];
 
 export declare type NodeParentType<NodeType extends ts.Node> = NodeType extends ts.SourceFile ? CompilerNodeToWrappedType<NodeType["parent"]> | undefined : ts.Node extends NodeType ? CompilerNodeToWrappedType<NodeType["parent"]> | undefined : CompilerNodeToWrappedType<NodeType["parent"]>;
@@ -4505,11 +4517,11 @@ export declare class Node<NodeType extends ts.Node = ts.Node> implements TextRan
     /**
      * Gets the parent if it's a syntax list or throws an error otherwise.
      */
-    getParentSyntaxListOrThrow(): Node<ts.Node>;
+    getParentSyntaxListOrThrow(): SyntaxList;
     /**
      * Gets the parent if it's a syntax list.
      */
-    getParentSyntaxList(): Node | undefined;
+    getParentSyntaxList(): SyntaxList | undefined;
     /**
      * Gets the child index of this node relative to the parent.
      */
@@ -4573,6 +4585,29 @@ export declare class Node<NodeType extends ts.Node = ts.Node> implements TextRan
      * @param settings - Format code settings.
      */
     formatText(settings?: FormatCodeSettings): void;
+    /**
+     * Transforms the node using the compiler api nodes and functions (experimental).
+     *
+     * WARNING: This will forget descendants of transformed nodes.
+     * @example Increments all the numeric literals in a source file.
+     * ```ts
+     * sourceFile.transform(traversal => {
+     *   const node = traversal.visitChildren(); // recommend always visiting the children first (post order)
+     *   if (ts.isNumericLiteral(node))
+     *     return ts.createNumericLiteral((parseInt(node.text, 10) + 1).toString());
+     *   return node;
+     * });
+     * ```
+     * @example Updates the class declaration node without visiting the children.
+     * ```ts
+     * const classDec = sourceFile.getClassOrThrow("MyClass");
+     * classDec.transform(traversal => {
+     *   const node = traversal.currentNode;
+     *   return ts.updateClassDeclaration(node, undefined, undefined, ts.createIdentifier("MyUpdatedClass"), undefined, undefined, []);
+     * });
+     * ```
+     */
+    transform(visitNode: (traversal: TransformTraversalControl) => ts.Node): this;
     /**
      * Gets the leading comment ranges of the current node.
      */

src/compiler/ast/common/Node.ts

@@ -2,7 +2,7 @@ import { CodeBlockWriter } from "../../../codeBlockWriter";
 import * as errors from "../../../errors";
 import { ProjectContext } from "../../../ProjectContext";
 import { getNextMatchingPos, getNextNonWhitespacePos, getPreviousNonWhitespacePos, getPreviousMatchingPos, getTextFromFormattingEdits,
-    insertIntoParentTextRange, replaceSourceFileTextForFormatting } from "../../../manipulation";
+    insertIntoParentTextRange, replaceSourceFileTextForFormatting, replaceSourceFileTextStraight } from "../../../manipulation";
 import { WriterFunction } from "../../../types";
 import { SyntaxKind, ts } from "../../../typescript";
 import { ArrayUtils, getParentSyntaxList, getSyntaxKindName, getTextFromStringOrWriter, isStringKind, printNode, PrintNodeOptions, StringUtils,
@@ -35,6 +35,18 @@ export interface ForEachDescendantTraversalControl extends ForEachChildTraversal
     up(): void;
 }
 
+export interface TransformTraversalControl {
+    /**
+     * The node currently being transformed.
+     * @remarks Use the result of `.visitChildren()` instead before transforming if visiting the children.
+     */
+    currentNode: ts.Node;
+    /**
+     * Visits the children of the current node and returns a new node for the current node.
+     */
+    visitChildren(): ts.Node;
+}
+
 export type NodePropertyToWrappedType<NodeType extends ts.Node, KeyName extends keyof NodeType, NonNullableNodeType = NonNullable<NodeType[KeyName]>> =
     NodeType[KeyName] extends ts.NodeArray<infer ArrayNodeTypeForNullable> | undefined ? CompilerNodeToWrappedType<ArrayNodeTypeForNullable>[] | undefined :
     NodeType[KeyName] extends ts.NodeArray<infer ArrayNodeType> ? CompilerNodeToWrappedType<ArrayNodeType>[] :
@@ -66,6 +78,8 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
     /** @internal */
     private _trailingCommentRanges: CommentRange[] | undefined;
     /** @internal */
+    _wrappedChildCount = 0;
+    /** @internal */
     _sourceFile: SourceFile;
 
     /**
@@ -133,6 +147,14 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
         if (this.wasForgotten())
             return;
 
+        const parent = this.getParent();
+        if (parent != null)
+            parent._wrappedChildCount--;
+
+        const parentSyntaxList = this._getParentSyntaxListIfWrapped();
+        if (parentSyntaxList != null)
+            parentSyntaxList._wrappedChildCount--;
+
         this._storeTextForForgetting();
         this._context.compilerFactory.removeNodeFromCache(this);
         this._clearInternals();
@@ -147,6 +169,14 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
         return this._compilerNode == null;
     }
 
+    /**
+     * Gets if this node has any wrapped children.
+     * @internal
+     */
+    _hasWrappedChildren() {
+        return this._wrappedChildCount > 0;
+    }
+
     /**
      * @internal
      *
@@ -384,21 +414,6 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
         return undefined;
     }
 
-    /**
-     * Offset this node's positions (pos and end) and all of its children by the given offset.
-     * @internal
-     * @param offset - Offset.
-     */
-    _offsetPositions(offset: number) {
-        // todo: remove? Not used internally anymore
-        this.compilerNode.pos += offset;
-        this.compilerNode.end += offset;
-
-        for (const child of this.getChildren()) {
-            child._offsetPositions(offset);
-        }
-    }
-
     /**
      * Gets the previous sibling or throws.
      * @param condition - Optional condition for getting the previous sibling.
@@ -502,7 +517,8 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
      * @internal
      */
     *_getChildrenInCacheIterator(): IterableIterator<Node> {
-        for (const child of this._getCompilerChildren()) {
+        const children = this._hasParsedTokens() ? this._getCompilerChildren() : this._getCompilerForEachChildren();
+        for (const child of children) {
             if (this._context.compilerFactory.hasCompilerNode(child))
                 yield this._context.compilerFactory.getExistingCompilerNode(child)!;
             else if (child.kind === SyntaxKind.SyntaxList) {
@@ -1077,11 +1093,22 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
     /**
      * Gets the parent if it's a syntax list.
      */
-    getParentSyntaxList(): Node | undefined {
+    getParentSyntaxList(): SyntaxList | undefined {
         const syntaxList = getParentSyntaxList(this.compilerNode);
         return this._getNodeFromCompilerNodeIfExists(syntaxList);
     }
 
+    /**
+     * Gets the parent syntax list if it's been wrapped.
+     * @internal
+     */
+    _getParentSyntaxListIfWrapped(): SyntaxList | undefined {
+        const parent = this.getParent();
+        if (parent == null || !parent._hasParsedTokens())
+            return undefined;
+        return this.getParentSyntaxList();
+    }
+
     /**
      * Gets the child index of this node relative to the parent.
      */
@@ -1250,6 +1277,108 @@ export class Node<NodeType extends ts.Node = ts.Node> implements TextRange {
         });
     }
 
+    /**
+     * Transforms the node using the compiler api nodes and functions (experimental).
+     *
+     * WARNING: This will forget descendants of transformed nodes.
+     * @example Increments all the numeric literals in a source file.
+     * ```ts
+     * sourceFile.transform(traversal => {
+     *   const node = traversal.visitChildren(); // recommend always visiting the children first (post order)
+     *   if (ts.isNumericLiteral(node))
+     *     return ts.createNumericLiteral((parseInt(node.text, 10) + 1).toString());
+     *   return node;
+     * });
+     * ```
+     * @example Updates the class declaration node without visiting the children.
+     * ```ts
+     * const classDec = sourceFile.getClassOrThrow("MyClass");
+     * classDec.transform(traversal => {
+     *   const node = traversal.currentNode;
+     *   return ts.updateClassDeclaration(node, undefined, undefined, ts.createIdentifier("MyUpdatedClass"), undefined, undefined, []);
+     * });
+     * ```
+     */
+    transform(visitNode: (traversal: TransformTraversalControl) => ts.Node) {
+        const compilerFactory = this._context.compilerFactory;
+        const printer = ts.createPrinter({
+            newLine: this._context.manipulationSettings.getNewLineKind(),
+            removeComments: false
+        });
+        const transformations: { start: number; end: number; compilerNode: ts.Node; }[] = [];
+        const compilerSourceFile = this._sourceFile.compilerNode;
+        const compilerNode = this.compilerNode;
+        const transformerFactory: ts.TransformerFactory<ts.Node> = context => {
+            return rootNode => innerVisit(rootNode, context);
+        };
+
+        ts.transform(compilerNode, [transformerFactory], this._context.compilerOptions.get());
+
+        replaceSourceFileTextStraight({
+            sourceFile: this._sourceFile,
+            newText: getTransformedText()
+        });
+
+        return this;
+
+        function innerVisit(node: ts.Node, context: ts.TransformationContext) {
+            const traversal: TransformTraversalControl = {
+                visitChildren() {
+                    node = ts.visitEachChild(node, child => innerVisit(child, context), context);
+                    return node;
+                },
+                currentNode: node
+            };
+            const resultNode = visitNode(traversal);
+            handleTransformation(node, resultNode);
+            return resultNode;
+        }
+
+        function handleTransformation(oldNode: ts.Node, newNode: ts.Node) {
+            if (oldNode === newNode)
+                return;
+
+            const start = oldNode.getStart(compilerSourceFile, true);
+            const end = oldNode.end;
+            const lastTransformation = transformations[transformations.length - 1];
+
+            // remove the last transformation if it's nested within this transformation
+            if (lastTransformation != null && lastTransformation.start > start)
+                transformations.pop();
+
+            const wrappedNode = compilerFactory.getExistingCompilerNode(oldNode);
+            transformations.push({
+                start,
+                end,
+                compilerNode: newNode
+            });
+
+            // It's very difficult and expensive to tell about changes that could have happened to the descendants
+            // via updating properties. For this reason, descendant nodes will always be forgotten.
+            if (wrappedNode != null) {
+                if (oldNode.kind !== newNode.kind)
+                    wrappedNode.forget();
+                else
+                    wrappedNode.forgetDescendants();
+            }
+        }
+
+        function getTransformedText() {
+            const fileText = compilerSourceFile.getFullText();
+            let finalText = "";
+            let lastPos = 0;
+
+            for (const transform of transformations) {
+                finalText += fileText.substring(lastPos, transform.start);
+                finalText += printer.printNode(ts.EmitHint.Unspecified, transform.compilerNode, compilerSourceFile);
+                lastPos = transform.end;
+            }
+
+            finalText += fileText.substring(lastPos);
+            return finalText;
+        }
+    }
+
     /**
      * Gets the leading comment ranges of the current node.
      */