Add builtin support for truncation

api.go

@@ -57,6 +57,9 @@ type SafeValue = i.SafeValue
 // TODO(knz): Remove this.
 type SafeMessager = i.SafeMessager
 
+// TruncatedValue represents a value that is truncated when printed.
+type TruncatedValue = i.TruncatedValue
+
 // SafePrinter is a stateful helper that abstracts an output stream in
 // the context of printf-like formatting, but with the ability to
 // separate safe and unsafe bits of data.
@@ -141,6 +144,11 @@ func Unsafe(a interface{}) interface{} { return w.Unsafe(a) }
 // The implementation is also slow.
 func Safe(a interface{}) SafeValue { return w.Safe(a) }
 
+// Truncate limits the rune length of the outputted value when
+// passed into the formatter. The excess runes are replaced
+// with the marker `…`.
+func Truncate(a interface{}, len int) TruncatedValue { return TruncatedValue{Value: a, Length: len} }
+
 // RegisterRedactErrorFn registers an error redaction function for use
 // during automatic redaction by this package.
 // Provided e.g. by cockroachdb/errors.

interfaces/interfaces.go

@@ -147,3 +147,9 @@ type SafeValue interface {
 type SafeMessager = interface {
 	SafeMessage() string
 }
+
+// TruncatedValue represents a value that is truncated when printed.
+type TruncatedValue struct {
+	Value  interface{}
+	Length int
+}

internal/buffer/buffer.go

@@ -31,6 +31,12 @@ type Buffer struct {
 	validUntil int // exclusive upper bound of data that's already validated
 	mode       OutputMode
 	markerOpen bool
+
+	// These fields relate to truncation and are used to enforce the limit of
+	// runes in the buffer. Markers do not contribute to the rune count.
+	runeCount int
+	runeLimit int
+	excess    bool
 }
 
 // OutputMode determines how writes are processed in the Buffer.
@@ -128,28 +134,96 @@ func (b *Buffer) Cap() int {
 // needed. The return value n is the length of p; err is always nil. If the
 // buffer becomes too large, Write will panic with ErrTooLarge.
 func (b *Buffer) Write(p []byte) (n int, err error) {
+	length := len(p)
+	if b.runeLimit != 0 {
+		if b.runeLimit == b.runeCount {
+			if len(p) != 0 {
+				b.excess = true
+			}
+			return
+		}
+
+		runes, end, excess := limitRunes(p, b.runeLimit-b.runeCount)
+		b.runeCount += runes
+		if end < length {
+			length = end
+		}
+		b.excess = excess
+	}
+
 	b.startWrite()
-	m, ok := b.tryGrowByReslice(len(p))
+	m, ok := b.tryGrowByReslice(length)
 	if !ok {
-		m = b.grow(len(p))
+		m = b.grow(length)
 	}
-	return copy(b.buf[m:], p), nil
+	return copy(b.buf[m:], p[:length]), nil
 }
 
 // WriteString appends the contents of s to the buffer, growing the buffer as
 // needed. The return value n is the length of s; err is always nil. If the
 // buffer becomes too large, WriteString will panic with ErrTooLarge.
 func (b *Buffer) WriteString(s string) (n int, err error) {
+	length := len(s)
+	if b.runeLimit != 0 {
+		if b.runeLimit == b.runeCount {
+			if len(s) != 0 {
+				b.excess = true
+			}
+			return
+		}
+
+		runes, end, excess := limitRunesInString(s, b.runeLimit-b.runeCount)
+		b.runeCount += runes
+		if end < length {
+			length = end
+		}
+		b.excess = excess
+	}
+
 	b.startWrite()
-	m, ok := b.tryGrowByReslice(len(s))
+	m, ok := b.tryGrowByReslice(length)
 	if !ok {
-		m = b.grow(len(s))
+		m = b.grow(length)
 	}
-	return copy(b.buf[m:], s), nil
+	return copy(b.buf[m:], s[:length]), nil
+}
+
+// limitRunes counts the number of bytes in the provided
+// byte slice up to a given rune limit.
+func limitRunes(p []byte, limit int) (runes int, end int, excess bool) {
+	for runes < limit && len(p) > 0 {
+		_, s := utf8.DecodeRune(p)
+		runes++
+		end += s
+		p = p[s:]
+	}
+	excess = len(p) != 0
+	return
+}
+
+// limitRunesInString counts the number of bytes in the provided
+// string up to a given rune limit.
+func limitRunesInString(str string, limit int) (runes int, end int, excess bool) {
+	for runes < limit && len(str) > 0 {
+		_, s := utf8.DecodeRuneInString(str)
+		runes++
+		end += s
+		str = str[s:]
+	}
+	excess = len(str) != 0
+	return
 }
 
 // WriteByte emits a single byte.
 func (b *Buffer) WriteByte(s byte) error {
+	if b.runeLimit != 0 {
+		if b.runeLimit == b.runeCount {
+			b.excess = true
+			return nil
+		}
+		b.runeCount++
+	}
+
 	b.startWrite()
 	if b.mode == UnsafeEscaped &&
 		(s >= utf8.RuneSelf ||
@@ -168,6 +242,14 @@ func (b *Buffer) WriteByte(s byte) error {
 
 // WriteRune emits a single rune.
 func (b *Buffer) WriteRune(s rune) error {
+	if b.runeLimit != 0 {
+		if b.runeLimit == b.runeCount {
+			b.excess = true
+			return nil
+		}
+		b.runeCount++
+	}
+
 	b.startWrite()
 	l := utf8.RuneLen(s)
 	m, ok := b.tryGrowByReslice(l)
@@ -178,6 +260,47 @@ func (b *Buffer) WriteRune(s rune) error {
 	return nil
 }
 
+// TruncateState is a struct used to store the previously set
+// rune limit. This is used to keep track of rune limits during
+// nested truncation operations.
+type TruncateState struct {
+	buf   *Buffer
+	limit int
+}
+
+// Truncate sets the rune limit and returns a TruncateState
+// that can be used to restore the previous rune limit. The
+// limit cannot exceed the previous rune limit.
+func (b *Buffer) Truncate(limit int) TruncateState {
+	prev := TruncateState{b, b.runeLimit}
+	if b.runeLimit == 0 || b.runeCount+limit < b.runeLimit {
+		b.runeLimit = b.runeCount + limit
+	}
+	return prev
+}
+
+// Restore restores the previous rune limit. If the current rune limit
+// is exceeded, a truncation marker is written.
+func (s TruncateState) Restore() {
+	if s.buf.excess && (s.limit == 0 || s.limit != s.buf.runeLimit) {
+		s.buf.excess = false
+		s.buf.markTruncate()
+	}
+	s.buf.runeLimit = s.limit
+}
+
+// markTruncate adds the truncation marker.
+func (b *Buffer) markTruncate() {
+	if len(b.buf) == 0 {
+		return
+	}
+	p, ok := b.tryGrowByReslice(len(m.Truncate))
+	if !ok {
+		p = b.grow(len(m.Truncate))
+	}
+	copy(b.buf[p:], m.Truncate)
+}
+
 // finalize ensures that all the buffer is properly
 // marked.
 func (b *Buffer) finalize() {

internal/markers/constants.go

@@ -25,6 +25,7 @@ const (
 	EscapeMark  = '?'
 	EscapeMarkS = string(EscapeMark)
 	RedactedS   = StartS + "×" + EndS
+	Truncate    = `…`
 )
 
 // Internal variables.

markers_test.go

@@ -204,6 +204,46 @@ func TestPrinter(t *testing.T) {
 		{func(w p) { w.Print(buf) }, "safe ‹unsafe›"},
 		{func(w p) { w.Printf("%v", &buf) }, "safe ‹unsafe›"},
 		{func(w p) { w.Print(&buf) }, "safe ‹unsafe›"},
+
+		// Truncation.
+		{func(w p) { w.Print(Truncate(Safe("säfe"), 3)) }, "säf…"},
+		{func(w p) { w.Print(Truncate(Safe("safe"), 5)) }, "safe"},
+		{func(w p) { w.Print(Truncate(Safe("safe"), 100)) }, "safe"},
+		{func(w p) { w.Print(Truncate("unsäfe", 3)) }, "‹uns›…"},
+		{func(w p) { w.Print(Truncate("unsafe", 100)) }, "‹unsafe›"},
+		{func(w p) {
+			w.Printf("Unsafe: %v Safe: %v", Truncate("unsafe", 3), Truncate(Safe("safe"), 2))
+		}, "Unsafe: ‹uns›… Safe: sa…"},
+		// Nested Truncation.
+		{func(w p) { w.Print(Truncate(Truncate(Safe("safe"), 5), 3)) }, "saf…"},
+		{func(w p) { w.Print(Truncate(Truncate(Safe("safe"), 3), 5)) }, "saf…"},
+		{func(w p) {
+			w.Print(Truncate(compose{func(p2 p) {
+				p2.Printf(
+					"Inside %v %v",
+					Truncate(Safe("safe"), 3),
+					Truncate("unsafe", 3),
+				)
+			}}, 10))
+		}, "Inside saf…"},
+		{func(w p) {
+			w.Print(Truncate(compose{func(p2 p) {
+				p2.Printf(
+					"Inside %v %v",
+					Truncate(Safe("safe"), 3),
+					Truncate("unsafe", 3),
+				)
+			}}, 20))
+		}, "Inside saf… ‹uns›…"},
+		{func(w p) {
+			w.Print(Truncate(compose{func(p2 p) {
+				p2.Printf(
+					"Inside %v %v",
+					Truncate(Safe("safe"), 3),
+					Truncate("unsafe", 20),
+				)
+			}}, 20))
+		}, "Inside saf… ‹unsafe›"},
 	}
 
 	var methods = []struct {