pretty: bring the code closer to Wadler's code

The Go code was mixing up the two types Doc and DOC from Philip
Wadler's prettier printer. This patch splits them, which opens an
optimization avenue.

Release note: None

Author: knz

pkg/util/pretty/document.go

@@ -28,13 +28,19 @@
 // side. The paper then describes various performance improvements that reduce
 // the search space of the best function such that it can complete in O(n)
 // instead of O(n^2) time, where n is the number of nodes.
+//
+// For example code with SQL to experiment further, refer to
+// https://github.com/knz/prettier/
+//
 package pretty
 
 import (
 	"fmt"
 )
 
-// Doc represents a document as described by the referenced paper.
+// Doc represents a document as described by the type "DOC" in the
+// referenced paper. This is the abstract representation constructed
+// by the pretty-printing code.
 type Doc interface {
 	// All Docs can uniquely convert themselves into a string so they can be
 	// memoized during better calculation.
@@ -48,81 +54,77 @@ func (nilDoc) isDoc() {}
 func (concat) isDoc() {}
 func (nest) isDoc()   {}
 func (union) isDoc()  {}
-func (textX) isDoc()  {}
-func (lineX) isDoc()  {}
 
-func (d text) String() string   { return fmt.Sprintf("(%q)", string(d)) }
+func (d text) String() string   { return fmt.Sprintf("(TEXT %q)", string(d)) }
 func (line) String() string     { return "LINE" }
 func (nilDoc) String() string   { return "NIL" }
 func (d concat) String() string { return fmt.Sprintf("(%s :<> %s)", d.a, d.b) }
 func (d nest) String() string   { return fmt.Sprintf("(NEST %d %s)", d.n, d.d) }
 func (d union) String() string  { return fmt.Sprintf("(%s :<|> %s)", d.x, d.y) }
-func (d textX) String() string  { return fmt.Sprintf("(%q TEXTX %s)", d.s, d.d) }
-func (d lineX) String() string  { return fmt.Sprintf("(%q LINEX %s)", d.s, d.d) }
 
+//
+// Implementations of Doc ("DOC" in paper).
+//
+
+// nilDoc represents NIL :: DOC -- the empty doc.
 type nilDoc struct{}
 
-// Nil is the empty Doc.
+// Nil is the NIL constructor.
 var Nil nilDoc
 
+// text represents (TEXT s) :: DOC -- a simple text string.
 type text string
 
-// Text creates a string Doc.
+// Text is the TEXT constructor.
 func Text(s string) Doc {
 	return text(s)
 }
 
+// line represents LINE :: DOC -- a "soft line" that can be flattened to a space.
 type line struct{}
 
-// Line is either a newline or a space if flattened onto a single line.
+// Line is the LINE constructor.
 var Line line
 
+// concat represents (DOC <> DOC) :: DOC -- the concatenation of two docs.
 type concat struct {
 	a, b Doc
 }
 
-// fnNotNil runs returns fn(a, b). nil (the Go value) is converted to Nil (the
-// Doc). If either Doc is Nil, the other Doc is returned without invoking fn.
-func concatFn(a, b Doc, fn func(Doc, Doc) Doc) Doc {
-	if a == nil {
-		a = Nil
-	}
-	if b == nil {
-		b = Nil
-	}
-	if a == Nil {
-		return b
-	}
-	if b == Nil {
-		return a
-	}
-	return fn(a, b)
+// Concat is the <> constructor.
+// This uses simplifyNil to avoid actually inserting NIL docs
+// in the abstract tree.
+func Concat(a, b Doc) Doc {
+	return simplifyNil(a, b, func(a, b Doc) Doc { return concat{a, b} })
 }
 
-// Concat concatenates two Docs.
-func Concat(a, b Doc) Doc {
-	return concatFn(a, b, func(a, b Doc) Doc {
-		return concat{a, b}
-	})
+// nest represents (NEST Int DOC) :: DOC -- nesting a doc "under" another.
+// NEST indents d by s at effective length n. len(s) does not have to be
+// n. This allows s to be a tab character and n can be a tab width.
+type nest struct {
+	n int
+	s string
+	d Doc
 }
 
-// Fill concatenates a list of docs with spaces if possible,
-// otherwise with newlines.
-// See linked paper pp 14-15.
-// func Fill(d ...Doc) Doc {
-// 	switch len(d) {
-// 	case 0:
-// 		return Nil
-// 	case 1:
-// 		return d[0]
-// 	default:
-// 		rest := Fill(d[1:]...)
-// 		return union{
-// 			ConcatSpace(flatten(d[0]), rest),
-// 			ConcatLine(d[0], rest),
-// 		}
-// 	}
-// }
+// Nest is the NEST constructor.
+func Nest(n int, s string, d Doc) Doc {
+	return nest{n, s, d}
+}
+
+// union represents (DOC <|> DOC) :: DOC -- the union of two docs.
+// <|> is really the union of two sets of layouts. x and y must flatten to the
+// same layout. Additionally, no first line of a document in x is shorter
+// than some first line of a document in y; or, equivalently, every first
+// line in x is at least as long as every first line in y.
+//
+// The main use of the union is via the Group operator defined below.
+//
+// We do not provide a public constructor as this type is not
+// exported.
+type union struct {
+	x, y Doc
+}
 
 // Group will format d on one line if possible.
 func Group(d Doc) Doc {
@@ -149,35 +151,3 @@ func flatten(d Doc) Doc {
 		panic(fmt.Errorf("unknown type: %T", d))
 	}
 }
-
-type nest struct {
-	n int
-	s string
-	d Doc
-}
-
-// Nest indents d by s at effective length n. len(s) does not have to be
-// n. This allows s to be a tab character and n can be a tab width.
-func Nest(n int, s string, d Doc) Doc {
-	return nest{n, s, d}
-}
-
-// Doc types below are not directly accessible to users.
-
-// union is the union of two sets of layouts. x and y must flatten to the
-// same layout. Additionally, no first line of a document in x is shorter
-// than some first line of a document in y; or, equivalently, every first
-// line in x is at least as long as every first line in y.
-type union struct {
-	x, y Doc
-}
-
-type textX struct {
-	s string
-	d Doc
-}
-
-type lineX struct {
-	s string
-	d Doc
-}

pkg/util/pretty/pretty.go

@@ -24,6 +24,37 @@ import (
 // of the below code. Methods, variables, and implementation details were
 // made to resemble it as close as possible.
 
+// docBest represents a selected document as described by the type
+// "Doc" in the referenced paper (not "DOC"). This is the
+// less-abstract representation constructed during "best layout"
+// selection.
+type docBest interface {
+	String() string
+	isDocBest()
+}
+
+func (nilDocB) String() string { return "Nil" }
+func (d textB) String() string { return fmt.Sprintf("(%q `Text` %s)", d.s, d.d) }
+func (d lineB) String() string { return fmt.Sprintf("(%q `Line` %s)", d.s, d.d) }
+
+func (nilDocB) isDocBest() {}
+func (textB) isDocBest()   {}
+func (lineB) isDocBest()   {}
+
+type nilDocB struct{}
+
+var nilB nilDocB
+
+type textB struct {
+	s string
+	d docBest
+}
+
+type lineB struct {
+	s string
+	d docBest
+}
+
 // Pretty returns a pretty-printed string for the Doc d at line length n.
 func Pretty(d Doc, n int) string {
 	var sb strings.Builder
@@ -33,10 +64,10 @@ func Pretty(d Doc, n int) string {
 }
 
 // w is the max line width.
-func best(w int, x Doc) Doc {
+func best(w int, x Doc) docBest {
 	b := beExec{
 		w:     w,
-		cache: make(map[cacheKey]Doc),
+		cache: make(map[cacheKey]docBest),
 	}
 	return b.be(0, iDoc{0, "", x})
 }
@@ -59,13 +90,13 @@ type cacheKey struct {
 type beExec struct {
 	w int
 	// cache is a memoized cache used during better calculation.
-	cache map[cacheKey]Doc
+	cache map[cacheKey]docBest
 	buf   bytes.Buffer
 }
 
-func (b beExec) be(k int, x ...iDoc) Doc {
+func (b beExec) be(k int, x ...iDoc) docBest {
 	if len(x) == 0 {
-		return Nil
+		return nilB
 	}
 	d := x[0]
 	z := x[1:]
@@ -82,12 +113,12 @@ func (b beExec) be(k int, x ...iDoc) Doc {
 		}
 		return b.be(k, x...)
 	case text:
-		return textX{
+		return textB{
 			s: string(t),
 			d: b.be(k+len(t), z...),
 		}
 	case line:
-		return lineX{
+		return lineB{
 			s: d.s,
 			d: b.be(d.i, z...),
 		}
@@ -113,7 +144,7 @@ func (b beExec) be(k int, x ...iDoc) Doc {
 		n := append([]iDoc{{d.i, d.s, t.x}}, z...)
 		res := better(b.w, k,
 			b.be(k, n...),
-			func() Doc {
+			func() docBest {
 				n[0].d = t.y
 				return b.be(k, n...)
 			},
@@ -125,37 +156,37 @@ func (b beExec) be(k int, x ...iDoc) Doc {
 	}
 }
 
-func better(w, k int, x Doc, y func() Doc) Doc {
+func better(w, k int, x docBest, y func() docBest) docBest {
 	if fits(w-k, x) {
 		return x
 	}
 	return y()
 }
 
-func fits(w int, x Doc) bool {
+func fits(w int, x docBest) bool {
 	if w < 0 {
 		return false
 	}
 	switch t := x.(type) {
-	case nilDoc:
+	case nilDocB:
 		return true
-	case textX:
+	case textB:
 		return fits(w-len(t.s), t.d)
-	case lineX:
+	case lineB:
 		return true
 	default:
 		panic(fmt.Errorf("unknown type: %T", x))
 	}
 }
 
-func layout(sb *strings.Builder, d Doc) {
+func layout(sb *strings.Builder, d docBest) {
 	switch d := d.(type) {
-	case nilDoc:
+	case nilDocB:
 		// ignore
-	case textX:
+	case textB:
 		sb.WriteString(d.s)
 		layout(sb, d.d)
-	case lineX:
+	case lineB:
 		sb.WriteString("\n")
 		sb.WriteString(d.s)
 		layout(sb, d.d)

pkg/util/pretty/util.go

@@ -62,7 +62,7 @@ func JoinNestedRight(n int, s string, sep Doc, nested ...Doc) Doc {
 
 // ConcatLine concatenates two Docs with a Line.
 func ConcatLine(a, b Doc) Doc {
-	return concatFn(a, b, func(a, b Doc) Doc {
+	return simplifyNil(a, b, func(a, b Doc) Doc {
 		return Concat(
 			a,
 			Concat(
@@ -75,7 +75,7 @@ func ConcatLine(a, b Doc) Doc {
 
 // ConcatSpace concatenates two Docs with a space.
 func ConcatSpace(a, b Doc) Doc {
-	return concatFn(a, b, func(a, b Doc) Doc {
+	return simplifyNil(a, b, func(a, b Doc) Doc {
 		return Concat(
 			a,
 			Concat(
@@ -146,3 +146,21 @@ func Bracket(n int, s string, l string, x Doc, r string) Doc {
 	)
 	return union{flatten(a), b}
 }
+
+// simplifyNil returns fn(a, b). nil (the Go value) is converted to Nil (the
+// Doc). If either Doc is Nil, the other Doc is returned without invoking fn.
+func simplifyNil(a, b Doc, fn func(Doc, Doc) Doc) Doc {
+	if a == nil {
+		a = Nil
+	}
+	if b == nil {
+		b = Nil
+	}
+	if a == Nil {
+		return b
+	}
+	if b == Nil {
+		return a
+	}
+	return fn(a, b)
+}