feat: Add code block indentation token support

This change incorporates the specified indentation token when rendering
code blocks. It includes fixes to how the margin and indentation tokens
are rendered globally. When both are specified, the margin is rendered
first as an empty space.

This PR is another attempt at adding the functionality from
#299

ansi/blockelement.go

@@ -36,7 +36,7 @@ func (e *BlockElement) Finish(w io.Writer, ctx RenderContext) error {
 			" ,.;-+|",
 		)
 
-		mw := NewMarginWriter(ctx, w, bs.Current().Style)
+		mw := NewMarginWriter(ctx, w, bs.Current().Style, true)
 		if _, err := io.WriteString(mw, s); err != nil {
 			return err
 		}

ansi/codeblock.go

@@ -1,13 +1,13 @@
 package ansi
 
 import (
+	"bytes"
 	"io"
 	"sync"
 
 	"github.com/alecthomas/chroma/v2"
 	"github.com/alecthomas/chroma/v2/quick"
 	"github.com/alecthomas/chroma/v2/styles"
-	"github.com/muesli/reflow/indent"
 	"github.com/muesli/termenv"
 )
 
@@ -62,82 +62,90 @@ func chromaStyle(style StylePrimitive) string {
 
 func (e *CodeBlockElement) Render(w io.Writer, ctx RenderContext) error {
 	bs := ctx.blockStack
-
-	var indentation uint
-	var margin uint
 	rules := ctx.options.Styles.CodeBlock
-	if rules.Indent != nil {
-		indentation = *rules.Indent
-	}
-	if rules.Margin != nil {
-		margin = *rules.Margin
+
+	be := BlockElement{
+		Block: &bytes.Buffer{},
+		Style: rules.StyleBlock,
 	}
-	theme := rules.Theme
+	bs.Push(be)
+	return nil
+}
+
+func (e *CodeBlockElement) Finish(w io.Writer, ctx RenderContext) error {
+	bs := ctx.blockStack
+	rules := bs.Current().Style
+
+	cb := ctx.options.Styles.CodeBlock
+	theme := cb.Theme
+	chromaRules := cb.Chroma
 
-	if rules.Chroma != nil && ctx.options.ColorProfile != termenv.Ascii {
+	if chromaRules != nil && ctx.options.ColorProfile != termenv.Ascii {
 		theme = chromaStyleTheme
 		mutex.Lock()
 		// Don't register the style if it's already registered.
 		_, ok := styles.Registry[theme]
 		if !ok {
 			styles.Register(chroma.MustNewStyle(theme,
 				chroma.StyleEntries{
-					chroma.Text:                chromaStyle(rules.Chroma.Text),
-					chroma.Error:               chromaStyle(rules.Chroma.Error),
-					chroma.Comment:             chromaStyle(rules.Chroma.Comment),
-					chroma.CommentPreproc:      chromaStyle(rules.Chroma.CommentPreproc),
-					chroma.Keyword:             chromaStyle(rules.Chroma.Keyword),
-					chroma.KeywordReserved:     chromaStyle(rules.Chroma.KeywordReserved),
-					chroma.KeywordNamespace:    chromaStyle(rules.Chroma.KeywordNamespace),
-					chroma.KeywordType:         chromaStyle(rules.Chroma.KeywordType),
-					chroma.Operator:            chromaStyle(rules.Chroma.Operator),
-					chroma.Punctuation:         chromaStyle(rules.Chroma.Punctuation),
-					chroma.Name:                chromaStyle(rules.Chroma.Name),
-					chroma.NameBuiltin:         chromaStyle(rules.Chroma.NameBuiltin),
-					chroma.NameTag:             chromaStyle(rules.Chroma.NameTag),
-					chroma.NameAttribute:       chromaStyle(rules.Chroma.NameAttribute),
-					chroma.NameClass:           chromaStyle(rules.Chroma.NameClass),
-					chroma.NameConstant:        chromaStyle(rules.Chroma.NameConstant),
-					chroma.NameDecorator:       chromaStyle(rules.Chroma.NameDecorator),
-					chroma.NameException:       chromaStyle(rules.Chroma.NameException),
-					chroma.NameFunction:        chromaStyle(rules.Chroma.NameFunction),
-					chroma.NameOther:           chromaStyle(rules.Chroma.NameOther),
-					chroma.Literal:             chromaStyle(rules.Chroma.Literal),
-					chroma.LiteralNumber:       chromaStyle(rules.Chroma.LiteralNumber),
-					chroma.LiteralDate:         chromaStyle(rules.Chroma.LiteralDate),
-					chroma.LiteralString:       chromaStyle(rules.Chroma.LiteralString),
-					chroma.LiteralStringEscape: chromaStyle(rules.Chroma.LiteralStringEscape),
-					chroma.GenericDeleted:      chromaStyle(rules.Chroma.GenericDeleted),
-					chroma.GenericEmph:         chromaStyle(rules.Chroma.GenericEmph),
-					chroma.GenericInserted:     chromaStyle(rules.Chroma.GenericInserted),
-					chroma.GenericStrong:       chromaStyle(rules.Chroma.GenericStrong),
-					chroma.GenericSubheading:   chromaStyle(rules.Chroma.GenericSubheading),
-					chroma.Background:          chromaStyle(rules.Chroma.Background),
+					chroma.Text:                chromaStyle(chromaRules.Text),
+					chroma.Error:               chromaStyle(chromaRules.Error),
+					chroma.Comment:             chromaStyle(chromaRules.Comment),
+					chroma.CommentPreproc:      chromaStyle(chromaRules.CommentPreproc),
+					chroma.Keyword:             chromaStyle(chromaRules.Keyword),
+					chroma.KeywordReserved:     chromaStyle(chromaRules.KeywordReserved),
+					chroma.KeywordNamespace:    chromaStyle(chromaRules.KeywordNamespace),
+					chroma.KeywordType:         chromaStyle(chromaRules.KeywordType),
+					chroma.Operator:            chromaStyle(chromaRules.Operator),
+					chroma.Punctuation:         chromaStyle(chromaRules.Punctuation),
+					chroma.Name:                chromaStyle(chromaRules.Name),
+					chroma.NameBuiltin:         chromaStyle(chromaRules.NameBuiltin),
+					chroma.NameTag:             chromaStyle(chromaRules.NameTag),
+					chroma.NameAttribute:       chromaStyle(chromaRules.NameAttribute),
+					chroma.NameClass:           chromaStyle(chromaRules.NameClass),
+					chroma.NameConstant:        chromaStyle(chromaRules.NameConstant),
+					chroma.NameDecorator:       chromaStyle(chromaRules.NameDecorator),
+					chroma.NameException:       chromaStyle(chromaRules.NameException),
+					chroma.NameFunction:        chromaStyle(chromaRules.NameFunction),
+					chroma.NameOther:           chromaStyle(chromaRules.NameOther),
+					chroma.Literal:             chromaStyle(chromaRules.Literal),
+					chroma.LiteralNumber:       chromaStyle(chromaRules.LiteralNumber),
+					chroma.LiteralDate:         chromaStyle(chromaRules.LiteralDate),
+					chroma.LiteralString:       chromaStyle(chromaRules.LiteralString),
+					chroma.LiteralStringEscape: chromaStyle(chromaRules.LiteralStringEscape),
+					chroma.GenericDeleted:      chromaStyle(chromaRules.GenericDeleted),
+					chroma.GenericEmph:         chromaStyle(chromaRules.GenericEmph),
+					chroma.GenericInserted:     chromaStyle(chromaRules.GenericInserted),
+					chroma.GenericStrong:       chromaStyle(chromaRules.GenericStrong),
+					chroma.GenericSubheading:   chromaStyle(chromaRules.GenericSubheading),
+					chroma.Background:          chromaStyle(chromaRules.Background),
 				}))
 		}
 		mutex.Unlock()
 	}
 
-	iw := indent.NewWriterPipe(w, indentation+margin, func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, " ")
-	})
-
+	mw := NewMarginWriter(ctx, w, bs.Current().Style, false)
+	renderText(mw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockPrefix)
 	if len(theme) > 0 {
-		renderText(iw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockPrefix)
-
-		err := quick.Highlight(iw, e.Code, e.Language, "terminal256", theme)
+		err := quick.Highlight(mw, e.Code, e.Language, "terminal256", theme)
 		if err != nil {
 			return err
 		}
-		renderText(iw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockSuffix)
-		return nil
-	}
+	} else {
+		// fallback rendering
+		el := &BaseElement{
+			Token: e.Code,
+			Style: rules.StylePrimitive,
+		}
 
-	// fallback rendering
-	el := &BaseElement{
-		Token: e.Code,
-		Style: rules.StylePrimitive,
+		err := el.Render(mw, ctx)
+		if err != nil {
+			return err
+		}
 	}
+	renderText(mw, ctx.options.ColorProfile, bs.Current().Style.StylePrimitive, rules.BlockSuffix)
 
-	return el.Render(iw, ctx)
+	bs.Current().Block.Reset()
+	bs.Pop()
+	return nil
 }

ansi/elements.go

@@ -282,12 +282,14 @@ func (tr *ANSIRenderer) NewElement(node ast.Node, source []byte) Element {
 			line := n.Lines().At(i)
 			s += string(line.Value(source))
 		}
+		e := &CodeBlockElement{
+			Code:     s,
+			Language: string(n.Language(source)),
+		}
 		return Element{
 			Entering: "\n",
-			Renderer: &CodeBlockElement{
-				Code:     s,
-				Language: string(n.Language(source)),
-			},
+			Renderer: e,
+			Finisher: e,
 		}
 
 	case ast.KindCodeBlock:
@@ -298,11 +300,13 @@ func (tr *ANSIRenderer) NewElement(node ast.Node, source []byte) Element {
 			line := n.Lines().At(i)
 			s += string(line.Value(source))
 		}
+		e := &CodeBlockElement{
+			Code: s,
+		}
 		return Element{
 			Entering: "\n",
-			Renderer: &CodeBlockElement{
-				Code: s,
-			},
+			Renderer: e,
+			Finisher: e,
 		}
 
 	case ast.KindCodeSpan:

ansi/heading.go

@@ -59,7 +59,7 @@ func (e *HeadingElement) Render(w io.Writer, ctx RenderContext) error {
 func (e *HeadingElement) Finish(w io.Writer, ctx RenderContext) error {
 	bs := ctx.blockStack
 	rules := bs.Current().Style
-	mw := NewMarginWriter(ctx, w, rules)
+	mw := NewMarginWriter(ctx, w, rules, true)
 
 	flow := wordwrap.NewWriter(int(bs.Width(ctx)))
 	_, err := flow.Write(bs.Current().Block.Bytes())

ansi/margin.go

@@ -5,48 +5,80 @@ import (
 
 	"github.com/muesli/reflow/indent"
 	"github.com/muesli/reflow/padding"
+	"github.com/muesli/termenv"
 )
 
 // MarginWriter is a Writer that applies indentation and padding around
 // whatever you write to it.
 type MarginWriter struct {
+	indentation, margin  uint
+	indentPos, marginPos uint
+	indentToken          string
+
+	profile            termenv.Profile
+	rules, parentRules StylePrimitive
+
 	w  io.Writer
 	pw *padding.Writer
 	iw *indent.Writer
 }
 
 // NewMarginWriter returns a new MarginWriter.
-func NewMarginWriter(ctx RenderContext, w io.Writer, rules StyleBlock) *MarginWriter {
+func NewMarginWriter(ctx RenderContext, w io.Writer, rules StyleBlock, padded bool) *MarginWriter {
 	bs := ctx.blockStack
+	mw := &MarginWriter{
+		w:           w,
+		profile:     ctx.options.ColorProfile,
+		rules:       rules.StylePrimitive,
+		parentRules: bs.Parent().Style.StylePrimitive,
+	}
 
-	var indentation uint
-	var margin uint
 	if rules.Indent != nil {
-		indentation = *rules.Indent
+		mw.indentation = *rules.Indent
+		mw.indentToken = " "
+		if rules.IndentToken != nil {
+			mw.indentToken = *rules.IndentToken
+		}
 	}
 	if rules.Margin != nil {
-		margin = *rules.Margin
+		mw.margin = *rules.Margin
 	}
 
-	pw := padding.NewWriterPipe(w, bs.Width(ctx), func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, rules.StylePrimitive, " ")
-	})
-
-	ic := " "
-	if rules.IndentToken != nil {
-		ic = *rules.IndentToken
-	}
-	iw := indent.NewWriterPipe(pw, indentation+margin, func(wr io.Writer) {
-		renderText(w, ctx.options.ColorProfile, bs.Parent().Style.StylePrimitive, ic)
-	})
-
-	return &MarginWriter{
-		w:  w,
-		pw: pw,
-		iw: iw,
+	fwd := mw.w
+	if padded {
+		mw.pw = padding.NewWriterPipe(mw.w, bs.Width(ctx), func(wr io.Writer) {
+			renderText(mw.w, mw.profile, mw.rules, "")
+		})
+		fwd = mw.pw
 	}
+
+	mw.iw = indent.NewWriterPipe(fwd, mw.indentation+(mw.margin*2), mw.indentFunc)
+	return mw
 }
 
 func (w *MarginWriter) Write(b []byte) (int, error) {
 	return w.iw.Write(b)
 }
+
+// indentFunc is called when writing each the margin and indentation tokens.
+// The margin is written first, using an empty space character as the token.
+// The indentation is written next, using the token specified in the rules.
+func (w *MarginWriter) indentFunc(iw io.Writer) {
+	ic := " "
+	switch {
+	case w.margin == 0 && w.indentation == 0:
+		return
+	case w.margin >= 1 && w.indentation == 0:
+		break
+	case w.margin >= 1 && w.marginPos < w.margin:
+		w.marginPos++
+	case w.indentation >= 1 && w.indentPos < w.indentation:
+		w.indentPos++
+		ic = w.indentToken
+		if w.indentPos == w.indentation {
+			w.marginPos = 0
+			w.indentPos = 0
+		}
+	}
+	renderText(w.w, w.profile, w.parentRules, ic)
+}

ansi/paragraph.go

@@ -35,7 +35,7 @@ func (e *ParagraphElement) Finish(w io.Writer, ctx RenderContext) error {
 	bs := ctx.blockStack
 	rules := bs.Current().Style
 
-	mw := NewMarginWriter(ctx, w, rules)
+	mw := NewMarginWriter(ctx, w, rules, true)
 	if len(strings.TrimSpace(bs.Current().Block.String())) > 0 {
 		flow := wordwrap.NewWriter(int(bs.Width(ctx)))
 		flow.KeepNewlines = ctx.options.PreserveNewLines