From 585ea5d6006c6be38b672b3d65356c044c7bbd90 Mon Sep 17 00:00:00 2001
From: Tobias Bosch <tbosch1009@gmail.com>
Date: Wed, 22 Jul 2015 14:24:57 -0700
Subject: [PATCH] feat(query): allow to query for `TemplateRef`

Part of #1989

Closes #3202
---
 .../src/core/compiler/element_injector.ts     |  4 +++
 .../src/core/compiler/proto_view_factory.ts   |  9 +++--
 .../src/core/compiler/template_ref.ts         |  2 +-
 .../core/compiler/element_injector_spec.ts    | 16 +++++++++
 .../core/compiler/query_integration_spec.ts   | 36 +++++++++++++++++--
 5 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/modules/angular2/src/core/compiler/element_injector.ts b/modules/angular2/src/core/compiler/element_injector.ts
index b01a3db79f..eb8f82dce1 100644
--- a/modules/angular2/src/core/compiler/element_injector.ts
+++ b/modules/angular2/src/core/compiler/element_injector.ts
@@ -714,6 +714,10 @@ export class ElementInjector extends TreeNode<ElementInjector> implements Depend
   }
 
   addDirectivesMatchingQuery(query: Query, list: any[]): void {
+    var templateRef = this._preBuiltObjects.templateRef;
+    if (query.selector === TemplateRef && isPresent(templateRef)) {
+      list.push(templateRef);
+    }
     this._strategy.addDirectivesMatchingQuery(query, list);
   }
 
