Inject styling in HTML AST (#9)

* Inject styling in HTML AST

* Restore optimize structure

8 files changed, 219 insertions, 141 deletions

diff --git a/src/@types/estree-walker.d.ts b/src/@types/estree-walker.d.ts
index 5afb476cb..a3b7da859 100644
--- a/src/@types/estree-walker.d.ts
+++ b/src/@types/estree-walker.d.ts
@@ -11,4 +11,15 @@ declare module 'estree-walker' {
       leave?: (this: { skip: () => void; remove: () => void; replace: (node: T) => void }, node: T, parent: T, key: string, index: number) => void;
     }
   ): T;
+
+  export function asyncWalk<T = BaseNode>(
+    ast: T,
+    {
+      enter,
+      leave,
+    }: {
+      enter?: (this: { skip: () => void; remove: () => void; replace: (node: T) => void }, node: T, parent: T, key: string, index: number) => void;
+      leave?: (this: { skip: () => void; remove: () => void; replace: (node: T) => void }, node: T, parent: T, key: string, index: number) => void;
+    }
+  ): T;
 }
diff --git a/src/optimize/types.ts b/src/@types/optimizer.ts
index e22700cba..c62976068 100644
--- a/src/optimize/types.ts
+++ b/src/@types/optimizer.ts
@@ -1,6 +1,5 @@
 import type { TemplateNode } from '../compiler/interfaces';
 
-
 export type VisitorFn = (node: TemplateNode) => void;
 
 export interface NodeVisitor {
@@ -10,8 +9,8 @@ export interface NodeVisitor {
 
 export interface Optimizer {
   visitors?: {
-    html?: Record<string, NodeVisitor>,
-    css?: Record<string, NodeVisitor>
-  },
-  finalize: () => Promise<void>
-}
\ No newline at end of file
+    html?: Record<string, NodeVisitor>;
+    css?: Record<string, NodeVisitor>;
+  };
+  finalize: () => Promise<void>;
+}
diff --git a/src/codegen/index.ts b/src/codegen/index.ts
index 3257d9936..9b3104f0a 100644
--- a/src/codegen/index.ts
+++ b/src/codegen/index.ts
@@ -190,7 +190,6 @@ export async function codegen(ast: Ast, { compileOptions }: CodeGenOptions): Pro
   let collectionItem: JsxItem | undefined;
   let currentItemName: string | undefined;
   let currentDepth = 0;
