Ensure template given to ShadyCSS is the one rendered by lit-html (#502)

This ensures ShadyCSS can update the style in the template that will be rendered by lit-html. ShadyCSS needs to do this, for example, when the set of properties changes in any `@apply`'s.

src/lib/shady-render.ts

@@ -107,53 +107,54 @@ const shadyRenderSet = new Set<string>();
  * not be scoped and the <style> will be left in the template and rendered
  * output.
  */
-const ensureStylesScoped =
+const styleTemplatesForScope =
     (fragment: DocumentFragment, template: Template, scopeName: string) => {
-      // only scope element template once per scope name
-      if (!shadyRenderSet.has(scopeName)) {
-        shadyRenderSet.add(scopeName);
-        const styleTemplate = document.createElement('template');
-        Array.from(fragment.querySelectorAll('style')).forEach((s: Element) => {
-          styleTemplate.content.appendChild(s);
-        });
-        window.ShadyCSS.prepareTemplateStyles(styleTemplate, scopeName);
-        // Fix templates: note the expectation here is that the given `fragment`
-        // has been generated from the given `template` which contains
-        // the set of templates rendered into this scope.
-        // It is only from this set of initial templates from which styles
-        // will be scoped and removed.
-        removeStylesFromLitTemplates(scopeName);
-        // ApplyShim case
-        if (window.ShadyCSS.nativeShadow) {
-          const style = styleTemplate.content.querySelector('style');
-          if (style !== null) {
-            // Insert style into rendered fragment
-            fragment.insertBefore(style, fragment.firstChild);
-            // Insert into lit-template (for subsequent renders)
-            insertNodeIntoTemplate(
-                template,
-                style.cloneNode(true),
-                template.element.content.firstChild);
-          }
+      shadyRenderSet.add(scopeName);
+      // Move styles out of rendered DOM and store.
+      const styles = fragment.querySelectorAll('style');
+      const styleFragment = document.createDocumentFragment();
+      for (let i = 0; i < styles.length; i++) {
+        styleFragment.appendChild(styles[i]);
+      }
+      // Remove styles from nested templates in this scope.
+      removeStylesFromLitTemplates(scopeName);
+      // And then put them into the "root" template passed in as `template`.
+      insertNodeIntoTemplate(
+          template, styleFragment, template.element.content.firstChild);
+      // Note, it's important that ShadyCSS gets the template that `lit-html`
+      // will actually render so that it can update the style inside when
+      // needed.
+      window.ShadyCSS.prepareTemplateStyles(template.element, scopeName);
+      // When using native Shadow DOM, replace the style in the rendered
+      // fragment.
+      if (window.ShadyCSS.nativeShadow) {
+        const style = template.element.content.querySelector('style');
+        if (style !== null) {
+          fragment.insertBefore(style.cloneNode(true), fragment.firstChild);
         }
       }
     };
 
-// NOTE: We're copying code from lit-html's `render` method here.
-// We're doing this explicitly because the API for rendering templates is likely
-// to change in the near term.
 export function render(
     result: TemplateResult,
     container: Element|DocumentFragment,
     scopeName: string) {
+  const shouldScope =
+      container instanceof ShadowRoot && compatibleShadyCSSVersion;
+  const hasScoped = shadyRenderSet.has(scopeName);
+  // Call `styleElement` to update element if it's already been processed.
+  // This ensures the template is up to date before stamping the template
+  // into the shadowRoot *or* updates the shadowRoot if we've already stamped
+  // and are just updating the template.
+  if (shouldScope && hasScoped) {
+    window.ShadyCSS.styleElement((container as ShadowRoot).host);
+  }
   litRender(result, container, shadyTemplateFactory(scopeName));
-
-  // If there's a shadow host, do ShadyCSS scoping...
-  if (container instanceof ShadowRoot && result instanceof TemplateResult &&
-      compatibleShadyCSSVersion) {
+  // When rendering a TemplateResult, scope the template with ShadyCSS
+  if (shouldScope && !hasScoped && result instanceof TemplateResult) {
     const part = parts.get(container)!;
     const instance = part.value as TemplateInstance;
-    ensureStylesScoped(container, instance.template, scopeName);
-    window.ShadyCSS.styleElement(container.host);
+    styleTemplatesForScope(
+        (container as ShadowRoot), instance.template, scopeName);
   }
 }

src/test/lib/shady-render-apply_test.ts

@@ -12,7 +12,10 @@
  * http://polymer.github.io/PATENTS.txt
  */
 
-import {html, render} from '../../lib/shady-render.js';
+// Rename the html tag so that CSS linting doesn't warn on the non-standard
+// @apply syntax
+import {until} from '../../directives/until.js';
+import {html as htmlWithApply, render} from '../../lib/shady-render.js';
 
 const assert = chai.assert;
 
@@ -21,9 +24,6 @@ suite('shady-render @apply', () => {
     const container = document.createElement('scope-5');
     document.body.appendChild(container);
     container.attachShadow({mode: 'open'});
-    // Rename the html tag so that CSS linting doesn't warn on the non-standard
-    // @apply syntax
-    const htmlWithApply = html;
     const result = htmlWithApply`
       <style>
         :host {
@@ -46,4 +46,65 @@ suite('shady-render @apply', () => {
     assert.equal(computedStyle.getPropertyValue('padding-top').trim(), '4px');
     document.body.removeChild(container);
   });
+
+  test(
+      'styles with css custom properties using @apply render in different contexts',
+      async () => {
+        const createApplyUser = () => {
+          const container = document.createElement('apply-user');
+          container.attachShadow({mode: 'open'});
+          const result = htmlWithApply`
+        <style>
+          div {
+            border-top: 2px solid black;
+            margin-top: 4px;
+            @apply --stuff;
+          }
+        </style>
+        <div>Testing...</div>
+      `;
+          render(result, container.shadowRoot!, 'apply-user');
+          return container;
+        };
+        const applyUser = createApplyUser();
+        document.body.appendChild(applyUser);
+        const applyUserDiv = (applyUser.shadowRoot!).querySelector('div');
+        const applyUserStyle = getComputedStyle(applyUserDiv!);
+        assert.equal(
+            applyUserStyle.getPropertyValue('border-top-width').trim(), '2px');
+        assert.equal(
+            applyUserStyle.getPropertyValue('margin-top').trim(), '4px');
+        // Render sub-element with a promise to ensure it's rendered after the
+        // containing scope.
+        const applyUserPromise = Promise.resolve().then(createApplyUser);
+        const producerResult = htmlWithApply`
+      <style>
+        :host {
+          --stuff: {
+            border-top: 10px solid orange;
+            padding-top: 20px;
+          };
+        }
+      </style>
+      ${until(applyUserPromise, 'loading')}
+    `;
+        const applyProducer = document.createElement('apply-producer');
+        applyProducer.attachShadow({mode: 'open'});
+        document.body.appendChild(applyProducer);
+        render(producerResult, applyProducer.shadowRoot!, 'apply-producer');
+        await applyUserPromise;
+        const applyProducerDiv =
+            applyProducer.shadowRoot!.querySelector('apply-user')!.shadowRoot!
+                .querySelector('div')!;
+        const applyProducerStyle = getComputedStyle(applyProducerDiv!);
+        assert.equal(
+            applyProducerStyle.getPropertyValue('border-top-width').trim(),
+            '10px');
+        assert.equal(
+            applyUserStyle.getPropertyValue('margin-top').trim(), '4px');
+        assert.equal(
+            applyProducerStyle.getPropertyValue('padding-top').trim(), '20px');
+        document.body.removeChild(applyUser);
+        document.body.removeChild(applyProducer);
+      });
 });

src/test/lib/shady-render_test.ts

@@ -12,6 +12,7 @@
  * http://polymer.github.io/PATENTS.txt
  */
 
+import {until} from '../../directives/until.js';
 import {html, render} from '../../lib/shady-render.js';
 
 const assert = chai.assert;
@@ -98,6 +99,46 @@ suite('shady-render', () => {
     document.body.removeChild(container);
   });
 
+  test(
+      'styles with css custom properties flow to nested shadowRoots',
+      async () => {
+        // promise for sub element
+        const elementPromise = Promise.resolve().then(() => {
+          const container = document.createElement('scope-4a-sub');
+          container.attachShadow({mode: 'open'});
+          const result = html`
+        <style>
+          :host {
+            display: block;
+            border: var(--border);
+          }
+        </style>
+        <div>Testing...</div>
+      `;
+          render(result, container.shadowRoot!, 'scope-4a-sub');
+          return container;
+        });
+
+        const container = document.createElement('scope-4a');
+        document.body.appendChild(container);
+        container.attachShadow({mode: 'open'});
+        const result = html`
+      <style>
+        :host {
+          --border: 2px solid orange;
+        }
+      </style>
+      ${until(elementPromise, '')}
+    `;
+        render(result, container.shadowRoot!, 'scope-4a');
+        await elementPromise;
+        const e = (container.shadowRoot!).querySelector('scope-4a-sub');
+        assert.equal(
+            getComputedStyle(e!).getPropertyValue('border-top-width').trim(),
+            '2px');
+        document.body.removeChild(container);
+      });
+
   test('parts around styles with parts render/update', () => {
     const container = document.createElement('div');
     document.body.appendChild(container);