Implement separate infrastructure for formatting and printing output

This fixes various issues around code folding and comments and reduces
duplicated logic for dealing with both.

Author: mcy

.github/workflows/go.yml

@@ -30,8 +30,7 @@ jobs:
         go-version: '^1.18.0'
     - name: go fmt
       run: |
-        go fmt | tee -a modified
-        go fmt ./cmd/protoscope | tee -a modified
+        go fmt . ./cmd/* ./internal/* | tee -a modified
         if [[ $(cat modified) ]]; then exit 1; fi;
 
   tests:

internal/print/print.go

@@ -0,0 +1,306 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://wwp.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// package print contains printing helpers used by the Protoscope disassembler.
+
+package print
+
+import (
+	"bytes"
+	"fmt"
+	"unicode/utf8"
+)
+
+// Stack is a wrapper over a slice type that provides helpers for pushing and
+// popping elements.
+//
+// Exported because of utility in the disassembler itself.
+type Stack[T any] []T
+
+// Push pushes an element.
+func (s *Stack[T]) Push(x T) {
+	*s = append(*s, x)
+}
+
+// Pop pops an element; panics if the stack is empty.
+func (s *Stack[T]) Pop() T {
+	popped := (*s)[len(*s)-1]
+	*s = (*s)[:len(*s)-1]
+	return popped
+}
+
+// Pop pops the top n elements off the stack and returns a slice containing
+// copies. Panics if the stack is too small
+func (s *Stack[T]) PopN(n int) []T {
+	popped := make([]T, n)
+	copy(popped, (*s)[len(*s)-n:])
+	*s = (*s)[:len(*s)-n]
+	return popped
+}
+
+// Peek returns a pointer to the top of the stack, or nil if the stack is
+// empty.
+func (s *Stack[T]) Peek() *T {
+	return &s.PeekN(1)[0]
+}
+
+// Peek returns the top n elements of the stack as another stack.
+//
+// Returns nil if the stack is too small.
+func (s *Stack[T]) PeekN(n int) Stack[T] {
+	if len(*s) < n {
+		return nil
+	}
+	return (*s)[len(*s)-n:]
+}
+
+// Line represents a single line in the output stream. A Printer buffers on a
+// line-by-line basis to be able to do indentation and brace collapse with
+// minimal difficulty.
+type Line struct {
+	// This line's in-progress text buffer.
+	bytes.Buffer
+
+	remarks []string
+	indent  int
+	folds   int
+}
+
+// Printer is an intelligent indentation and codeblock aware printer.
+type Printer struct {
+	// The number of spaces to use per indentation level.
+	Indent int
+	// The number of nested folded blocks allowed, < 0 means infinity.
+	MaxFolds int
+
+	lines  Stack[Line]
+	blocks Stack[BlockInfo]
+}
+
+// Current returns the current line being processed.
+func (p *Printer) Current() *Line {
+	return p.Prev(0)
+}
+
+// Discards the current line
+func (p *Printer) DiscardLine() {
+	p.lines.Pop()
+}
+
+type Mark int
+
+// Makes a mark on the line buffer.
+func (p *Printer) Mark() Mark {
+	return Mark(len(p.lines))
+}
+
+// Discards all lines after the mark.
+func (p *Printer) Reset(m Mark) {
+	p.lines = p.lines[:m]
+}
+
+// Prev returns the nth most recent line.
+//
+// Returns nil if there are not enough lines.
+func (p *Printer) Prev(n int) *Line {
+	return &p.lines.PeekN(n + 1)[0]
+}
+
+// NewLine pushes a new line.
+func (p *Printer) NewLine() {
+	p.lines.Push(Line{})
+}
+
+// Writes to the current line's buffer with Fprint.
+func (p *Printer) Write(args ...any) {
+	fmt.Fprint(p.Current(), args...)
+}
+
+// Writes to the current line's buffer with Fprintf.
+func (p *Printer) Writef(f string, args ...any) {
+	fmt.Fprintf(p.Current(), f, args...)
+}
+
+// Adds a new remark made from stringifying args.
+func (p *Printer) Remark(args ...any) {
+	l := p.Current()
+	l.remarks = append(l.remarks, fmt.Sprint(args...))
+}
+
+// Adds a new remark made from stringifying args.
+func (p *Printer) Remarkf(f string, args ...any) {
+	l := p.Current()
+	l.remarks = append(l.remarks, fmt.Sprintf(f, args...))
+}
+
+// Finish dumps the entire contents of the Printer into a byte array.
+func (p *Printer) Finish() []byte {
+	if len(p.blocks) != 0 {
+		panic("called Finish() without closing all blocks")
+	}
+
+	var out bytes.Buffer
+	indent := 0
+	commentCol := -1
+	commentColUntil := -1
+	for i, line := range p.lines {
+		if len(line.remarks) != 0 && commentColUntil < i {
+			// Comments are aligned to the same column if they are contiguous, unless
+			// crossing an indentation boundary would cause the remark column to be
+			// further than it would have been without crossing the boundary.
+			//
+			// This allows the column finding algorithm to be linear.
+			indent2 := indent
+			commentCol = -1
+			for j, line := range p.lines[i:] {
+				if len(line.remarks) == 0 {
+					commentColUntil = j + i
+					break
+				}
+
+				lineLen := indent2*p.Indent + utf8.RuneCount(line.Bytes())
+				indent2 += line.indent
+				if lineLen > commentCol {
+					if j > 1 && line.indent != 0 {
+						commentColUntil = j + i
+						break
+					}
+					commentCol = lineLen
+				}
+			}
+			if extra := commentCol % p.Indent; extra != 0 {
+				commentCol += p.Indent - extra
+			}
+		}
+
+		for i := 0; i < indent*p.Indent; i++ {
+			out.WriteString(" ")
+		}
+
+		out.Write(line.Bytes())
+		if len(line.remarks) > 0 {
+			needed := commentCol - indent*p.Indent - line.Len()
+			for i := 0; i < needed; i++ {
+				out.WriteString(" ")
+			}
+
+			out.WriteString("  # ")
+			for i, remark := range line.remarks {
+				if i != 0 {
+					out.WriteString(", ")
+				}
+				out.WriteString(remark)
+			}
+		}
+
+		indent += line.indent
+		out.WriteString("\n")
+	}
+
+	return out.Bytes()
+}
+
+type BlockInfo struct {
+	// Whether this block will start and end with delimiters that do not need to
+	// have spaces placed before/after them, allowing for output like {x} instead
+	// of { x }.
+	HasDelimiters bool
+	// The maximum height of the block that will be folded into a single line.
+	HeightToFoldAt int
+	// The line (zero-indexed, starting from the last line) that should be the
+	// final indented line. If there are not enough lines, the block is not
+	// indented at all.
+	UnindentAt int
+
+	start int
+}
+
+// Starts a new indentation block.
+func (p *Printer) StartBlock(bi BlockInfo) {
+	if p.Current().indent != 0 {
+		panic("called StartBlock() too many times; this is a bug")
+	}
+	bi.start = len(p.lines) - 1
+	p.blocks.Push(bi)
+	p.Current().indent++
+}
+
+// Discards the current block and undoes its indentation.
+func (p *Printer) DropBlock() *Line {
+	bi := p.blocks.Pop()
+	start := &p.lines[bi.start]
+	start.indent--
+	return start
+}
+
+// Finishes an indentation block; a call to this function must match up to
+// a corresponding previous StartBlock() call. Returns the starting line for the
+// block.
+//
+// This function will perform folding of small blocks as appropriate.
+func (p *Printer) EndBlock() *Line {
+	bi := p.blocks.Pop()
+	start := &p.lines[bi.start]
+	height := len(p.lines) - bi.start
+
+	// Does the unindentation operation. Because this may run after a successful
+	// fold, we need to make sure that it re-computes the height.
+	defer func() {
+		height := len(p.lines) - bi.start
+		if height <= bi.UnindentAt {
+			p.lines[bi.start].indent--
+		} else {
+			p.lines.PeekN(bi.UnindentAt + 1)[0].indent--
+		}
+	}()
+
+	// Decide whether to fold this block.
+	if height > bi.HeightToFoldAt || height < 2 {
+		return start
+	}
+
+	folds := 0
+	remarks := 0
+	for _, line := range p.lines[bi.start:] {
+		folds += line.folds
+		if len(line.remarks) > 0 {
+			remarks++
+		}
+	}
+
+	if folds > p.MaxFolds {
+		return start
+	}
+
+	// Do not mix remarks from different lines.
+	if remarks > 1 {
+		return start
+	}
+
+	// We are ok to unindent.
+	for i, line := range p.lines[bi.start+1:] {
+		if (i != 0 && i != height-2) || !bi.HasDelimiters {
+			start.WriteString(" ")
+		}
+		start.Write(line.Bytes())
+		if len(line.remarks) != 0 {
+			// This will execute at most once per loop.
+			start.remarks = line.remarks
+		}
+	}
+
+	start.folds = folds
+	p.lines = p.lines[:bi.start+1]
+	return start
+}

testdata/explicit-wire-types.pb.golden

@@ -9,8 +9,8 @@
 8:I64 108i64
 9:I32 109i32
 10:I64 110i64
-11:I32 111.0i32  # 0x42de0000i32
-12:I64 112.0  # 0x405c000000000000i64
+11:I32 111.0i32   # 0x42de0000i32
+12:I64 112.0      # 0x405c000000000000i64
 13:VARINT 1
 14:LEN {"115"}
 15:LEN {"116"}
@@ -48,10 +48,10 @@
 39:I32 309i32
 40:I64 210i64
 40:I64 310i64
-41:I32 211.0i32  # 0x43530000i32
-41:I32 311.0i32  # 0x439b8000i32
-42:I64 212.0  # 0x406a800000000000i64
-42:I64 312.0  # 0x4073800000000000i64
+41:I32 211.0i32   # 0x43530000i32
+41:I32 311.0i32   # 0x439b8000i32
+42:I64 212.0      # 0x406a800000000000i64
+42:I64 312.0      # 0x4073800000000000i64
 43:VARINT 1
 43:VARINT 0
 44:LEN {"215"}
@@ -92,8 +92,8 @@
 68:I64 408i64
 69:I32 409i32
 70:I64 410i64
-71:I32 411.0i32  # 0x43cd8000i32
-72:I64 412.0  # 0x4079c00000000000i64
+71:I32 411.0i32   # 0x43cd8000i32
+72:I64 412.0      # 0x4079c00000000000i64
 73:VARINT 0
 74:LEN {"415"}
 75:LEN {"416"}

testdata/groups.pb.golden

@@ -9,7 +9,10 @@
 4:EGROUP
 5: !{6: !{}}
 6: !{long-form:5}
-7: !{1: 1 long-form:5}
+7: !{
+  1: 1
+  long-form:5
+}
 7: !{
   1: 1
   1: 1