zs/d2.go

package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os/exec"
	"strings"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"github.com/yuin/goldmark"
	"github.com/yuin/goldmark/ast"
	"github.com/yuin/goldmark/parser"
	"github.com/yuin/goldmark/renderer"
	"github.com/yuin/goldmark/text"
	"github.com/yuin/goldmark/util"
)

// d2Timeout bounds a single diagram render so a wedged d2 process cannot hang
// the whole build.
const d2Timeout = 30 * time.Second

// d2Lang is the fenced code block info string handled by this extension.
var d2Lang = []byte("d2")

// errNoD2 reports that the d2 binary is not installed, which is warned about
// once per process rather than once per diagram.
var errNoD2 = errors.New("not found in $PATH")

// warnNoD2 keeps the "d2 is not installed" notice to once per process: in watch
// mode every rebuild would otherwise repeat it for every diagram.
var warnNoD2 sync.Once

// D2Extender renders ```d2 fenced code blocks as inline SVG by piping them
// through the d2 binary (https://d2lang.com). When d2 is not installed the
// block is rendered as an ordinary code block, so a site still builds on a
// machine without it.
//
// zs shells out rather than linking d2 as a library: the library accounted for
// 14MB of a 31MB binary, and pulled in a JavaScript engine, a plotting package
// and a PDF writer to do it.
type D2Extender struct{}

// Extend implements goldmark.Extender.
func (e *D2Extender) Extend(m goldmark.Markdown) {
	m.Parser().AddOptions(parser.WithASTTransformers(
		util.Prioritized(&d2Transformer{}, 100),
	))
	m.Renderer().AddOptions(renderer.WithNodeRenderers(
		util.Prioritized(&d2Renderer{}, 0),
	))
}

// d2Block is a fenced code block tagged as d2, replaced in the AST so that the
// default (and highlighting) renderers leave every other fence alone.
type d2Block struct {
	ast.BaseBlock
}

var kindD2Block = ast.NewNodeKind("D2Block")

// Kind implements ast.Node.
func (n *d2Block) Kind() ast.NodeKind { return kindD2Block }

// Dump implements ast.Node.
func (n *d2Block) Dump(source []byte, level int) {
	ast.DumpHelper(n, source, level, nil, nil)
}

// d2Transformer swaps ```d2 fenced code blocks for d2Block nodes.
type d2Transformer struct{}

// Transform implements parser.ASTTransformer.
func (t *d2Transformer) Transform(doc *ast.Document, reader text.Reader, _ parser.Context) {
	var blocks []*ast.FencedCodeBlock

	// Collect first: replacing nodes while walking would modify the tree
	// underneath the walk.
	ast.Walk(doc, func(node ast.Node, entering bool) (ast.WalkStatus, error) {
		if !entering {
			return ast.WalkContinue, nil
		}
		cb, ok := node.(*ast.FencedCodeBlock)
		if !ok {
			return ast.WalkContinue, nil
		}
		if !bytes.Equal(cb.Language(reader.Source()), d2Lang) {
			return ast.WalkContinue, nil
		}
		blocks = append(blocks, cb)
		return ast.WalkContinue, nil
	})

	for _, cb := range blocks {
		b := new(d2Block)
		b.SetLines(cb.Lines())
		if parent := cb.Parent(); parent != nil {
			parent.ReplaceChild(parent, cb, b)
		}
	}
}

// d2Renderer renders d2Block nodes.
type d2Renderer struct{}

// RegisterFuncs implements renderer.NodeRenderer.
func (r *d2Renderer) RegisterFuncs(reg renderer.NodeRendererFuncRegisterer) {
	reg.Register(kindD2Block, r.render)
}

func (r *d2Renderer) render(w util.BufWriter, src []byte, node ast.Node, entering bool) (ast.WalkStatus, error) {
	if !entering {
		return ast.WalkContinue, nil
	}

	n := node.(*d2Block)
	var buf bytes.Buffer
	lines := n.Lines()
	for i := 0; i < lines.Len(); i++ {
		line := lines.At(i)
		buf.Write(line.Value(src))
	}
	source := buf.Bytes()

	if len(bytes.TrimSpace(source)) == 0 {
		return ast.WalkSkipChildren, nil
	}

	svg, err := renderD2(source)
	if err != nil {
		// A diagram that cannot be rendered must not fail the page: show the
		// source so the content is still readable and the problem is visible.
		// A missing d2 binary has already been reported, once.
		if !errors.Is(err, errNoD2) {
			log.WithError(err).Warn("error rendering d2 diagram, falling back to a code block")
		}
		writeD2Source(w, source)
		return ast.WalkSkipChildren, nil
	}

	w.WriteString(`<div class="d2">`)
	w.Write(svg)
	w.WriteString("</div>\n")

	return ast.WalkSkipChildren, nil
}

// renderD2 renders d2 source to SVG with the d2 binary. The command and any
// extra arguments (themes, layout engine, sketch mode) are configurable via
// the d2.command and d2.args keys.
func renderD2(source []byte) ([]byte, error) {
	name := viper.GetString("d2.command")
	if name == "" {
		name = "d2"
	}

	if _, err := exec.LookPath(name); err != nil {
		warnNoD2.Do(func() {
			log.Warnf("%s not found in $PATH: d2 diagrams will render as code blocks", name)
		})
		return nil, fmt.Errorf("%s: %w", name, errNoD2)
	}

	ctx, cancel := context.WithTimeout(context.Background(), d2Timeout)
	defer cancel()

	// "-" reads the diagram from stdin and writes the SVG to stdout.
	args := append([]string{}, viper.GetStringSlice("d2.args")...)
	args = append(args, "-", "-")

	var out, errs bytes.Buffer
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.Stdin = bytes.NewReader(source)
	cmd.Stdout = &out
	cmd.Stderr = &errs
	if err := cmd.Run(); err != nil {
		if msg := strings.TrimSpace(errs.String()); msg != "" {
			return nil, fmt.Errorf("%s: %w: %s", name, err, msg)
		}
		return nil, fmt.Errorf("%s: %w", name, err)
	}

	return stripXMLProlog(out.Bytes()), nil
}

// stripXMLProlog removes the XML declaration the d2 CLI emits, which is not
// valid partway through an HTML document.
func stripXMLProlog(b []byte) []byte {
	b = bytes.TrimSpace(b)
	if !bytes.HasPrefix(b, []byte("<?xml")) {
		return b
	}
	if i := bytes.Index(b, []byte("?>")); i != -1 {
		return bytes.TrimSpace(b[i+len("?>"):])
	}
	return b
}

// writeD2Source writes the diagram source as a code block, matching what
// goldmark renders for a fenced block it has no special handling for.
func writeD2Source(w util.BufWriter, source []byte) {
	w.WriteString(`<pre><code class="language-d2">`)
	w.Write(util.EscapeHTML(source))
	w.WriteString("</code></pre>\n")
}