Adds hydrationPolyfills config for renderers (#594)

Some renderers, such as Lit, need special polyfills only for hydration. We have the `polyfills` array, but that is intended for polyfills that always need to run. This adds a second type hydrationPolyfills that only run on elements that are `:load`, `:idle`, etc.

7 files changed, 31 insertions, 9 deletions

diff --git a/docs/core-concepts/ui-renderers.md b/docs/core-concepts/ui-renderers.md
index a83d3f49b..2d91f1811 100644
--- a/docs/core-concepts/ui-renderers.md
+++ b/docs/core-concepts/ui-renderers.md
@@ -106,6 +106,7 @@ export default {
   knownEntrypoint: ['framework'], // optional, entrypoint modules that will be used by compiled source
   external: ['dep'] // optional, dependencies that should not be built by snowpack
   polyfills: ['./shadow-dom-polyfill.js'] // optional, module scripts that should be loaded before client hydration.
+  hydrationPolyfills: ['./hydrate-framework.js'] // optional, polyfills that need to run before hydration ever occurs.
 };
 ```
 
diff --git a/packages/astro/src/config_manager.ts b/packages/astro/src/config_manager.ts
index e4a4fa8e1..b349cb8c7 100644
--- a/packages/astro/src/config_manager.ts
+++ b/packages/astro/src/config_manager.ts
@@ -16,6 +16,7 @@ interface RendererInstance {
   knownEntrypoints: string[] | undefined;
   external: string[] | undefined;
   polyfills: string[];
+  hydrationPolyfills: string[];
 }
 
 const CONFIG_MODULE_BASE_NAME = '__astro_config.js';
@@ -106,6 +107,7 @@ export class ConfigManager {
       }
 
       const polyfillsNormalized = (raw.polyfills || []).map((p: string) => p.startsWith('.') ? path.join(name, p) : p);
+      const hydrationPolyfillsNormalized = (raw.hydrationPolyfills || []).map((p: string) => p.startsWith('.') ? path.join(name, p) : p);
 
       return {
         name,
@@ -115,7 +117,8 @@ export class ConfigManager {
         server: path.join(name, raw.server),
         knownEntrypoints: raw.knownEntrypoints,
         external: raw.external,
-        polyfills: polyfillsNormalized
+        polyfills: polyfillsNormalized,
+        hydrationPolyfills: hydrationPolyfillsNormalized
       };
     });
 
@@ -127,7 +130,7 @@ export class ConfigManager {
     const rendererServerPackages = renderers.map(({ server }) => server);
     const rendererClientPackages = await Promise.all(renderers.filter(({client}) => client).map(({ client }) => this.resolvePackageUrl(client!)));
     const rendererPolyfills = await Promise.all(renderers.map(({ polyfills }) => Promise.all(polyfills.map(src => this.resolvePackageUrl(src)))));
-
+    const rendererHydrationPolyfills = await Promise.all(renderers.map(({ hydrationPolyfills }) => Promise.all(hydrationPolyfills.map(src => this.resolvePackageUrl(src)))));
 
     const result = /* js */ `${rendererServerPackages.map((pkg, i) => `import __renderer_${i} from "${pkg}";`).join('\n')}
 
@@ -137,7 +140,8 @@ let rendererInstances = [${renderers.map((r, i) => `{
   source: ${rendererClientPackages[i] ? `"${rendererClientPackages[i]}"` : 'null'},
   renderer: __renderer_${i},
   options: ${r.options ? JSON.stringify(r.options) : 'null'},
-  polyfills: ${JSON.stringify(rendererPolyfills[i])}
+  polyfills: ${JSON.stringify(rendererPolyfills[i])},
+  hydrationPolyfills: ${JSON.stringify(rendererHydrationPolyfills[i])}
 }`).join(', ')}];
 
 ${contents}
diff --git a/packages/astro/src/internal/__astro_component.ts b/packages/astro/src/internal/__astro_component.ts
index 2e98d55dc..d57deb462 100644
--- a/packages/astro/src/internal/__astro_component.ts
+++ b/packages/astro/src/internal/__astro_component.ts
@@ -14,20 +14,23 @@ export interface RendererInstance {
   renderer: Renderer;
   options: any;
   polyfills: string[];
+  hydrationPolyfills: string[];
 }
 
 const astroRendererInstance: RendererInstance = {
   source: '',
   renderer: astro as Renderer,
   options: null,
-  polyfills: []
+  polyfills: [],
+  hydrationPolyfills: []
 };
 
 const astroHtmlRendererInstance: RendererInstance = {
   source: '',
   renderer: astroHtml as Renderer,
   options: null,
-  polyfills: []
+  polyfills: [],
+  hydrationPolyfills: []
 };
 
 let rendererInstances: RendererInstance[] = [];
@@ -90,13 +93,18 @@ interface HydrateScriptOptions {
 async function generateHydrateScript({ instance, astroId, props }: HydrateScriptOptions, { hydrate, componentUrl, componentExport }: Required<AstroComponentProps>) {
   const { source } = instance;
 
-  const hydrationSource = source ? `
+  let hydrationSource = '';
+  if(instance.hydrationPolyfills.length) {
+    hydrationSource += `await Promise.all([${instance.hydrationPolyfills.map(src => `import("${src}")`).join(', ')}]);\n`;
+  }
+
+  hydrationSource += source ? `
   const [{ ${componentExport.value}: Component }, { default: hydrate }] = await Promise.all([import("${componentUrl}"), import("${source}")]);
   return (el, children) => hydrate(el)(Component, ${serialize(props)}, children);
-`.trim() : `
+` : `
   await import("${componentUrl}");
   return () => {};
-`.trim()
+`;
 
   const hydrationScript = `<script type="module">
 import setup from '/_astro_frontend/hydrate/${hydrate}.js';
diff --git a/packages/astro/src/runtime.ts b/packages/astro/src/runtime.ts
index d6ef89a5c..53a5fa8bb 100644
--- a/packages/astro/src/runtime.ts
+++ b/packages/astro/src/runtime.ts
@@ -364,7 +364,7 @@ async function createSnowpack(astroConfig: AstroConfig, options: CreateSnowpackO
 
   // Make sure that Snowpack builds our renderer plugins
   const rendererInstances = await configManager.buildRendererInstances();
-  const knownEntrypoints: string[] = ['astro/dist/internal/__astro_component.js'];
+  const knownEntrypoints: string[] = ['astro/dist/internal/__astro_component.js', 'astro/dist/internal/element-registry.js'];
   for (const renderer of rendererInstances) {
     knownEntrypoints.push(renderer.server);
     if(renderer.client) {
@@ -373,6 +373,8 @@ async function createSnowpack(astroConfig: AstroConfig, options: CreateSnowpackO
     if (renderer.knownEntrypoints) {
       knownEntrypoints.push(...renderer.knownEntrypoints);
     }
+    knownEntrypoints.push(...renderer.polyfills);
+    knownEntrypoints.push(...renderer.hydrationPolyfills);
   }
   const external = snowpackExternals.concat([]);
   for(const renderer of rendererInstances) {
diff --git a/packages/astro/test/custom-elements.test.js b/packages/astro/test/custom-elements.test.js
index 2d45e0241..22a16d9aa 100644
--- a/packages/astro/test/custom-elements.test.js
+++ b/packages/astro/test/custom-elements.test.js
@@ -49,6 +49,7 @@ CustomElements('Polyfills are added before the hydration script', async ({ runti
 
   assert.equal($('script[type=module]').length, 2);
   assert.equal($('script[type=module]').attr('src'), '/_snowpack/link/packages/astro/test/fixtures/custom-elements/my-component-lib/polyfill.js');
+  assert.match($($('script[type=module]').get(1)).html(), new RegExp('/_snowpack/link/packages/astro/test/fixtures/custom-elements/my-component-lib/hydration-polyfill.js'));
 });
 
 CustomElements('Polyfills are added even if not hydrating', async ({ runtime }) => {
@@ -60,6 +61,7 @@ CustomElements('Polyfills are added even if not hydrating', async ({ runtime })
 
   assert.equal($('script[type=module]').length, 1);
   assert.equal($('script[type=module]').attr('src'), '/_snowpack/link/packages/astro/test/fixtures/custom-elements/my-component-lib/polyfill.js');
+  assert.not.match($($('script[type=module]').get(1)).html(), new RegExp('/_snowpack/link/packages/astro/test/fixtures/custom-elements/my-component-lib/hydration-polyfill.js'));
 });
 
 CustomElements('Custom elements not claimed by renderer are rendered as regular HTML', async ({ runtime }) => {
diff --git a/packages/astro/test/fixtures/custom-elements/my-component-lib/hydration-polyfill.js b/packages/astro/test/fixtures/custom-elements/my-component-lib/hydration-polyfill.js
new file mode 100644
index 000000000..f3ae60285
--- /dev/null
+++ b/packages/astro/test/fixtures/custom-elements/my-component-lib/hydration-polyfill.js
@@ -0,0 +1,2 @@
+
+globalThis.somePolyfillHere = '';
\ No newline at end of file
diff --git a/packages/astro/test/fixtures/custom-elements/my-component-lib/index.js b/packages/astro/test/fixtures/custom-elements/my-component-lib/index.js
index b67b9906d..0d1ed17a5 100644
--- a/packages/astro/test/fixtures/custom-elements/my-component-lib/index.js
+++ b/packages/astro/test/fixtures/custom-elements/my-component-lib/index.js
@@ -4,5 +4,8 @@ export default {
   server: './server',
   polyfills: [
     './polyfill.js'
+  ],
+  hydrationPolyfills: [
+    './hydration-polyfill.js'
   ]
 };