From 5d96f9856441e7a2a02d306dc67c74d9a6b32fb5 Mon Sep 17 00:00:00 2001
From: Nate Moore <nate@skypack.dev>
Date: Mon, 1 Nov 2021 12:45:25 -0500
Subject: [PATCH] feat: add automatic customElements.define/@customElement
 detection for lit

---
 .changeset/gorgeous-parents-build.md          |   6 +
 packages/astro/src/runtime/server/metadata.ts |   2 +-
 packages/renderers/renderer-lit/index.js      |   5 +
 .../renderers/renderer-lit/vite-plugin-lit.js | 106 ++++++++++++++++++
 4 files changed, 118 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/gorgeous-parents-build.md
 create mode 100644 packages/renderers/renderer-lit/vite-plugin-lit.js

diff --git a/.changeset/gorgeous-parents-build.md b/.changeset/gorgeous-parents-build.md
new file mode 100644
index 000000000..df1170308
--- /dev/null
+++ b/.changeset/gorgeous-parents-build.md
@@ -0,0 +1,6 @@
+---
+'@astrojs/renderer-lit': minor
+'astro': patch
+---
+
+Add support for automatically detecting @customElement("my-element")`and`customElements.define("my-element", MyElement)` when using Lit SSR.
diff --git a/packages/astro/src/runtime/server/metadata.ts b/packages/astro/src/runtime/server/metadata.ts
index 2b318ad80..c5348707d 100644
--- a/packages/astro/src/runtime/server/metadata.ts
+++ b/packages/astro/src/runtime/server/metadata.ts
@@ -41,7 +41,7 @@ class Metadata {
       const id = specifier.startsWith('.') ? new URL(specifier, this.fileURL).pathname : specifier;
       for (const [key, value] of Object.entries(module)) {
         if (isCustomElement) {
-          if (key === 'tagName' && Component === value) {
+          if ((key === 'tagName' || key === '__astroTagName') && Component === value) {
             return {
               componentExport: key,
               componentUrl: id,
diff --git a/packages/renderers/renderer-lit/index.js b/packages/renderers/renderer-lit/index.js
index f08c2fef8..d7d446745 100644
--- a/packages/renderers/renderer-lit/index.js
+++ b/packages/renderers/renderer-lit/index.js
@@ -1,3 +1,5 @@
+import pluginLit from './vite-plugin-lit.js';
+
 // NOTE: @lit-labs/ssr uses syntax incompatible with anything < Node v13.9.0.
 // Throw an error if using that Node version.
 
@@ -13,6 +15,9 @@ export default {
   hydrationPolyfills: ['./hydration-support.js'],
   viteConfig() {
     return {
+      plugins: [
+        pluginLit(),
+      ],
       optimizeDeps: {
         include: [
           '@astrojs/renderer-lit/client-shim.js',
diff --git a/packages/renderers/renderer-lit/vite-plugin-lit.js b/packages/renderers/renderer-lit/vite-plugin-lit.js
new file mode 100644
index 000000000..8a08d9a94
--- /dev/null
+++ b/packages/renderers/renderer-lit/vite-plugin-lit.js
@@ -0,0 +1,106 @@
+import { parse } from 'acorn';
+import { walk } from 'estree-walker';
+import MagicString from 'magic-string';
+
+/**
+ * Determine if Vite is in SSR mode based on options
+ * https://github.com/vitejs/vite/discussions/5109#discussioncomment-1450726
+ *
+ * @param options boolean | { ssr: boolean }
+ * @returns boolean
+ */
+function isSSR(options) {
+  if (options === undefined) {
+    return false;
+  }
+  if (typeof options === 'boolean') {
+    return options;
+  }
+  if (typeof options == 'object') {
+    return !!options.ssr;
+  }
+  return false;
+}
+
+// This matches any JS-like file (that we know of)
+// See https://regex101.com/r/Cgofir/1
+const SUPPORTED_FILES = /\.([cm]?js|jsx|[cm]?ts|tsx)$/;
+const IGNORED_MODULES = [/astro\/dist\/runtime\/server/, /\/renderer-lit\/server/, /\/@lit\//];
+
+function scanForTagName(code) {
+  const ast = parse(code, {
+    sourceType: 'module'
+  })
+
+  let tagName;
+  walk(ast, {
+    enter(node, parent) {
+      if (tagName) {
+        return this.skip();
+      }
+      // Matches `customElement("my-component")`, which is Lit's @customElement("my-component") decorator
+      if (node.type === 'CallExpression' && node.callee.type === 'Identifier' && node.callee.name === 'customElement') {
+        const arg = node.arguments[0];
+        if (arg.type === 'Literal') {
+          tagName = arg.raw
+        }
+      }
+      // Matches `customElements.define("my-component", thing)`
+      if (node.type === 'MemberExpression' && node.object.name === 'customElements' && node.property.name === 'define') {
+        const arg = parent.arguments[0];
+        if (arg.type === 'Literal') {
+          tagName = arg.raw
+        } else {
+          tagName = arg.name
+        }
+      }
+    }
+  })
+
+  return tagName;
+}
+
+/**
+ * @returns {import('vite').Plugin}
+ */
+export default function pluginLit() {
+  return {
+    name: '@astrojs/vite-plugin-lit',
+    enforce: 'post',
+    async transform(code, id, opts) {
+      const ssr = isSSR(opts);
+      // If this isn't an SSR pass, `fetch` will already be available!
+      if (!ssr) {
+        return null;
+      }
+      // Only transform JS-like files
+      if (!id.match(SUPPORTED_FILES)) {
+        return null;
+      }
+      // Optimization: only run on probable matches
+      if (!code.includes('customElement') && !code.includes('lit')) {
+        return null;
+      }
+      // Ignore specific modules
+      for (const ignored of IGNORED_MODULES) {
+        if (id.match(ignored)) {
+          return null;
+        }
+      }
+
+      const tagName = scanForTagName(code);
+      if (!tagName) {
+        return null;
+      }
+    
+      const s = new MagicString(code);
+      s.append(`export const __astroTagName = ${tagName};`);
+      const result = s.toString();
+      const map = s.generateMap({
+        source: id,
+        includeContent: true,
+      });
+      return { code: result, map };
+    },
+  };
+}