From 4202936bbf18e1a7e937ce6de077e487c63d8c19 Mon Sep 17 00:00:00 2001
From: Chuck Jazdzewski <chuckj@google.com>
Date: Mon, 10 Oct 2016 09:13:50 -0700
Subject: [PATCH] refactor(compiler): add optional `visit()` to html AST
 `Visitor` (#12135)

---
 .../@angular/compiler/src/ml_parser/ast.ts    | 10 ++-
 .../test/ml_parser/html_parser_spec.ts        | 64 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/modules/@angular/compiler/src/ml_parser/ast.ts b/modules/@angular/compiler/src/ml_parser/ast.ts
index ddd342f5ef..d15e1fa042 100644
--- a/modules/@angular/compiler/src/ml_parser/ast.ts
+++ b/modules/@angular/compiler/src/ml_parser/ast.ts
@@ -54,6 +54,10 @@ export class Comment implements Node {
 }
 
 export interface Visitor {
+  // Returning a truthy value from `visit()` will prevent `visitAll()` from the call to the typed
+  // method and result returned will become the result included in `visitAll()`s result array.
+  visit?(node: Node, context: any): any;
+
   visitElement(element: Element, context: any): any;
   visitAttribute(attribute: Attribute, context: any): any;
   visitText(text: Text, context: any): any;
@@ -64,8 +68,12 @@ export interface Visitor {
 
 export function visitAll(visitor: Visitor, nodes: Node[], context: any = null): any[] {
   let result: any[] = [];
+
+  let visit = visitor.visit ?
+      (ast: Node) => visitor.visit(ast, context) || ast.visit(visitor, context) :
+      (ast: Node) => ast.visit(visitor, context);
   nodes.forEach(ast => {
-    const astResult = ast.visit(visitor, context);
+    const astResult = visit(ast);
     if (astResult) {
       result.push(astResult);
     }
diff --git a/modules/@angular/compiler/test/ml_parser/html_parser_spec.ts b/modules/@angular/compiler/test/ml_parser/html_parser_spec.ts
index d510a892bf..f70a87a1e0 100644
--- a/modules/@angular/compiler/test/ml_parser/html_parser_spec.ts
+++ b/modules/@angular/compiler/test/ml_parser/html_parser_spec.ts
@@ -390,6 +390,70 @@ export function main() {
         });
       });
 
+      describe('visitor', () => {
+        it('should visit text nodes', () => {
+          const result = humanizeDom(parser.parse('text', 'TestComp'));
+          expect(result).toEqual([[html.Text, 'text', 0]]);
+        });
+
+        it('should visit element nodes', () => {
+          const result = humanizeDom(parser.parse('<div></div>', 'TestComp'));
+          expect(result).toEqual([[html.Element, 'div', 0]]);
+        });
+
+        it('should visit attribute nodes', () => {
+          const result = humanizeDom(parser.parse('<div id="foo"></div>', 'TestComp'));
+          expect(result).toContain([html.Attribute, 'id', 'foo']);
+        });
+
+        it('should visit all nodes', () => {
+          const result =
+              parser.parse('<div id="foo"><span id="bar">a</span><span>b</span></div>', 'TestComp');
+          const accumulator: html.Node[] = [];
+          const visitor = new class {
+            visit(node: html.Node, context: any) { accumulator.push(node); }
+            visitElement(element: html.Element, context: any): any {
+              html.visitAll(this, element.attrs);
+              html.visitAll(this, element.children);
+            }
+            visitAttribute(attribute: html.Attribute, context: any): any {}
+            visitText(text: html.Text, context: any): any {}
+            visitComment(comment: html.Comment, context: any): any {}
+            visitExpansion(expansion: html.Expansion, context: any): any {
+              html.visitAll(this, expansion.cases);
+            }
+            visitExpansionCase(expansionCase: html.ExpansionCase, context: any): any {}
+          };
+
+          html.visitAll(visitor, result.rootNodes);
+          expect(accumulator.map(n => n.constructor)).toEqual([
+            html.Element, html.Attribute, html.Element, html.Attribute, html.Text, html.Element,
+            html.Text
+          ]);
+        });
+
+        it('should skip typed visit if visit() returns a truthy value', () => {
+          const visitor = new class {
+            visit(node: html.Node, context: any) { return true; }
+            visitElement(element: html.Element, context: any): any { throw Error('Unexpected'); }
+            visitAttribute(attribute: html.Attribute, context: any): any {
+              throw Error('Unexpected');
+            }
+            visitText(text: html.Text, context: any): any { throw Error('Unexpected'); }
+            visitComment(comment: html.Comment, context: any): any { throw Error('Unexpected'); }
+            visitExpansion(expansion: html.Expansion, context: any): any {
+              throw Error('Unexpected');
+            }
+            visitExpansionCase(expansionCase: html.ExpansionCase, context: any): any {
+              throw Error('Unexpected');
+            }
+          };
+          const result = parser.parse('<div id="foo"></div><div id="bar"></div>', 'TestComp');
+          const traversal = html.visitAll(visitor, result.rootNodes);
+          expect(traversal).toEqual([true, true]);
+        });
+      });
+
       describe('errors', () => {
         it('should report unexpected closing tags', () => {
           let errors = parser.parse('<div></p></div>', 'TestComp').errors;