Add support for doctype (#37)

* Add support for doctype

* Automatically prepend doctype

2 files changed, 41 insertions, 5 deletions

diff --git a/src/compiler/optimize/doctype.ts b/src/compiler/optimize/doctype.ts
new file mode 100644
index 000000000..a666876bf
--- /dev/null
+++ b/src/compiler/optimize/doctype.ts
@@ -0,0 +1,33 @@
+import { Optimizer } from '../../@types/optimizer';
+
+export default function (_opts: { filename: string; fileID: string }): Optimizer {
+  let hasDoctype = false;
+
+  return {
+    visitors: {
+      html: {
+        Element: {
+          enter(node, parent, _key, index) {
+            if(node.name === '!doctype') {
+              hasDoctype = true;
+            }
+            if(node.name === 'html' && !hasDoctype) {
+              const dtNode = {
+                start: 0, end: 0,
+                attributes: [{ type: 'Attribute', name: 'html', value: true, start: 0, end: 0 }],
+                children: [],
+                name: '!doctype',
+                type: 'Element'
+              };
+              parent.children!.splice(index, 0, dtNode);
+              hasDoctype = true;
+            }
+          }
+        }
+      }
+    },
+    async finalize() {
+      // Nothing happening here.
+    }
+  }
+}
\ No newline at end of file
diff --git a/src/compiler/optimize/index.ts b/src/compiler/optimize/index.ts
index aa5ca58f3..d86ce3c24 100644
--- a/src/compiler/optimize/index.ts
+++ b/src/compiler/optimize/index.ts
@@ -1,7 +1,10 @@
 import { walk } from 'estree-walker';
 import type { Ast, TemplateNode } from '../../parser/interfaces';
 import { NodeVisitor, Optimizer, VisitorFn } from '../../@types/optimizer';
+
+// Optimizers
 import optimizeStyles from './styles.js';
+import optimizeDoctype from './doctype.js';
 
 interface VisitorCollection {
   enter: Map<string, VisitorFn[]>;
@@ -44,19 +47,19 @@ function createVisitorCollection() {
 
 function walkAstWithVisitors(tmpl: TemplateNode, collection: VisitorCollection) {
   walk(tmpl, {
-    enter(node) {
+    enter(node, parent, key, index) {
       if (collection.enter.has(node.type)) {
         const fns = collection.enter.get(node.type)!;
         for (let fn of fns) {
-          fn(node);
+          fn(node, parent, key, index);
         }
       }
     },
-    leave(node) {
+    leave(node, parent, key, index) {
       if (collection.leave.has(node.type)) {
         const fns = collection.leave.get(node.type)!;
         for (let fn of fns) {
-          fn(node);
+          fn(node, parent, key, index);
         }
       }
     },
@@ -73,7 +76,7 @@ export async function optimize(ast: Ast, opts: OptimizeOptions) {
   const cssVisitors = createVisitorCollection();
   const finalizers: Array<() => Promise<void>> = [];
 
-  const optimizers = [optimizeStyles(opts)];
+  const optimizers = [optimizeStyles(opts), optimizeDoctype(opts)];
 
   for (const optimizer of optimizers) {
     collectVisitors(optimizer, htmlVisitors, cssVisitors, finalizers);