Refactor to enable optimizer modules (#8)

* Refactor to enable optimizer modules

This refactors HMX compilation into steps:

1. Parse - Turn HMX string into an AST.
2. Optimize - Walk the AST making modifications.
3. Codegen - Turn the AST into hyperscript function calls.

There's still more logic in (3) than we probably want. The nice there here is it gives a Visitor API that you can implement to do optimizations. See src/optimize/styles.ts for an example.

* Allow multiple visitors per optimizer

6 files changed, 515 insertions, 357 deletions

diff --git a/src/@types/compiler.ts b/src/@types/compiler.ts
new file mode 100644
index 000000000..343aa548b
--- /dev/null
+++ b/src/@types/compiler.ts
@@ -0,0 +1,6 @@
+import type { LogOptions } from '../logger';
+
+export interface CompileOptions {
+  logging: LogOptions;
+  resolve: (p: string) => string;
+}
\ No newline at end of file
diff --git a/src/codegen/index.ts b/src/codegen/index.ts
new file mode 100644
index 000000000..c0f4199c6
--- /dev/null
+++ b/src/codegen/index.ts
@@ -0,0 +1,342 @@
+import type { CompileOptions } from '../@types/compiler';
+import type { Ast, TemplateNode } from '../compiler/interfaces';
+import type { JsxItem } from '../@types/astro.js';
+
+import eslexer from 'es-module-lexer';
+import esbuild from 'esbuild';
+import path from 'path';
+import { walk } from 'estree-walker';
+
+const { transformSync } = esbuild;
+
+interface Attribute {
+  start: number;
+  end: number;
+  type: 'Attribute';
+  name: string;
+  value: any
+}
+
+interface CodeGenOptions {
+  compileOptions: CompileOptions;
+  filename: string;
+  fileID: string
+}
+
+function internalImport(internalPath: string) {
+  return `/__hmx_internal__/${internalPath}`;
+}
+
+function getAttributes(attrs: Attribute[]): Record<string, string> {
+  let result: Record<string, string> = {};
+  for (const attr of attrs) {
+    if (attr.value === true) {
+      result[attr.name] = JSON.stringify(attr.value);
+      continue;
+    }
+    if (attr.value === false) {
+      continue;
+    }
+    if (attr.value.length > 1) {
+      result[attr.name] =
+        '(' +
+        attr.value
+          .map((v: TemplateNode) => {
+            if (v.expression) {
+              return v.expression;
+            } else {
+              return JSON.stringify(getTextFromAttribute(v));
+            }
+          })
+          .join('+') +
+        ')';
+      continue;
+    }
+    const val: TemplateNode = attr.value[0];
+    switch (val.type) {
+      case 'MustacheTag':
+        result[attr.name] = '(' + val.expression + ')';
+        continue;
+      case 'Text':
+        result[attr.name] = JSON.stringify(getTextFromAttribute(val));
+        continue;
+      default:
+        console.log(val);
+        throw new Error('UNKNOWN V');
+    }
+  }
+  return result;
+}
+
+function getTextFromAttribute(attr: any): string {
+  if (attr.raw !== undefined) {
+    return attr.raw;
+  }
+  if (attr.data !== undefined) {
+    return attr.data;
+  }
+  console.log(attr);
+  throw new Error('UNKNOWN attr');
+}
+
+function generateAttributes(attrs: Record<string, string>): string {
+  let result = '{';
+  for (const [key, val] of Object.entries(attrs)) {
+    result += JSON.stringify(key) + ':' + val + ',';
+  }
+  return result + '}';
+}
+
+function getComponentWrapper(_name: string, { type, url }: { type: string; url: string }, { resolve }: CompileOptions) {
+  const [name, kind] = _name.split(':');
+  switch (type) {
+    case '.hmx': {
+      if (kind) {
+        throw new Error(`HMX does not support :${kind}`);
+      }
+      return {
+        wrapper: name,
+        wrapperImport: ``,
+      };
+    }
+    case '.jsx': {
+      if (kind === 'dynamic') {
+        return {
+          wrapper: `__preact_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.js'))}, \`http://TEST\${import.meta.url}\`).pathname, '${resolve('preact')}')`,
+          wrapperImport: `import {__preact_dynamic} from '${internalImport('render/preact.js')}';`,
+        };
+      } else {
+        return {
+          wrapper: `__preact_static(${name})`,
+          wrapperImport: `import {__preact_static} from '${internalImport('render/preact.js')}';`,
+        };
+      }
+    }
+    case '.svelte': {
+      if (kind === 'dynamic') {
+        return {
+          wrapper: `__svelte_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.svelte.js'))}, \`http://TEST\${import.meta.url}\`).pathname)`,
+          wrapperImport: `import {__svelte_dynamic} from '${internalImport('render/svelte.js')}';`,
+        };
+      } else {
+        return {
+          wrapper: `__svelte_static(${name})`,
+          wrapperImport: `import {__svelte_static} from '${internalImport('render/svelte.js')}';`,
+        };
+      }
+    }
+    case '.vue': {
+      if (kind === 'dynamic') {
+        return {
+          wrapper: `__vue_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.vue.js'))}, \`http://TEST\${import.meta.url}\`).pathname, '${resolve('vue')}')`,
+          wrapperImport: `import {__vue_dynamic} from '${internalImport('render/vue.js')}';`,
+        };
+      } else {
+        return {
+          wrapper: `__vue_static(${name})`,
+          wrapperImport: `
+            import {__vue_static} from '${internalImport('render/vue.js')}';
+          `,
+        };
+      }
+    }
+  }
+  throw new Error('Unknown Component Type: ' + name);
+}
+
+const patternImport = new RegExp(/import(?:["'\s]*([\w*${}\n\r\t, ]+)from\s*)?["'\s]["'\s](.*[@\w_-]+)["'\s].*;$/, 'mg');
+function compileScriptSafe(raw: string, loader: 'jsx' | 'tsx'): string {
+  // esbuild treeshakes unused imports. In our case these are components, so let's keep them.
+  const imports: Array<string> = [];
+  raw.replace(patternImport, (value: string) => {
+    imports.push(value);
+    return value;
+  });
+
+  let { code } = transformSync(raw, {
+    loader,
+    jsxFactory: 'h',
+    jsxFragment: 'Fragment',
+    charset: 'utf8',
+  });
+
+  for (let importStatement of imports) {
+    if (!code.includes(importStatement)) {
+      code = importStatement + '\n' + code;
+    }
+  }
+
+  return code;
+}
+
+export async function codegen(ast: Ast, { compileOptions }: CodeGenOptions) {
+  const script = compileScriptSafe(ast.instance ? ast.instance.content : '', 'tsx');
+
+  // Compile scripts as TypeScript, always
+
+  // Todo: Validate that `h` and `Fragment` aren't defined in the script
+
+  const [scriptImports] = eslexer.parse(script, 'optional-sourcename');
+  const components = Object.fromEntries(
+    scriptImports.map((imp) => {
+      const componentType = path.posix.extname(imp.n!);
+      const componentName = path.posix.basename(imp.n!, componentType);
+      return [componentName, { type: componentType, url: imp.n! }];
+    })
+  );
+
+  const additionalImports = new Set<string>();
+  let items: JsxItem[] = [];
+  let mode: 'JSX' | 'SCRIPT' | 'SLOT' = 'JSX';
+  let collectionItem: JsxItem | undefined;
+  let currentItemName: string | undefined;
+  let currentDepth = 0;
+  const classNames: Set<string> = new Set();
+
+  walk(ast.html, {
+    enter(node: TemplateNode) {
+      //   console.log("enter", node.type);
+      switch (node.type) {
+        case 'MustacheTag':
+          let code = compileScriptSafe(node.expression, 'jsx');
+
+          let matches: RegExpExecArray[] = [];
+          let match: RegExpExecArray | null | undefined;
+          const H_COMPONENT_SCANNER = /h\(['"]?([A-Z].*?)['"]?,/gs;
+          const regex = new RegExp(H_COMPONENT_SCANNER);
+          while ((match = regex.exec(code))) {
+            matches.push(match);
+          }
+          for (const match of matches.reverse()) {
+            const name = match[1];
+            const [componentName, componentKind] = name.split(':');
+            if (!components[componentName]) {
+              throw new Error(`Unknown Component: ${componentName}`);
+            }
+            const { wrapper, wrapperImport } = getComponentWrapper(name, components[componentName], compileOptions);
+            if (wrapperImport) {
+              additionalImports.add(wrapperImport);
+            }
+            if (wrapper !== name) {
+              code = code.slice(0, match.index + 2) + wrapper + code.slice(match.index + match[0].length - 1);
+            }
+          }
+          collectionItem!.jsx += `,(${code.trim().replace(/\;$/, '')})`;
+          return;
+        case 'Slot':
+          mode = 'SLOT';
+          collectionItem!.jsx += `,child`;
+          return;
+        case 'Comment':
+          return;
+        case 'Fragment':
+          // Ignore if its the top level fragment
+          // This should be cleaned up, but right now this is how the old thing worked
+          if (!collectionItem) {
+            return;
+          }
+          break;
+        case 'InlineComponent':
+        case 'Element':
+          const name: string = node.name;
+          if (!name) {
+            console.log(node);
+            throw new Error('AHHHH');
+          }
+          const attributes = getAttributes(node.attributes);
+          currentDepth++;
+          currentItemName = name;
+          if (!collectionItem) {
+            collectionItem = { name, jsx: '' };
+            items.push(collectionItem);
+          }
+          collectionItem.jsx += collectionItem.jsx === '' ? '' : ',';
+          const COMPONENT_NAME_SCANNER = /^[A-Z]/;
+          if (!COMPONENT_NAME_SCANNER.test(name)) {
+            collectionItem.jsx += `h("${name}", ${attributes ? generateAttributes(attributes) : 'null'}`;
+            return;
+          }
+          if (name === 'Component') {
+            collectionItem.jsx += `h(Fragment, null`;
+            return;
+          }
+          const [componentName, componentKind] = name.split(':');
+          const componentImportData = components[componentName];
+          if (!componentImportData) {
+            throw new Error(`Unknown Component: ${componentName}`);
+          }
+          const { wrapper, wrapperImport } = getComponentWrapper(name, components[componentName], compileOptions);
+          if (wrapperImport) {
+            additionalImports.add(wrapperImport);
+          }
+
+          collectionItem.jsx += `h(${wrapper}, ${attributes ? generateAttributes(attributes) : 'null'}`;
+          return;
+        case 'Attribute': {
+          this.skip();
+          return;
+        }
+        case 'Text': {
+          const text = getTextFromAttribute(node);
+          if (mode === 'SLOT') {
+            return;
+          }
+          if (!text.trim()) {
+            return;
+          }
+          if (!collectionItem) {
+            throw new Error('Not possible! TEXT:' + text);
+          }
+          if (currentItemName === 'script' || currentItemName === 'code') {
+            collectionItem.jsx += ',' + JSON.stringify(text);
+            return;
+          }
+          collectionItem.jsx += ',' + JSON.stringify(text);
+          return;
+        }
+        default:
+          console.log(node);
+          throw new Error('Unexpected node type: ' + node.type);
+      }
+    },
+    leave(node, parent, prop, index) {
+      //   console.log("leave", node.type);
+      switch (node.type) {
+        case 'Text':
+        case 'MustacheTag':
+        case 'Attribute':
+        case 'Comment':
+          return;
+        case 'Slot': {
+          const name = node.name;
+          if (name === 'slot') {
+            mode = 'JSX';
+          }
+          return;
+        }
+        case 'Fragment':
+          if (!collectionItem) {
+            return;
+          }
+        case 'Element':
+        case 'InlineComponent':
+          if (!collectionItem) {
+            throw new Error('Not possible! CLOSE ' + node.name);
+          }
+          collectionItem.jsx += ')';
+          currentDepth--;
+          if (currentDepth === 0) {
+            collectionItem = undefined;
+          }
+          return;
+        default:
+          throw new Error('Unexpected node type: ' + node.type);
+      }
+    },
+  });
+
+  return {
+    script: script + '\n' + Array.from(additionalImports).join('\n'),
+    items,
+  };
+}
\ No newline at end of file
diff --git a/src/optimize/index.ts b/src/optimize/index.ts
new file mode 100644
index 000000000..d22854a32
--- /dev/null
+++ b/src/optimize/index.ts
@@ -0,0 +1,85 @@
+import type { Ast, TemplateNode } from '../compiler/interfaces';
+import { NodeVisitor, Optimizer, VisitorFn } from './types';
+import { walk } from 'estree-walker';
+
+import optimizeStyles from './styles.js';
+
+interface VisitorCollection {
+  enter: Map<string, VisitorFn[]>;
+  leave: Map<string, VisitorFn[]>;
+}
+
+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]!);
+    }
+
+    collection.enter.set(nodeName, [visitor[event]!]);
+  }
+}
+
+function collectVisitors(optimizer: Optimizer, htmlVisitors: VisitorCollection, cssVisitors: VisitorCollection, finalizers: Array<() => Promise<void>>) {
+  if(optimizer.visitors) {
+    if(optimizer.visitors.html) {
+      for(const [nodeName, visitor] of Object.entries(optimizer.visitors.html)) {
+        addVisitor(visitor, htmlVisitors, nodeName, 'enter');
+        addVisitor(visitor, htmlVisitors, nodeName, 'leave');
+      }
+    }
+    if(optimizer.visitors.css) {
+      for(const [nodeName, visitor] of Object.entries(optimizer.visitors.css)) {
+        addVisitor(visitor, cssVisitors, nodeName, 'enter');
+        addVisitor(visitor, cssVisitors, nodeName, 'leave');
+      }
+    }
+  }
+  finalizers.push(optimizer.finalize);
+}
+
+function createVisitorCollection() {
+  return {
+    enter: new Map<string, VisitorFn[]>(),
+    leave: new Map<string, VisitorFn[]>(),
+  };
+}
+
+function walkAstWithVisitors(tmpl: TemplateNode, collection: VisitorCollection) {
+  walk(tmpl, {
+    enter(node) {
+      if(collection.enter.has(node.type)) {
+        const fns = collection.enter.get(node.type)!;
+        for(let fn of fns) {
+          fn(node);
+        }
+      }
+    },
+    leave(node) {
+      if(collection.leave.has(node.type)) {
+        const fns = collection.leave.get(node.type)!;
+        for(let fn of fns) {
+          fn(node);
+        }
+      }
+    }
+  });
+}
+
+interface OptimizeOptions { 
+  filename: string,
+  fileID: string
+}
+
+export async function optimize(ast: Ast, opts: OptimizeOptions) {
+  const htmlVisitors = createVisitorCollection();
+  const cssVisitors = createVisitorCollection();
+  const finalizers: Array<() => Promise<void>> = [];
+
+  collectVisitors(optimizeStyles(opts), htmlVisitors, cssVisitors, finalizers);
+
+  walkAstWithVisitors(ast.html, htmlVisitors);
+  walkAstWithVisitors(ast.css, cssVisitors);
+
+  // Run all of the finalizer functions in parallel because why not.
+  await Promise.all(finalizers.map(fn => fn()));
+}
\ No newline at end of file
diff --git a/src/optimize/styles.ts b/src/optimize/styles.ts
new file mode 100644
index 000000000..b654ca7d1
--- /dev/null
+++ b/src/optimize/styles.ts
@@ -0,0 +1,51 @@
+import type { Ast, TemplateNode } from '../compiler/interfaces';
+import type { Optimizer } from './types'
+import { transformStyle } from '../style.js';
+
+export default function({ filename, fileID }: { filename: string, fileID: string }): Optimizer {
+  const classNames: Set<string> = new Set();
+  let stylesPromises: any[] = [];
+
+  return {
+    visitors: {
+      html: {
+        Element: {
+          enter(node) {
+            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);
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      },
+      css: {
+        Style: {
+          enter(node: TemplateNode) {
+            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>
+          }
+        }
+      }
+    },
+    async finalize() {
+      const styles = await Promise.all(stylesPromises); // TODO: clean this up
+      console.log({ styles });
+    }
+  };
+}
\ No newline at end of file
diff --git a/src/optimize/types.ts b/src/optimize/types.ts
new file mode 100644
index 000000000..e22700cba
--- /dev/null
+++ b/src/optimize/types.ts
@@ -0,0 +1,17 @@
+import type { TemplateNode } from '../compiler/interfaces';
+
+
+export type VisitorFn = (node: TemplateNode) => void;
+
+export interface NodeVisitor {
+  enter?: VisitorFn;
+  leave?: VisitorFn;
+}
+
+export interface Optimizer {
+  visitors?: {
+    html?: Record<string, NodeVisitor>,
+    css?: Record<string, NodeVisitor>
+  },
+  finalize: () => Promise<void>
+}
\ No newline at end of file
diff --git a/src/transform2.ts b/src/transform2.ts
index e54845baa..84277efdf 100644
--- a/src/transform2.ts
+++ b/src/transform2.ts
@@ -7,23 +7,11 @@ import micromark from 'micromark';
 import gfmSyntax from 'micromark-extension-gfm';
 import matter from 'gray-matter';
 import gfmHtml from 'micromark-extension-gfm/html.js';
-import { walk } from 'estree-walker';
 import { parse } from './compiler/index.js';
 import markdownEncode from './markdown-encode.js';
-import { TemplateNode } from './compiler/interfaces.js';
-import { defaultLogOptions, info } from './logger.js';
-import { transformStyle } from './style.js';
-import { JsxItem } from './@types/astro.js';
-
-const { transformSync } = esbuild;
-
-interface Attribute {
-  start: 574;
-  end: 595;
-  type: 'Attribute';
-  name: 'class';
-  value: any;
-}
+import { defaultLogOptions } from './logger.js';
+import { optimize } from './optimize/index.js';
+import { codegen } from './codegen/index.js';
 
 interface CompileOptions {
   logging: LogOptions;
@@ -39,357 +27,26 @@ function internalImport(internalPath: string) {
   return `/__hmx_internal__/${internalPath}`;
 }
 
-function getAttributes(attrs: Attribute[]): Record<string, string> {
-  let result: Record<string, string> = {};
-  for (const attr of attrs) {
-    if (attr.value === true) {
-      result[attr.name] = JSON.stringify(attr.value);
-      continue;
-    }
-    if (attr.value === false) {
-      continue;
-    }
-    if (attr.value.length > 1) {
-      result[attr.name] =
-        '(' +
-        attr.value
-          .map((v: TemplateNode) => {
-            if (v.expression) {
-              return v.expression;
-            } else {
-              return JSON.stringify(getTextFromAttribute(v));
-            }
-          })
-          .join('+') +
-        ')';
-      continue;
-    }
-    const val: TemplateNode = attr.value[0];
-    switch (val.type) {
-      case 'MustacheTag':
-        result[attr.name] = '(' + val.expression + ')';
-        continue;
-      case 'Text':
-        result[attr.name] = JSON.stringify(getTextFromAttribute(val));
-        continue;
-      default:
-        console.log(val);
-        throw new Error('UNKNOWN V');
-    }
-  }
-  return result;
-}
-
-function getTextFromAttribute(attr: any): string {
-  if (attr.raw !== undefined) {
-    return attr.raw;
-  }
-  if (attr.data !== undefined) {
-    return attr.data;
-  }
-  console.log(attr);
-  throw new Error('UNKNOWN attr');
-}
-
-function generateAttributes(attrs: Record<string, string>): string {
-  let result = '{';
-  for (const [key, val] of Object.entries(attrs)) {
-    result += JSON.stringify(key) + ':' + val + ',';
-  }
-  return result + '}';
+interface ConvertHmxOptions {
+  compileOptions: CompileOptions;
+  filename: string;
+  fileID: string
 }
 
-function getComponentWrapper(_name: string, { type, url }: { type: string; url: string }, { resolve }: CompileOptions) {
-  const [name, kind] = _name.split(':');
-  switch (type) {
-    case '.hmx': {
-      if (kind) {
-        throw new Error(`HMX does not support :${kind}`);
-      }
-      return {
-        wrapper: name,
-        wrapperImport: ``,
-      };
-    }
-    case '.jsx': {
-      if (kind === 'dynamic') {
-        return {
-          wrapper: `__preact_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.js'))}, \`http://TEST\${import.meta.url}\`).pathname, '${resolve('preact')}')`,
-          wrapperImport: `import {__preact_dynamic} from '${internalImport('render/preact.js')}';`,
-        };
-      } else {
-        return {
-          wrapper: `__preact_static(${name})`,
-          wrapperImport: `import {__preact_static} from '${internalImport('render/preact.js')}';`,
-        };
-      }
-    }
-    case '.svelte': {
-      if (kind === 'dynamic') {
-        return {
-          wrapper: `__svelte_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.svelte.js'))}, \`http://TEST\${import.meta.url}\`).pathname)`,
-          wrapperImport: `import {__svelte_dynamic} from '${internalImport('render/svelte.js')}';`,
-        };
-      } else {
-        return {
-          wrapper: `__svelte_static(${name})`,
-          wrapperImport: `import {__svelte_static} from '${internalImport('render/svelte.js')}';`,
-        };
-      }
-    }
-    case '.vue': {
-      if (kind === 'dynamic') {
-        return {
-          wrapper: `__vue_dynamic(${name}, new URL(${JSON.stringify(url.replace(/\.[^.]+$/, '.vue.js'))}, \`http://TEST\${import.meta.url}\`).pathname, '${resolve('vue')}')`,
-          wrapperImport: `import {__vue_dynamic} from '${internalImport('render/vue.js')}';`,
-        };
-      } else {
-        return {
-          wrapper: `__vue_static(${name})`,
-          wrapperImport: `
-            import {__vue_static} from '${internalImport('render/vue.js')}';
-          `,
-        };
-      }
-    }
-  }
-  throw new Error('Unknown Component Type: ' + name);
-}
-
-const patternImport = new RegExp(/import(?:["'\s]*([\w*${}\n\r\t, ]+)from\s*)?["'\s]["'\s](.*[@\w_-]+)["'\s].*;$/, 'mg');
-function compileScriptSafe(raw: string, loader: 'jsx' | 'tsx'): string {
-  // esbuild treeshakes unused imports. In our case these are components, so let's keep them.
-  const imports: Array<string> = [];
-  raw.replace(patternImport, (value: string) => {
-    imports.push(value);
-    return value;
-  });
-
-  let { code } = transformSync(raw, {
-    loader,
-    jsxFactory: 'h',
-    jsxFragment: 'Fragment',
-    charset: 'utf8',
-  });
-
-  for (let importStatement of imports) {
-    if (!code.includes(importStatement)) {
-      code = importStatement + '\n' + code;
-    }
-  }
-
-  return code;
-}
-
-async function convertHmxToJsx(template: string, { compileOptions, filename, fileID }: { compileOptions: CompileOptions; filename: string; fileID: string }) {
+async function convertHmxToJsx(template: string, opts: ConvertHmxOptions) {
+  const { filename } = opts;
   await eslexer.init;
 
+  // 1. Parse
   const ast = parse(template, {
     filename,
   });
-  const script = compileScriptSafe(ast.instance ? ast.instance.content : '', 'tsx');
-
-  // Compile scripts as TypeScript, always
-
-  // Todo: Validate that `h` and `Fragment` aren't defined in the script
-
-  const [scriptImports] = eslexer.parse(script, 'optional-sourcename');
-  const components = Object.fromEntries(
-    scriptImports.map((imp) => {
-      const componentType = path.posix.extname(imp.n!);
-      const componentName = path.posix.basename(imp.n!, componentType);
-      return [componentName, { type: componentType, url: imp.n! }];
-    })
-  );
-
-  const additionalImports = new Set<string>();
-  let items: JsxItem[] = [];
-  let mode: 'JSX' | 'SCRIPT' | 'SLOT' = 'JSX';
-  let collectionItem: JsxItem | undefined;
-  let currentItemName: string | undefined;
-  let currentDepth = 0;
-  const classNames: Set<string> = new Set();
-
-  walk(ast.html, {
-    enter(node, parent, prop, index) {
-      //   console.log("enter", node.type);
-      switch (node.type) {
-        case 'MustacheTag':
-          let code = compileScriptSafe(node.expression, 'jsx');
 
-          let matches: RegExpExecArray[] = [];
-          let match: RegExpExecArray | null | undefined;
-          const H_COMPONENT_SCANNER = /h\(['"]?([A-Z].*?)['"]?,/gs;
-          const regex = new RegExp(H_COMPONENT_SCANNER);
-          while ((match = regex.exec(code))) {
-            matches.push(match);
-          }
-          for (const match of matches.reverse()) {
-            const name = match[1];
-            const [componentName, componentKind] = name.split(':');
-            if (!components[componentName]) {
-              throw new Error(`Unknown Component: ${componentName}`);
-            }
-            const { wrapper, wrapperImport } = getComponentWrapper(name, components[componentName], compileOptions);
-            if (wrapperImport) {
-              additionalImports.add(wrapperImport);
-            }
-            if (wrapper !== name) {
-              code = code.slice(0, match.index + 2) + wrapper + code.slice(match.index + match[0].length - 1);
-            }
-          }
-          collectionItem!.jsx += `,(${code.trim().replace(/\;$/, '')})`;
-          return;
-        case 'Slot':
-          mode = 'SLOT';
-          collectionItem!.jsx += `,child`;
-          return;
-        case 'Comment':
-          return;
-        case 'Fragment':
-          // Ignore if its the top level fragment
-          // This should be cleaned up, but right now this is how the old thing worked
-          if (!collectionItem) {
-            return;
-          }
-          break;
-        case 'InlineComponent':
-        case 'Element':
-          const name: string = node.name;
-          if (!name) {
-            console.log(node);
-            throw new Error('AHHHH');
-          }
-          const attributes = getAttributes(node.attributes);
-          currentDepth++;
-          currentItemName = name;
-          if (!collectionItem) {
-            collectionItem = { name, jsx: '' };
-            items.push(collectionItem);
-          }
-          if (attributes.class) {
-            attributes.class
-              .replace(/^"/, '')
-              .replace(/"$/, '')
-              .split(' ')
-              .map((c) => c.trim())
-              .forEach((c) => classNames.add(c));
-          }
-          collectionItem.jsx += collectionItem.jsx === '' ? '' : ',';
-          const COMPONENT_NAME_SCANNER = /^[A-Z]/;
-          if (!COMPONENT_NAME_SCANNER.test(name)) {
-            collectionItem.jsx += `h("${name}", ${attributes ? generateAttributes(attributes) : 'null'}`;
-            return;
-          }
-          if (name === 'Component') {
-            collectionItem.jsx += `h(Fragment, null`;
-            return;
-          }
-          const [componentName, componentKind] = name.split(':');
-          const componentImportData = components[componentName];
-          if (!componentImportData) {
-            throw new Error(`Unknown Component: ${componentName}`);
-          }
-          const { wrapper, wrapperImport } = getComponentWrapper(name, components[componentName], compileOptions);
-          if (wrapperImport) {
-            additionalImports.add(wrapperImport);
-          }
+  // 2. Optimize the AST
+  await optimize(ast, opts);
 
-          collectionItem.jsx += `h(${wrapper}, ${attributes ? generateAttributes(attributes) : 'null'}`;
-          return;
-        case 'Attribute': {
-          this.skip();
-          return;
-        }
-        case 'Text': {
-          const text = getTextFromAttribute(node);
-          if (mode === 'SLOT') {
-            return;
-          }
-          if (!text.trim()) {
-            return;
-          }
-          if (!collectionItem) {
-            throw new Error('Not possible! TEXT:' + text);
-          }
-          if (currentItemName === 'script' || currentItemName === 'code') {
-            collectionItem.jsx += ',' + JSON.stringify(text);
-            return;
-          }
-          collectionItem.jsx += ',' + JSON.stringify(text);
-          return;
-        }
-        default:
-          console.log(node);
-          throw new Error('Unexpected node type: ' + node.type);
-      }
-    },
-    leave(node, parent, prop, index) {
-      //   console.log("leave", node.type);
-      switch (node.type) {
-        case 'Text':
-        case 'MustacheTag':
-        case 'Attribute':
-        case 'Comment':
-          return;
-        case 'Slot': {
-          const name = node.name;
-          if (name === 'slot') {
-            mode = 'JSX';
-          }
-          return;
-        }
-        case 'Fragment':
-          if (!collectionItem) {
-            return;
-          }
-        case 'Element':
-        case 'InlineComponent':
-          if (!collectionItem) {
-            throw new Error('Not possible! CLOSE ' + node.name);
-          }
-          collectionItem.jsx += ')';
-          currentDepth--;
-          if (currentDepth === 0) {
-            collectionItem = undefined;
-          }
-          return;
-        default:
-          throw new Error('Unexpected node type: ' + node.type);
-      }
-    },
-  });
-
-  let stylesPromises: any[] = [];
-  walk(ast.css, {
-    enter(node) {
-      if (node.type !== 'Style') return;
-
-      const code = node.content.styles;
-      const typeAttr = node.attributes && node.attributes.find(({ name }) => 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 styles = await Promise.all(stylesPromises); // TODO: clean this up
-  console.log({ styles });
-
-  // console.log({
-  //   additionalImports,
-  //   script,
-  //   items,
-  // });
-
-  return {
-    script: script + '\n' + Array.from(additionalImports).join('\n'),
-    items,
-  };
+  // Turn AST into JSX
+  return await codegen(ast, opts);
 }
 
 async function convertMdToJsx(contents: string, { compileOptions, filename, fileID }: { compileOptions: CompileOptions; filename: string; fileID: string }) {