use transformer and renderer to correctly parse and render shortcode …

…blocks

extensions/shortcode/node.go

@@ -25,3 +25,14 @@ func (s *ShortCodeNode) Dump(source []byte, level int) {
 func (h *ShortCodeNode) Kind() ast.NodeKind {
 	return KindShortCode
 }
+
+var KindShortCodeBlock = ast.NewNodeKind("ShortCodeBlock")
+
+type ShortCodeBlock struct {
+	ast.FencedCodeBlock
+	fun ShortCodeFunc
+}
+
+func (s *ShortCodeBlock) Kind() ast.NodeKind {
+	return KindShortCodeBlock
+}

extensions/shortcode/renderer.go

@@ -17,6 +17,7 @@ type shortCodeRenderer struct{}
 
 func (s *shortCodeRenderer) RegisterFuncs(reg renderer.NodeRendererFuncRegisterer) {
 	reg.Register(KindShortCode, s.render)
+	reg.Register(KindShortCodeBlock, s.renderBlock)
 }
 
 func (s *shortCodeRenderer) render(w util.BufWriter, source []byte, n ast.Node, entering bool) (ast.WalkStatus, error) {
@@ -35,3 +36,26 @@ func (s *shortCodeRenderer) render(w util.BufWriter, source []byte, n ast.Node,
 
 	return ast.WalkContinue, nil
 }
+
+func (s *shortCodeRenderer) renderBlock(w util.BufWriter, source []byte, n ast.Node, entering bool) (ast.WalkStatus, error) {
+	if !entering {
+		return ast.WalkContinue, nil
+	}
+
+	node, ok := n.(*ShortCodeBlock)
+	if !ok {
+		return ast.WalkContinue, nil
+	}
+
+	lines := node.Lines()
+	content := ""
+	for i := 0; i < lines.Len(); i++ {
+		line := lines.At(i)
+		content += string(line.Value(source))
+	}
+
+	output := node.fun(Markdown(content))
+	w.Write([]byte(output))
+
+	return ast.WalkContinue, nil
+}

extensions/shortcode/shortcode.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"fmt"
 	"html/template"
-	"regexp"
 
 	. "github.com/emad-elsaid/xlog"
 )
@@ -29,26 +28,6 @@ var shortcodes = map[string]ShortCodeFunc{
 	"alert":   func(c Markdown) template.HTML { return container("is-danger", c) },
 }
 
-func init() {
-	for k, v := range shortcodes {
-		ShortCode(k, v)
-	}
-}
-
 func ShortCode(name string, shortcode ShortCodeFunc) {
 	shortcodes[name] = shortcode
-
-	headerSkip := len("```" + name + "\n")
-	multireg := regexp.MustCompile("(?imUs)^```" + regexp.QuoteMeta(name) + "$(.*)^```$")
-	multilinePreprocessor := func(skip int, v ShortCodeFunc) Preprocessor {
-		return func(c Markdown) Markdown {
-			output := multireg.ReplaceAllStringFunc(string(c), func(i string) string {
-				input := i[skip : len(i)-4]
-				return string(v(Markdown(input)))
-			})
-			return Markdown(output)
-		}
-	}(headerSkip, shortcode)
-
-	RegisterPreprocessor(multilinePreprocessor)
 }

extensions/shortcode/transformer.go

@@ -0,0 +1,58 @@
+package shortcode
+
+import (
+	. "github.com/emad-elsaid/xlog"
+	"github.com/yuin/goldmark/ast"
+	"github.com/yuin/goldmark/parser"
+	"github.com/yuin/goldmark/text"
+	"github.com/yuin/goldmark/util"
+)
+
+func init() {
+	MarkDownRenderer.Parser().AddOptions(
+		parser.WithASTTransformers(
+			util.Prioritized(transformShortCodeBlocks(0), 0),
+		),
+	)
+}
+
+type transformShortCodeBlocks int
+
+func (t transformShortCodeBlocks) Transform(doc *ast.Document, reader text.Reader, pc parser.Context) {
+	source := reader.Source()
+	blocks := []*ast.FencedCodeBlock{}
+
+	ast.Walk(doc, func(node ast.Node, entering bool) (ast.WalkStatus, error) {
+		if !entering {
+			return ast.WalkContinue, nil
+		}
+
+		for c := node.FirstChild(); c != nil; c = c.NextSibling() {
+			n, ok := c.(*ast.FencedCodeBlock)
+			if !ok {
+				continue
+			}
+
+			lang := string(n.Language(source))
+			if _, ok := shortcodes[lang]; !ok {
+				continue
+			}
+
+			blocks = append(blocks, n)
+		}
+
+		return ast.WalkContinue, nil
+	})
+
+	for _, b := range blocks {
+		lang := string(b.Language(source))
+
+		replacement := ShortCodeBlock{
+			FencedCodeBlock: *b,
+			fun:             shortcodes[lang],
+		}
+
+		parent := b.Parent()
+		parent.ReplaceChild(parent, b, &replacement)
+	}
+}