diff --git a/modules/angular2/src/core/compiler/proto_view_factory.ts b/modules/angular2/src/core/compiler/proto_view_factory.ts
index 56b912ac06..a510865afe 100644
--- a/modules/angular2/src/core/compiler/proto_view_factory.ts
+++ b/modules/angular2/src/core/compiler/proto_view_factory.ts
@@ -366,10 +366,13 @@ function _createProtoElementInjector(binderIndex, parentPeiWithDistance, renderE
                                      componentDirectiveBinding, directiveBindings) {
   var protoElementInjector = null;
   // Create a protoElementInjector for any element that either has bindings *or* has one
-  // or more var- defined. Elements with a var- defined need a their own element injector
-  // so that, when hydrating, $implicit can be set to the element.
+  // or more var- defined *or* for <template> elements:
+  // - Elements with a var- defined need a their own element injector
+  //   so that, when hydrating, $implicit can be set to the element.
+  // - <template> elements need their own ElementInjector so that we can query their TemplateRef
   var hasVariables = MapWrapper.size(renderElementBinder.variableBindings) > 0;
-  if (directiveBindings.length > 0 || hasVariables) {
+  if (directiveBindings.length > 0 || hasVariables ||
+      isPresent(renderElementBinder.nestedProtoView)) {
     var directiveVariableBindings =
         createDirectiveVariableBindings(renderElementBinder, directiveBindings);
     protoElementInjector =
diff --git a/modules/angular2/src/core/compiler/template_ref.ts b/modules/angular2/src/core/compiler/template_ref.ts
index 0491791afb..e01c601b63 100644
--- a/modules/angular2/src/core/compiler/template_ref.ts
+++ b/modules/angular2/src/core/compiler/template_ref.ts
@@ -28,5 +28,5 @@ export class TemplateRef {
   /**
    * Whether this template has a local variable with the given name
    */
-  hasLocal(name: string): boolean { return this._getProtoView().protoLocals.has(name); }
+  hasLocal(name: string): boolean { return this._getProtoView().variableBindings.has(name); }
 }
diff --git a/modules/angular2/test/core/compiler/element_injector_spec.ts b/modules/angular2/test/core/compiler/element_injector_spec.ts
index 1e582629e2..4aede997fc 100644
--- a/modules/angular2/test/core/compiler/element_injector_spec.ts
+++ b/modules/angular2/test/core/compiler/element_injector_spec.ts
@@ -166,6 +166,12 @@ class NeedsQueryByVarBindings {
   constructor(@Query("one,two") query: QueryList<any>) { this.query = query; }
 }
 
+@Injectable()
+class NeedsTemplateRefQuery {
+  query: QueryList<TemplateRef>;
+  constructor(@Query(TemplateRef) query: QueryList<TemplateRef>) { this.query = query; }
+}
+
 @Injectable()
 class NeedsElementRef {
   elementRef;
@@ -1009,6 +1015,16 @@ export function main() {
             expectDirectives(inj.get(NeedsQuery).query, CountingDirective, [0]);
           })
 
+          it('should contain PreBuiltObjects on the same injector', () => {
+            var preBuiltObjects = new PreBuiltObjects(null, null, null, new TemplateRef(<any>new DummyElementRef()));
+            var inj = injector(ListWrapper.concat([
+                NeedsTemplateRefQuery
+              ], extraBindings), null,
+              false, preBuiltObjects);
+
+            expect(inj.get(NeedsTemplateRefQuery).query.first).toBe(preBuiltObjects.templateRef);
+          });
+
           it('should contain multiple directives from the same injector', () => {
             var inj = injector(ListWrapper.concat([
                 NeedsQuery,
diff --git a/modules/angular2/test/core/compiler/query_integration_spec.ts b/modules/angular2/test/core/compiler/query_integration_spec.ts
index 40ecbe85f3..3b2c116c9c 100644
--- a/modules/angular2/test/core/compiler/query_integration_spec.ts
+++ b/modules/angular2/test/core/compiler/query_integration_spec.ts
@@ -10,12 +10,13 @@ import {
   it,
   xit,
   TestComponentBuilder,
-  asNativeElements
+  asNativeElements,
+  By
 } from 'angular2/test_lib';
 
 
 import {Injectable, Optional} from 'angular2/di';
-import {QueryList} from 'angular2/core';
+import {QueryList, TemplateRef} from 'angular2/core';
 import {Query, ViewQuery, Component, Directive, View} from 'angular2/annotations';
 
 import {NgIf, NgFor} from 'angular2/angular2';
@@ -126,6 +127,24 @@ export function main() {
          }));
     });
 
+    describe('query for TemplateRef', () => {
+      it('should find TemplateRefs in the light and shadow dom',
+         inject([TestComponentBuilder, AsyncTestCompleter], (tcb: TestComponentBuilder, async) => {
+           var template = '<needs-tpl><template var-x="light"></template></needs-tpl>';
+           tcb.overrideTemplate(MyComp, template)
+               .createAsync(MyComp)
+               .then((view) => {
+                 view.detectChanges();
+                 var needsTpl: NeedsTpl = view.componentViewChildren[0].inject(NeedsTpl);
+                 expect(needsTpl.query.first.hasLocal('light')).toBe(true);
+                 expect(needsTpl.viewQuery.first.hasLocal('shadow')).toBe(true);
+
+                 async.done();
+               });
+         }));
+
+    });
+
     describe("onChange", () => {
       it('should notify query on change',
          inject([TestComponentBuilder, AsyncTestCompleter], (tcb: TestComponentBuilder, async) => {
@@ -545,6 +564,18 @@ class NeedsViewQueryOrder {
   constructor(@ViewQuery(TextDirective) query: QueryList<TextDirective>) { this.query = query; }
 }
 
+@Component({selector: 'needs-tpl'})
+@View({template: '<template var-x="shadow"></template>'})
+class NeedsTpl {
+  viewQuery: QueryList<TemplateRef>;
+  query: QueryList<TemplateRef>;
+  constructor(@ViewQuery(TemplateRef) viewQuery: QueryList<TemplateRef>,
+              @Query(TemplateRef) query: QueryList<TemplateRef>) {
+    this.viewQuery = viewQuery;
+    this.query = query;
+  }
+}
+
 @Component({selector: 'my-comp'})
 @View({
   directives: [
@@ -558,6 +589,7 @@ class NeedsViewQueryOrder {
     NeedsViewQueryIf,
     NeedsViewQueryNestedIf,
     NeedsViewQueryOrder,
+    NeedsTpl,
     TextDirective,
     InertDirective,
     NgIf,