Add support for array return values from visitors

src/compiler/transformers/destructuring.ts

@@ -16,7 +16,7 @@ namespace ts {
         node: BinaryExpression,
         needsValue: boolean,
         recordTempVariable: (node: Identifier) => void,
-        visitor?: (node: Node) => Node) {
+        visitor?: (node: Node) => OneOrMany<Node>) {
 
         if (isEmptyObjectLiteralOrArrayLiteral(node.left)) {
             return node.right;
@@ -79,7 +79,7 @@ namespace ts {
      * @param value The rhs value for the binding pattern.
      * @param visitor An optional visitor to use to visit expressions.
      */
-    export function flattenParameterDestructuring(node: ParameterDeclaration, value: Expression, visitor?: (node: Node) => Node) {
+    export function flattenParameterDestructuring(node: ParameterDeclaration, value: Expression, visitor?: (node: Node) => OneOrMany<Node>) {
         const declarations: VariableDeclaration[] = [];
 
         flattenDestructuring(node, value, node, emitAssignment, emitTempVariableAssignment, visitor);
@@ -110,7 +110,7 @@ namespace ts {
      * @param value An optional rhs value for the binding pattern.
      * @param visitor An optional visitor to use to visit expressions.
      */
-    export function flattenVariableDestructuring(node: VariableDeclaration, value?: Expression, visitor?: (node: Node) => Node) {
+    export function flattenVariableDestructuring(node: VariableDeclaration, value?: Expression, visitor?: (node: Node) => OneOrMany<Node>) {
         const declarations: VariableDeclaration[] = [];
 
         flattenDestructuring(node, value, node, emitAssignment, emitTempVariableAssignment, visitor);
@@ -151,7 +151,7 @@ namespace ts {
         node: VariableDeclaration,
         recordTempVariable: (name: Identifier) => void,
         nameSubstitution?: (name: Identifier) => Expression,
-        visitor?: (node: Node) => Node) {
+        visitor?: (node: Node) => OneOrMany<Node>) {
 
         const pendingAssignments: Expression[] = [];
 
@@ -191,7 +191,7 @@ namespace ts {
         location: TextRange,
         emitAssignment: (name: Identifier, value: Expression, location: TextRange, original: Node) => void,
         emitTempVariableAssignment: (value: Expression, location: TextRange) => Identifier,
-        visitor?: (node: Node) => Node) {
+        visitor?: (node: Node) => OneOrMany<Node>) {
         if (value && visitor) {
             value = visitNode(value, visitor, isExpression);
         }

src/compiler/transformers/es6.ts

@@ -55,7 +55,7 @@ namespace ts {
             return visitEachChild(node, visitor, context);
         }
 
-        function visitor(node: Node): Node {
+        function visitor(node: Node): OneOrMany<Node> {
             const savedContainingNonArrowFunction = containingNonArrowFunction;
             const savedCurrentParent = currentParent;
             const savedCurrentNode = currentNode;
@@ -63,17 +63,18 @@ namespace ts {
             const savedEnclosingBlockScopeContainerParent = enclosingBlockScopeContainerParent;
 
             onBeforeVisitNode(node);
-            node = visitorWorker(node);
+
+            const visited = visitorWorker(node);
 
             containingNonArrowFunction = savedContainingNonArrowFunction;
             currentParent = savedCurrentParent;
             currentNode = savedCurrentNode;
             enclosingBlockScopeContainer = savedEnclosingBlockScopeContainer;
             enclosingBlockScopeContainerParent = savedEnclosingBlockScopeContainerParent;
-            return node;
+            return visited;
         }
 
-        function visitorWorker(node: Node): Node {
+        function visitorWorker(node: Node): OneOrMany<Node> {
             if (node.transformFlags & TransformFlags.ES6) {
                 return visitJavaScript(node);
             }
@@ -85,7 +86,7 @@ namespace ts {
             }
         }
 
-        function visitJavaScript(node: Node): Node {
+        function visitJavaScript(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.ClassDeclaration:
                     return visitClassDeclaration(<ClassDeclaration>node);

src/compiler/transformers/es7.ts

@@ -12,7 +12,7 @@ namespace ts {
             return visitEachChild(node, visitor, context);
         }
 
-        function visitor(node: Node): Node {
+        function visitor(node: Node): OneOrMany<Node> {
             if (node.transformFlags & TransformFlags.ES7) {
                 return visitorWorker(node);
             }
@@ -24,7 +24,7 @@ namespace ts {
             }
         }
 
-        function visitorWorker(node: Node) {
+        function visitorWorker(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.BinaryExpression:
                     return visitBinaryExpression(<BinaryExpression>node);

src/compiler/transformers/jsx.ts

@@ -18,7 +18,7 @@ namespace ts {
             return visitEachChild(node, visitor, context);
         }
 
-        function visitor(node: Node): Node {
+        function visitor(node: Node): OneOrMany<Node> {
             if (node.transformFlags & TransformFlags.Jsx) {
                 return visitorWorker(node);
             }
@@ -30,7 +30,7 @@ namespace ts {
             }
         }
 
-        function visitorWorker(node: Node): Node {
+        function visitorWorker(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.JsxElement:
                     return visitJsxElement(<JsxElement>node);

src/compiler/transformers/module/module.ts

@@ -173,7 +173,7 @@ namespace ts {
          *
          * @param node The node.
          */
-        function visitor(node: Node) {
+        function visitor(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.ImportDeclaration:
                     return visitImportDeclaration(<ImportDeclaration>node);

src/compiler/transformers/module/system.ts

@@ -408,7 +408,7 @@ namespace ts {
             return createArrayLiteral(setters);
         }
 
-        function visitSourceElement(node: Node): Node {
+        function visitSourceElement(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.ImportDeclaration:
                     return visitImportDeclaration(<ImportDeclaration>node);
@@ -427,7 +427,7 @@ namespace ts {
             }
         }
 
-        function visitNestedNode(node: Node): Node {
+        function visitNestedNode(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.VariableStatement:
                     return visitVariableStatement(<VariableStatement>node);

src/compiler/transformers/ts.ts

@@ -92,7 +92,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function visitWithStack(node: Node, visitor: (node: Node) => Node): Node {
+        function visitWithStack(node: Node, visitor: (node: Node) => OneOrMany<Node>): OneOrMany<Node> {
             // Save state
             const savedCurrentNamespace = currentNamespace;
             const savedCurrentScope = currentScope;
@@ -102,7 +102,7 @@ namespace ts {
             // Handle state changes before visiting a node.
             onBeforeVisitNode(node);
 
-            node = visitor(node);
+            const visited = visitor(node);
 
             // Restore state
             currentNamespace = savedCurrentNamespace;
@@ -110,7 +110,7 @@ namespace ts {
             currentParent = savedCurrentParent;
             currentNode = savedCurrentNode;
 
-            return node;
+            return visited;
         }
 
         /**
@@ -118,7 +118,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function visitor(node: Node): Node {
+        function visitor(node: Node): OneOrMany<Node> {
             return visitWithStack(node, visitorWorker);
         }
 
@@ -127,14 +127,14 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function visitorWorker(node: Node): Node {
+        function visitorWorker(node: Node): OneOrMany<Node> {
             if (node.transformFlags & TransformFlags.TypeScript) {
                 // This node is explicitly marked as TypeScript, so we should transform the node.
-                node = visitTypeScript(node);
+                return visitTypeScript(node);
             }
             else if (node.transformFlags & TransformFlags.ContainsTypeScript) {
                 // This node contains TypeScript, so we should visit its children.
-                node = visitEachChild(node, visitor, context);
+                return visitEachChild(node, visitor, context);
             }
 
             return node;
@@ -145,7 +145,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function namespaceElementVisitor(node: Node): Node {
+        function namespaceElementVisitor(node: Node): OneOrMany<Node> {
             return visitWithStack(node, namespaceElementVisitorWorker);
         }
 
@@ -154,15 +154,15 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function namespaceElementVisitorWorker(node: Node): Node {
+        function namespaceElementVisitorWorker(node: Node): OneOrMany<Node> {
             if (node.transformFlags & TransformFlags.TypeScript || isExported(node)) {
                 // This node is explicitly marked as TypeScript, or is exported at the namespace
                 // level, so we should transform the node.
-                node = visitTypeScript(node);
+                return visitTypeScript(node);
             }
             else if (node.transformFlags & TransformFlags.ContainsTypeScript) {
                 // This node contains TypeScript, so we should visit its children.
-                node = visitEachChild(node, visitor, context);
+                return visitEachChild(node, visitor, context);
             }
 
             return node;
@@ -173,7 +173,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function classElementVisitor(node: Node) {
+        function classElementVisitor(node: Node): OneOrMany<Node> {
             return visitWithStack(node, classElementVisitorWorker);
         }
 
@@ -182,7 +182,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function classElementVisitorWorker(node: Node) {
+        function classElementVisitorWorker(node: Node): OneOrMany<Node> {
             switch (node.kind) {
                 case SyntaxKind.Constructor:
                     // TypeScript constructors are transformed in `transformClassDeclaration`.
@@ -212,7 +212,7 @@ namespace ts {
          *
          * @param node The node to visit.
          */
-        function visitTypeScript(node: Node): Node {
+        function visitTypeScript(node: Node): OneOrMany<Node> {
             if (hasModifier(node, ModifierFlags.Ambient)) {
                 // TypeScript ambient declarations are elided.
                 return undefined;

src/compiler/visitor.ts

@@ -3,7 +3,7 @@
 
 /* @internal */
 namespace ts {
-    export type OneOrMany<T extends Node> = T | NodeArrayNode<T>;
+    export type OneOrMany<T extends Node> = T | NodeArrayNode<T> | T[];
 
     /**
      * Describes an edge of a Node, used when traversing a syntax tree.
@@ -478,7 +478,7 @@ namespace ts {
      * @param optional An optional value indicating whether the Node is itself optional.
      * @param lift An optional callback to execute to lift a NodeArrayNode into a valid Node.
      */
-    export function visitNode<T extends Node>(node: T, visitor: (node: Node) => Node, test: (node: Node) => boolean, optional?: boolean, lift?: (node: NodeArray<Node>) => T): T {
+    export function visitNode<T extends Node>(node: T, visitor: (node: Node) => OneOrMany<Node>, test: (node: Node) => boolean, optional?: boolean, lift?: (node: NodeArray<Node>) => T): T {
         return <T>visitNodeWorker(node, visitor, test, optional, lift, /*parenthesize*/ undefined, /*parentNode*/ undefined);
     }
 
@@ -493,7 +493,7 @@ namespace ts {
      * @param parenthesize A callback used to parenthesize the node if needed.
      * @param parentNode A parentNode for the node.
      */
-    function visitNodeWorker(node: Node, visitor: (node: Node) => Node, test: (node: Node) => boolean, optional: boolean, lift: (node: NodeArray<Node>) => Node, parenthesize: (node: Node, parentNode: Node) => Node, parentNode: Node): Node {
+    function visitNodeWorker(node: Node, visitor: (node: Node) => OneOrMany<Node>, test: (node: Node) => boolean, optional: boolean, lift: (node: NodeArray<Node>) => Node, parenthesize: (node: Node, parentNode: Node) => Node, parentNode: Node): Node {
         if (node === undefined) {
             return undefined;
         }
@@ -503,22 +503,29 @@ namespace ts {
             return node;
         }
 
-        if (visited !== undefined && isNodeArrayNode(visited)) {
-            visited = (lift || extractSingleNode)((<NodeArrayNode<Node>>visited).nodes);
-        }
-
-        if (parenthesize !== undefined && visited !== undefined) {
-            visited = parenthesize(visited, parentNode);
-        }
-
         if (visited === undefined) {
             Debug.assert(optional, "Node not optional.");
             return undefined;
         }
 
-        Debug.assert(test === undefined || test(visited), "Wrong node type after visit.", () => `Node ${formatSyntaxKind(visited.kind)} did not pass test ${(<any>test).name}.`);
-        aggregateTransformFlags(visited);
-        return visited;
+        let visitedNode: Node;
+        if (isArray(visited)) {
+            visitedNode = (lift || extractSingleNode)(<NodeArray<Node>>visited);
+        }
+        else if (isNodeArrayNode(visited)) {
+            visitedNode = (lift || extractSingleNode)((<NodeArrayNode<Node>>visited).nodes);
+        }
+        else {
+            visitedNode = visited;
+        }
+
+        if (parenthesize !== undefined) {
+            visitedNode = parenthesize(visitedNode, parentNode);
+        }
+
+        Debug.assert(test === undefined || test(visitedNode), "Wrong node type after visit.", () => `Node ${formatSyntaxKind(visitedNode.kind)} did not pass test ${(<any>test).name}.`);
+        aggregateTransformFlags(visitedNode);
+        return visitedNode;
     }
 
     /**
@@ -530,7 +537,7 @@ namespace ts {
      * @param start An optional value indicating the starting offset at which to start visiting.
      * @param count An optional value indicating the maximum number of nodes to visit.
      */
-    export function visitNodes<T extends Node, TArray extends NodeArray<T>>(nodes: TArray, visitor: (node: Node) => Node, test: (node: Node) => boolean, start?: number, count?: number): TArray {
+    export function visitNodes<T extends Node, TArray extends NodeArray<T>>(nodes: TArray, visitor: (node: Node) => OneOrMany<Node>, test: (node: Node) => boolean, start?: number, count?: number): TArray {
         return <TArray>visitNodesWorker(nodes, visitor, test, /*parenthesize*/ undefined, /*parentNode*/ undefined, start, count);
     }
 
@@ -543,7 +550,7 @@ namespace ts {
      * @param start An optional value indicating the starting offset at which to start visiting.
      * @param count An optional value indicating the maximum number of nodes to visit.
      */
-    function visitNodesWorker(nodes: NodeArray<Node>, visitor: (node: Node) => Node, test: (node: Node) => boolean, parenthesize: (node: Node, parentNode: Node) => Node, parentNode: Node, start: number, count: number): NodeArray<Node> {
+    function visitNodesWorker(nodes: NodeArray<Node>, visitor: (node: Node) => OneOrMany<Node>, test: (node: Node) => boolean, parenthesize: (node: Node, parentNode: Node) => Node, parentNode: Node, start: number, count: number): NodeArray<Node> {
         if (nodes === undefined) {
             return undefined;
         }
@@ -593,8 +600,8 @@ namespace ts {
      * @param visitor The callback used to visit each child.
      * @param context A lexical environment context for the visitor.
      */
-    export function visitEachChild<T extends Node>(node: T, visitor: (node: Node) => Node, context: LexicalEnvironment): T;
-    export function visitEachChild<T extends Node>(node: T & Map<any>, visitor: (node: Node) => Node, context: LexicalEnvironment): T {
+    export function visitEachChild<T extends Node>(node: T, visitor: (node: Node) => OneOrMany<Node>, context: LexicalEnvironment): T;
+    export function visitEachChild<T extends Node>(node: T & Map<any>, visitor: (node: Node) => OneOrMany<Node>, context: LexicalEnvironment): T {
         if (node === undefined) {
             return undefined;
         }
@@ -661,7 +668,7 @@ namespace ts {
         if (nodes) {
             for (let i = 0; i < nodes.length; i++) {
                 const node = nodes[i];
-                if (result || node === undefined || isNodeArrayNode(node)) {
+                if (result || node === undefined || isArray(node) || isNodeArrayNode(node)) {
                     if (!result) {
                         result = <T[]>nodes.slice(0, i);
                     }
@@ -695,26 +702,26 @@ namespace ts {
 
     function addNodeWorker(to: Node[], from: OneOrMany<Node>, startOnNewLine: boolean, test: (node: Node) => boolean, parenthesize: (node: Node, parentNode: Node) => Node, parentNode: Node, isVisiting: boolean) {
         if (to && from) {
-            if (isNodeArrayNode(from)) {
+            if (isArray(from)) {
+                addNodesWorker(to, from, startOnNewLine, test, parenthesize, parentNode, isVisiting);
+            }
+            else if (isNodeArrayNode(from)) {
                 addNodesWorker(to, from.nodes, startOnNewLine, test, parenthesize, parentNode, isVisiting);
-                return;
             }
+            else {
+                const node = parenthesize !== undefined ? parenthesize(from, parentNode) : from;
+                Debug.assert(test === undefined || test(node), "Wrong node type after visit.", () => `Node ${formatSyntaxKind(node.kind)} did not pass test ${(<any>test).name}.`);
 
-            if (parenthesize !== undefined) {
-                from = parenthesize(from, parentNode);
+                if (startOnNewLine) {
+                    node.startsOnNewLine = true;
+                }
+
+                if (isVisiting) {
+                    aggregateTransformFlags(node);
+                }
+
+                to.push(node);
             }
-
-            Debug.assert(test === undefined || test(from), "Wrong node type after visit.", () => `Node ${formatSyntaxKind(from.kind)} did not pass test ${(<any>test).name}.`);
-
-            if (startOnNewLine) {
-                from.startsOnNewLine = true;
-            }
-
-            if (isVisiting) {
-                aggregateTransformFlags(from);
-            }
-
-            to.push(from);
         }
     }