1 files changed, 103 insertions, 43 deletions

diff --git a/packages/gitbook-markdown/src/page.js b/packages/gitbook-markdown/src/page.js
index 70da324..5bd20e2 100644
--- a/packages/gitbook-markdown/src/page.js
+++ b/packages/gitbook-markdown/src/page.js
@@ -1,21 +1,99 @@
-const { State } = require('markup-it');
+const { State, Block, Text, Inline, INLINES, BLOCKS, MARKS } = require('markup-it');
 const markdown = require('markup-it/lib/markdown');
 
-const RAW_START = '{% raw %}';
-const RAW_END   = '{% endraw %}';
+const RAW_START = 'raw';
+const RAW_END   = 'endraw';
 
 /**
- * Escape a code block's content using raw blocks
- *
- * @param {String}
- * @return {String}
+ * Create a templating node.
+ * @param {String} expr
+ * @return {Node}
+ */
+function createTemplatingNode(expr) {
+    return Inline.create({
+        type: INLINES.TEMPLATE,
+        data: {
+            type: 'expr',
+            text: expr
+        }
+    });
+}
+
+/**
+ * Escape a code block.
+ * @param  {Block} block
+ * @return {Array<Node>} blocks
+ */
+function escapeCodeBlock(block) {
+    return [
+        Block.create({
+            type: BLOCKS.TEXT,
+            nodes: [
+                createTemplatingNode(RAW_START)
+            ]
+        }),
+        block,
+        Block.create({
+            type: BLOCKS.TEXT,
+            nodes: [
+                createTemplatingNode(RAW_END)
+            ]
+        })
+    ];
+}
+
+
+/**
+ * Escape a text node.
+ * @param  {Text} node
+ * @return {Array<Node>} nodes
  */
-function escape(str) {
-    return RAW_START + str + RAW_END;
+function escapeTextNode(node) {
+    const ranges = node.getRanges();
+
+    const nodes = ranges.reduce((result, range) => {
+        const hasCode = range.marks.some(mark => mark.type == MARKS.CODE);
+        const text = Text.createFromRanges([ range ]);
+
+        if (hasCode) {
+            return result.concat([
+                createTemplatingNode(RAW_START),
+                text,
+                createTemplatingNode(RAW_END)
+            ]);
+        }
+
+        return result.concat([ text ]);
+    }, []);
+
+    return nodes;
 }
 
 
 /**
+ * Annotate a block container.
+ * @param  {Node} parent
+ * @return {Node}
+ */
+function annotateNode(parent, levelRaw) {
+    let { nodes } = parent;
+
+    nodes = nodes.reduce((out, node) => {
+        if (node.type === BLOCKS.CODE) {
+            return out.concat(escapeCodeBlock(node));
+        }
+
+        else if (node.kind == 'text') {
+            return out.concat(escapeTextNode(node));
+        }
+
+        return out.concat([node]);
+    }, []);
+
+    return parent.merge({ nodes });
+}
+
+/**
  * Add templating "raw" to code blocks to
  * avoid nunjucks processing their content.
  *
@@ -25,48 +103,30 @@ function escape(str) {
 function preparePage(src) {
     let levelRaw = 0;
 
-    const fromMD = State.create(markdown);
-    const document = fromMD.deserializeToDocument(src);
+    const state = State.create(markdown);
+    let document = state.deserializeToDocument(src);
 
     document = document.mapDescendants((node) => {
-
-    });
-
-
-    const content  = markdown.toContent(src, {
-        math:     true,
-        template: true
-    });
-
-    const textMarkdown = markdown.toText(content, {
-        annotate(state, raw, token) {
-            const tokenType = token.getType();
-
-            if (tokenType === MarkupIt.ENTITIES.TEMPLATE) {
-                const type = token.getData().get('type');
-                const expr = token.getAsPlainText();
-
-                if (type === 'expr') {
-                    if (expr === 'raw') {
-                        levelRaw = levelRaw + 1;
-                    } else if (expr == 'endraw') {
-                        levelRaw = 0;
-                    }
+        if (node.type === INLINES.TEMPLATE) {
+            const { type, text } = node.data;
+
+            if (type === 'expr') {
+                if (text === 'raw') {
+                    levelRaw = levelRaw + 1;
+                } else if (text == 'endraw') {
+                    levelRaw = 0;
                 }
             }
 
-            if (
-                (tokenType === MarkupIt.BLOCKS.CODE || tokenType === MarkupIt.STYLES.CODE)
-                && levelRaw === 0
-            ) {
-                return escape(raw);
-            }
-
-            return raw;
+            return node;
+        } else if (node.kind !== 'text') {
+            return annotateNode(node, levelRaw);
+        } else {
+            return node;
         }
     });
 
-    return textMarkdown;
+    return state.serializeDocument(document);
 }
 
 module.exports = {