-  const classNames: Set<string> = new Set();
 
   walk(ast.html, {
     enter(node: TemplateNode) {
@@ -275,6 +274,11 @@ export async function codegen(ast: Ast, { compileOptions }: CodeGenOptions): Pro
           this.skip();
           return;
         }
+        case 'Style': {
+          const attributes = getAttributes(node.attributes);
+          items.push({ name: 'style', jsx: `h("style", ${attributes ? generateAttributes(attributes) : 'null'}, ${JSON.stringify(node.content.styles)})` });
+          break;
+        }
         case 'Text': {
           const text = getTextFromAttribute(node);
           if (mode === 'SLOT') {
@@ -328,6 +332,9 @@ export async function codegen(ast: Ast, { compileOptions }: CodeGenOptions): Pro
             collectionItem = undefined;
           }
           return;
+        case 'Style': {
+          return;
+        }
         default:
           throw new Error('Unexpected node type: ' + node.type);
       }
diff --git a/src/dev.ts b/src/dev.ts
index 524379dd1..0872ffe74 100644
--- a/src/dev.ts
+++ b/src/dev.ts
@@ -105,6 +105,7 @@ export default async function (astroConfig: AstroConfig) {
           break;
         }
         default: {
+          console.error(err.code, err);
           error(logging, 'running hmx', err);
           break;
         }
diff --git a/src/optimize/index.ts b/src/optimize/index.ts
index a0604b1c8..9f8ec2f05 100644
--- a/src/optimize/index.ts
+++ b/src/optimize/index.ts
@@ -1,7 +1,6 @@
-import type { Ast, TemplateNode } from '../compiler/interfaces';
-import { NodeVisitor, Optimizer, VisitorFn } from './types';
 import { walk } from 'estree-walker';
-
+import type { Ast, TemplateNode } from '../compiler/interfaces';
+import { NodeVisitor, Optimizer, VisitorFn } from '../@types/optimizer';
 import optimizeStyles from './styles.js';
 
 interface VisitorCollection {
@@ -10,13 +9,12 @@ interface VisitorCollection {
 }
 
 function addVisitor(visitor: NodeVisitor, collection: VisitorCollection, nodeName: string, event: 'enter' | 'leave') {
-  if (event in visitor) {
-    if (collection[event].has(nodeName)) {
-      collection[event].get(nodeName)!.push(visitor[event]!);
-    }
+  if (typeof visitor[event] !== 'function') return;
+  if (!collection[event]) collection[event] = new Map<string, VisitorFn[]>();
 
-    collection.enter.set(nodeName, [visitor[event]!]);
-  }
+  const visitors = collection[event].get(nodeName) || [];
+  visitors.push(visitor[event] as any);
+  collection[event].set(nodeName, visitors);
 }
 
 function collectVisitors(optimizer: Optimizer, htmlVisitors: VisitorCollection, cssVisitors: VisitorCollection, finalizers: Array<() => Promise<void>>) {
@@ -77,8 +75,8 @@ export async function optimize(ast: Ast, opts: OptimizeOptions) {
 
   collectVisitors(optimizeStyles(opts), htmlVisitors, cssVisitors, finalizers);
 
-  walkAstWithVisitors(ast.html, htmlVisitors);
   walkAstWithVisitors(ast.css, cssVisitors);
+  walkAstWithVisitors(ast.html, htmlVisitors);
 
   // Run all of the finalizer functions in parallel because why not.
   await Promise.all(finalizers.map((fn) => fn()));
diff --git a/src/optimize/styles.ts b/src/optimize/styles.ts
index 6d15cb602..1353cb006 100644
--- a/src/optimize/styles.ts
+++ b/src/optimize/styles.ts
@@ -1,51 +1,200 @@
-import type { Ast, TemplateNode } from '../compiler/interfaces';
-import type { Optimizer } from './types';
-import { transformStyle } from '../style.js';
+import crypto from 'crypto';
+import path from 'path';
+import autoprefixer from 'autoprefixer';
+import postcss from 'postcss';
+import postcssModules from 'postcss-modules';
+import sass from 'sass';
+import { Optimizer } from '../@types/optimizer';
+import type { TemplateNode } from '../compiler/interfaces';
+
+type StyleType = 'text/css' | 'text/scss' | 'text/sass' | 'text/postcss';
+
+const getStyleType: Map<string, StyleType> = new Map([
+  ['.css', 'text/css'],
+  ['.pcss', 'text/postcss'],
+  ['.sass', 'text/sass'],
+  ['.scss', 'text/scss'],
+  ['css', 'text/css'],
+  ['postcss', 'text/postcss'],
+  ['sass', 'text/sass'],
+  ['scss', 'text/scss'],
+  ['text/css', 'text/css'],
+  ['text/postcss', 'text/postcss'],
+  ['text/sass', 'text/sass'],
+  ['text/scss', 'text/scss'],
+]);
+
+const SASS_OPTIONS: Partial<sass.Options> = {
+  outputStyle: 'compressed',
+};
+
+/** Should be deterministic, given a unique filename */
+function hashFromFilename(filename: string): string {
+  const hash = crypto.createHash('sha256');
+  return hash
+    .update(filename.replace(/\\/g, '/'))
+    .digest('base64')
+    .toString()
+    .replace(/[^A-Za-z0-9-]/g, '')
+    .substr(0, 8);
+}
+
+export interface StyleTransformResult {
+  css: string;
+  cssModules: Map<string, string>;
+  type: StyleType;
+}
+
+async function transformStyle(code: string, { type, filename, fileID }: { type?: string; filename: string; fileID: string }): Promise<StyleTransformResult> {
+  let styleType: StyleType = 'text/css'; // important: assume CSS as default
+  if (type) {
+    styleType = getStyleType.get(type) || styleType;
+  }
+
+  let css = '';
+  switch (styleType) {
+    case 'text/css': {
+      css = code;
+      break;
+    }
+    case 'text/sass':
+    case 'text/scss': {
+      css = sass
+        .renderSync({
+          ...SASS_OPTIONS,
+          data: code,
+          includePaths: [path.dirname(filename)],
+        })
+        .css.toString('utf8');
+      break;
+    }
+    case 'text/postcss': {
+      css = code; // TODO
+      break;
+    }
+    default: {
+      throw new Error(`Unsupported: <style type="${styleType}">`);
+    }
+  }
+
+  const cssModules = new Map<string, string>();
+
+  css = await postcss([
+    postcssModules({
+      generateScopedName(name: string) {
+        return `${name}__${hashFromFilename(fileID)}`;
+      },
+      getJSON(_: string, json: any) {
+        Object.entries(json).forEach(([k, v]: any) => {
+          if (k !== v) cssModules.set(k, v);
+        });
+      },
+    }),
+    autoprefixer(),
+  ])
+    .process(css, { from: filename, to: undefined })
+    .then((result) => result.css);
+
+  return {
+    css,
+    cssModules,
+    type: styleType,
+  };
+}
 
 export default function ({ filename, fileID }: { filename: string; fileID: string }): Optimizer {
-  const classNames: Set<string> = new Set();
-  let stylesPromises: any[] = [];
+  const elementNodes: TemplateNode[] = []; //  elements that need CSS Modules class names
+  const styleNodes: TemplateNode[] = []; // <style> tags to be updated
+  const styleTransformPromises: Promise<StyleTransformResult>[] = []; // async style transform results to be finished in finalize();
+  let rootNode: TemplateNode; // root node which needs <style> tags
 
   return {
     visitors: {
       html: {
         Element: {
           enter(node) {
+            // Find the root node to inject the <style> tag in later
+            if (node.name === 'head') {
+              rootNode = node; // If this is <head>, this is what we want. Always take this if found. However, this may not always exist (it won’t for Component subtrees).
+            } else if (!rootNode) {
+              rootNode = node; // If no <head> (yet), then take the first element we come to and assume it‘s the “root” (but if we find a <head> later, then override this per the above)
+            }
+
             for (let attr of node.attributes) {
-              if (attr.name === 'class') {
-                for (let value of attr.value) {
-                  if (value.type === 'Text') {
-                    const classes = value.data.split(' ');
-                    for (const className in classes) {
-                      classNames.add(className);
-                    }
-                  }
-                }
-              }
+              if (attr.name !== 'class') continue;
+              elementNodes.push(node);
             }
           },
         },
       },
+      // CSS: compile styles, apply CSS Modules scoping
       css: {
         Style: {
-          enter(node: TemplateNode) {
+          enter(node) {
             const code = node.content.styles;
-            const typeAttr = node.attributes && node.attributes.find(({ name }: { name: string }) => name === 'type');
-            stylesPromises.push(
-              transformStyle(code, {
-                type: (typeAttr.value[0] && typeAttr.value[0].raw) || undefined,
-                classNames,
-                filename,
-                fileID,
-              })
-            ); // TODO: styles needs to go in <head>
+            const typeAttr = (node.attributes || []).find(({ name }: { name: string }) => name === 'type');
+            styleNodes.push(node);
+            styleTransformPromises.push(transformStyle(code, { type: (typeAttr.value[0] && typeAttr.value[0].raw) || undefined, filename, fileID }));
+
+            // TODO: we should delete the old untransformed <style> node after we’re done.
+            // However, the svelte parser left it in ast.css, not ast.html. At the final step, this just gets ignored, so it will be deleted, in a sense.
+            // If we ever end up scanning ast.css for something else, then we’ll need to actually delete the node (or transform it to the processed version)
           },
         },
       },
     },
     async finalize() {
-      const styles = await Promise.all(stylesPromises); // TODO: clean this up
-      // console.log({ styles });
+      const allCssModules = new Map<string, string>(); // note: this may theoretically have conflicts, but when written, it shouldn’t because we’re processing everything per-component (if we change this to run across the whole document at once, revisit this)
+      const styleTransforms = await Promise.all(styleTransformPromises);
+
+      if (!rootNode) {
+        throw new Error(`No root node found`); // TODO: remove this eventually; we should always find it, but for now alert if there’s a bug in our code
+      }
+
+      // 1. transform <style> tags
+      styleTransforms.forEach((result, n) => {
+        if (styleNodes[n].attributes) {
+          // Add to global CSS Module class list for step 2
+          for (const [k, v] of result.cssModules) {
+            allCssModules.set(k, v);
+          }
+
+          // Update original <style> node with finished results
+          styleNodes[n].attributes = styleNodes[n].attributes.map((attr: any) => {
+            if (attr.name === 'type') {
+              attr.value[0].raw = 'text/css';
+              attr.value[0].data = 'text/css';
+            }
+            return attr;
+          });
+        }
+        styleNodes[n].content.styles = result.css;
+      });
+
+      // 2. inject finished <style> tags into root node
+      rootNode.children = [...styleNodes, ...(rootNode.children || [])];
+
+      // 3. update HTML classes
+      for (let i = 0; i < elementNodes.length; i++) {
+        if (!elementNodes[i].attributes) continue;
+        const node = elementNodes[i];
+        for (let j = 0; j < node.attributes.length; j++) {
+          if (node.attributes[j].name !== 'class') continue;
+          const attr = node.attributes[j];
+          for (let k = 0; k < attr.value.length; k++) {
+            if (attr.value[k].type !== 'Text') continue;
+            const elementClassNames = (attr.value[k].raw as string)
+              .split(' ')
+              .map((c) => {
+                let className = c.trim();
+                return allCssModules.get(className) || className; // if className matches exactly, replace; otherwise keep original
+              })
+              .join(' ');
+            attr.value[k].raw = elementClassNames;
+            attr.value[k].data = elementClassNames;
+          }
+        }
+      }
     },
   };
 }
diff --git a/src/style.ts b/src/style.ts
deleted file mode 100644
index 489f22ce8..000000000
--- a/src/style.ts
+++ /dev/null
@@ -1,92 +0,0 @@
-import crypto from 'crypto';
-import path from 'path';
-import autoprefixer from 'autoprefixer';
-import postcss from 'postcss';
-import postcssModules from 'postcss-modules';
-import sass from 'sass';
-
-type StyleType = 'text/css' | 'text/scss' | 'text/sass' | 'text/postcss';
-
-const getStyleType: Map<string, StyleType> = new Map([
-  ['.css', 'text/css'],
-  ['.pcss', 'text/postcss'],
-  ['.sass', 'text/sass'],
-  ['.scss', 'text/scss'],
-  ['css', 'text/css'],
-  ['postcss', 'text/postcss'],
-  ['sass', 'text/sass'],
-  ['scss', 'text/scss'],
-  ['text/css', 'text/css'],
-  ['text/postcss', 'text/postcss'],
-  ['text/sass', 'text/sass'],
-  ['text/scss', 'text/scss'],
-]);
-
-const SASS_OPTIONS: Partial<sass.Options> = {
-  outputStyle: 'compressed',
-};
-
-/** Should be deterministic, given a unique filename */
-function hashFromFilename(filename: string): string {
-  const hash = crypto.createHash('sha256');
-  return hash.update(filename.replace(/\\/g, '/')).digest('base64').toString().substr(0, 8);
-}
-
-export async function transformStyle(
-  code: string,
-  { type, classNames, filename, fileID }: { type?: string; classNames?: Set<string>; filename: string; fileID: string }
-): Promise<{ css: string; cssModules: Map<string, string> }> {
-  let styleType: StyleType = 'text/css'; // important: assume CSS as default
-  if (type) {
-    styleType = getStyleType.get(type) || styleType;
-  }
-
-  let css = '';
-  switch (styleType) {
-    case 'text/css': {
-      css = code;
-      break;
-    }
-    case 'text/sass':
-    case 'text/scss': {
-      css = sass
-        .renderSync({
-          ...SASS_OPTIONS,
-          data: code,
-          includePaths: [path.dirname(filename)],
-        })
-        .css.toString('utf8');
-      break;
-    }
-    case 'text/postcss': {
-      css = code; // TODO
-      break;
-    }
-    default: {
-      throw new Error(`Unsupported: <style type="${styleType}">`);
-    }
-  }
-
-  const cssModules = new Map<string, string>();
-
-  css = await postcss([
-    postcssModules({
-      generateScopedName(name: string) {
-        if (classNames && classNames.has(name)) {
-          return `${name}__${hashFromFilename(fileID)}`;
-        }
-        return name;
-      },
-      getJSON(_: string, json: any) {
-        Object.entries(json).forEach(([k, v]: any) => {
-          if (k !== v) cssModules.set(k, v);
-        });
-      },
-    }),
-    autoprefixer(),
-  ])
-    .process(css, { from: filename, to: undefined })
-    .then((result) => result.css);
-
-  return { css, cssModules };
-}
diff --git a/src/transform2.ts b/src/transform2.ts
index 047ccc3d0..0ccdc6b55 100644
--- a/src/transform2.ts
+++ b/src/transform2.ts
@@ -126,13 +126,18 @@ export async function compileComponent(
   { compileOptions = defaultCompileOptions, filename, projectRoot }: { compileOptions: CompileOptions; filename: string; projectRoot: string }
 ): Promise<CompileResult> {
   const sourceJsx = await transformFromSource(source, { compileOptions, filename, projectRoot });
-  const componentJsx = sourceJsx.items.find((item) => item.name === 'Component');
-  if (!componentJsx) {
-    throw new Error(`${filename} <Component> expected!`);
-  }
+
+  // throw error if <Component /> missing
+  if (!sourceJsx.items.find(({ name }) => name === 'Component')) throw new Error(`${filename} <Component> expected!`);
+
+  // sort <style> tags first
+  // TODO: remove these and inject in <head>
+  sourceJsx.items.sort((a, b) => (a.name === 'style' && b.name !== 'style' ? -1 : 0));
+
+  // return template
   const modJsx = `
       import { h, Fragment } from '${internalImport('h.js')}';
-      export default function(props) { return h(Fragment, null, ${componentJsx.jsx}); }
+      export default function(props) { return h(Fragment, null, ${sourceJsx.items.map(({ jsx }) => jsx).join(',')}); }
       `.trim();
   return {
     contents: modJsx,