PgTidy/pkg/format/doc.go

package format

import "strings"

// Doc is a Wadler/Prettier-style layout document. Build one from Text, Line,
// SoftLine, Group, Indent, IfBreak and Concat, then Render it at a line width:
// a Group is printed flat (Lines become spaces, SoftLines vanish) when it fits
// in the remaining width, otherwise broken (Lines become newlines).
//
// Rendering is relative to column 0; callers that splice the result mid-line
// re-indent continuation lines themselves, so width checks ignore that offset.
type Doc interface{ isDoc() }

type (
	docText   string
	docLine   struct{ soft, hard bool }
	docConcat []Doc
	docIndent struct{ d Doc }
	docGroup  struct{ d Doc }
	docBreak  struct{ broken, flat Doc }
)

func (docText) isDoc()   {}
func (docLine) isDoc()   {}
func (docConcat) isDoc() {}
func (docIndent) isDoc() {}
func (docGroup) isDoc()  {}
func (docBreak) isDoc()  {}

// Text is literal text. If it contains a newline the enclosing group can never
// be printed flat.
func Text(s string) Doc { return docText(s) }

// Line is a space when flat and a newline (plus indentation) when broken.
func Line() Doc { return docLine{} }

// SoftLine is nothing when flat and a newline (plus indentation) when broken.
func SoftLine() Doc { return docLine{soft: true} }

// HardLine is always a newline (plus indentation), even inside a flat group; a
// group containing one can never be printed flat.
func HardLine() Doc { return docLine{hard: true} }

// Lines turns multi-line text into Text pieces joined by HardLines so that each
// line picks up the surrounding Indent. Empty lines carry no indentation.
func Lines(s, nl string) Doc {
	parts := strings.Split(s, nl)
	ds := make([]Doc, 0, 2*len(parts))
	for i, p := range parts {
		if i > 0 {
			ds = append(ds, HardLine())
		}
		ds = append(ds, Text(p))
	}
	return Concat(ds...)
}

// Concat joins docs in order.
func Concat(ds ...Doc) Doc { return docConcat(ds) }

// Indent indents every line break inside d by one indent unit.
func Indent(d Doc) Doc { return docIndent{d} }

// Group lays d out flat if it fits on the current line, otherwise broken.
func Group(d Doc) Doc { return docGroup{d} }

// IfBreak renders broken when the enclosing group is broken, flat otherwise.
func IfBreak(broken, flat Doc) Doc { return docBreak{broken, flat} }

type docCmd struct {
	indent string
	flat   bool
	d      Doc
}

// Render lays out d. width <= 0 means unlimited (every group stays flat).
func Render(d Doc, width int, unit, nl string) string {
	var b strings.Builder
	col := 0
	pending := "" // indentation owed to the next non-empty text
	stack := []docCmd{{"", false, d}}
	for len(stack) > 0 {
		c := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		switch v := c.d.(type) {
		case docText:
			if len(v) > 0 {
				b.WriteString(pending)
				pending = ""
			}
			b.WriteString(string(v))
			if i := strings.LastIndexByte(string(v), '\n'); i >= 0 {
				col = len(v) - i - 1
			} else {
				col += len(v)
			}
		case docConcat:
			for i := len(v) - 1; i >= 0; i-- {
				stack = append(stack, docCmd{c.indent, c.flat, v[i]})
			}
		case docIndent:
			stack = append(stack, docCmd{c.indent + unit, c.flat, v.d})
		case docBreak:
			if c.flat {
				stack = append(stack, docCmd{c.indent, c.flat, v.flat})
			} else {
				stack = append(stack, docCmd{c.indent, c.flat, v.broken})
			}
		case docLine:
			if c.flat && !v.hard {
				if !v.soft {
					b.WriteByte(' ')
					col++
				}
				break
			}
			b.WriteString(nl)
			pending = c.indent
			col = len(c.indent)
		case docGroup:
			flat := c.flat || width <= 0 || fitsFlat(v.d, width-col, c.indent, unit, stack)
			stack = append(stack, docCmd{c.indent, flat, v.d})
		}
	}
	return b.String()
}

// fitsFlat reports whether d, printed flat, plus whatever follows it up to the
// next possible line break, fits in rem columns.
func fitsFlat(d Doc, rem int, indent, unit string, rest []docCmd) bool {
	if rem < 0 {
		return false
	}
	type item struct {
		d    Doc
		flat bool
	}
	work := []item{{d, true}}
	ri := len(rest) - 1
	for rem >= 0 {
		if len(work) == 0 {
			if ri < 0 {
				return true
			}
			work = append(work, item{rest[ri].d, rest[ri].flat})
			ri--
			continue
		}
		it := work[len(work)-1]
		work = work[:len(work)-1]
		switch v := it.d.(type) {
		case docText:
			if strings.ContainsRune(string(v), '\n') {
				return false
			}
			rem -= len(v)
		case docConcat:
			for i := len(v) - 1; i >= 0; i-- {
				work = append(work, item{v[i], it.flat})
			}
		case docIndent:
			work = append(work, item{v.d, it.flat})
		case docGroup:
			work = append(work, item{v.d, it.flat})
		case docBreak:
			if it.flat {
				work = append(work, item{v.flat, it.flat})
			} else {
				work = append(work, item{v.broken, it.flat})
			}
		case docLine:
			if v.hard {
				return false
			}
			if !it.flat {
				return true // a real line break ends the measured run
			}
			if !v.soft {
				rem--
			}
		}
	}
	return false
}