Add support for syntax highlighting of code blocks (#65)

* Add support for syntax highlighting of code blocks

* Escape usage of backtick strings

* Add workspace root for snowpack

* Use prismjs/components as an external module

4 files changed, 91 insertions, 5 deletions

diff --git a/src/@types/optimizer.ts b/src/@types/optimizer.ts
index a076ae9b8..b6459ab51 100644
--- a/src/@types/optimizer.ts
+++ b/src/@types/optimizer.ts
@@ -1,7 +1,7 @@
 import type { TemplateNode } from '../parser/interfaces';
 import type { CompileOptions } from './compiler';
 
-export type VisitorFn = (node: TemplateNode, parent: TemplateNode, type: string, index: number) => void;
+export type VisitorFn<T = TemplateNode> = (this: { skip: () => void; remove: () => void; replace: (node: T) => void }, node: T, parent: T, type: string, index: number) => void;
 
 export interface NodeVisitor {
   enter?: VisitorFn;
diff --git a/src/compiler/optimize/index.ts b/src/compiler/optimize/index.ts
index a7bf828e0..fcbd6e950 100644
--- a/src/compiler/optimize/index.ts
+++ b/src/compiler/optimize/index.ts
@@ -7,6 +7,7 @@ import { walk } from 'estree-walker';
 import optimizeStyles from './styles.js';
 import optimizeDoctype from './doctype.js';
 import optimizeModuleScripts from './module-scripts.js';
+import optimizeCodeBlocks from './prism.js';
 
 interface VisitorCollection {
   enter: Map<string, VisitorFn[]>;
@@ -57,7 +58,7 @@ function walkAstWithVisitors(tmpl: TemplateNode, collection: VisitorCollection)
       if (collection.enter.has(node.type)) {
         const fns = collection.enter.get(node.type)!;
         for (let fn of fns) {
-          fn(node, parent, key, index);
+          fn.call(this, node, parent, key, index);
         }
       }
     },
@@ -65,7 +66,7 @@ function walkAstWithVisitors(tmpl: TemplateNode, collection: VisitorCollection)
       if (collection.leave.has(node.type)) {
         const fns = collection.leave.get(node.type)!;
         for (let fn of fns) {
-          fn(node, parent, key, index);
+          fn.call(this, node, parent, key, index);
         }
       }
     },
@@ -83,7 +84,7 @@ export async function optimize(ast: Ast, opts: OptimizeOptions) {
   const cssVisitors = createVisitorCollection();
   const finalizers: Array<() => Promise<void>> = [];
 
-  const optimizers = [optimizeStyles(opts), optimizeDoctype(opts), optimizeModuleScripts(opts)];
+  const optimizers = [optimizeStyles(opts), optimizeDoctype(opts), optimizeModuleScripts(opts), optimizeCodeBlocks(ast.module)];
 
   for (const optimizer of optimizers) {
     collectVisitors(optimizer, htmlVisitors, cssVisitors, finalizers);
diff --git a/src/compiler/optimize/prism.ts b/src/compiler/optimize/prism.ts
new file mode 100644
index 000000000..2a96ab73d
--- /dev/null
+++ b/src/compiler/optimize/prism.ts
@@ -0,0 +1,85 @@
+import type { Optimizer } from '../../@types/optimizer';
+import type { Script } from '../../parser/interfaces';
+import { getAttrValue } from '../../ast.js';
+
+const PRISM_IMPORT = `import Prism from 'astro/components/Prism.astro';\n`;
+const prismImportExp = /import Prism from ['"]astro\/components\/Prism.astro['"]/;
+
+function escape(code: string) {
+  return code.replace(/[`$]/g, match => {
+    return '\\' + match;
+  });
+}
+
+export default function (module: Script): Optimizer {
+  let usesPrism = false;
+
+  return {
+    visitors: {
+      html: {
+        Element: {
+          enter(node) {
+            if (node.name !== 'code') return;
+            const className = getAttrValue(node.attributes, 'class') || '';
+            const classes = className.split(' ');
+
+            let lang;
+            for (let cn of classes) {
+              const matches = /language-(.+)/.exec(cn);
+              if (matches) {
+                lang = matches[1];
+              }
+            }
+
+            if (!lang) return;
+
+            let code;
+            if (node.children?.length) {
+              code = node.children[0].data;
+            }
+
+            const repl = {
+              start: 0,
+              end: 0,
+              type: 'InlineComponent',
+              name: 'Prism',
+              attributes: [
+                {
+                  type: 'Attribute',
+                  name: 'lang',
+                  value: [
+                    {
+                      type: 'Text',
+                      raw: lang,
+                      data: lang,
+                    },
+                  ],
+                },
+                {
+                  type: 'Attribute',
+                  name: 'code',
+                  value: [
+                    {
+                      type: 'MustacheTag',
+                      content: '`' + escape(code) + '`',
+                    },
+                  ],
+                },
+              ],
+              children: [],
+            };
+
+            this.replace(repl);
+            usesPrism = true;
+          },
+        },
+      },
+    },
+    async finalize() {
+      // Add the Prism import if needed.
+      if (usesPrism && !prismImportExp.test(module.content)) {
+        module.content = PRISM_IMPORT + module.content;
+      }
+    },
+  };
+}
diff --git a/src/runtime.ts b/src/runtime.ts
index 800c4d40a..f524a2bf5 100644
--- a/src/runtime.ts
+++ b/src/runtime.ts
@@ -164,7 +164,7 @@ async function createSnowpack(astroConfig: AstroConfig, env: Record<string, any>
     },
     packageOptions: {
       knownEntrypoints: ['preact-render-to-string'],
-      external: ['@vue/server-renderer', 'node-fetch'],
+      external: ['@vue/server-renderer', 'node-fetch', 'prismjs/components/'],
     },
   });