Add builtin support for truncation

api.go

@@ -57,6 +57,9 @@ type SafeValue = i.SafeValue
 // TODO(knz): Remove this.
 type SafeMessager = i.SafeMessager
 
+// TODO(kzh): write comment
+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,9 @@ func Unsafe(a interface{}) interface{} { return w.Unsafe(a) }
 // The implementation is also slow.
 func Safe(a interface{}) SafeValue { return w.Safe(a) }
 
+// TODO(kzh): write comment
+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,8 @@ type SafeValue interface {
 type SafeMessager = interface {
 	SafeMessage() string
 }
+
+type TruncatedValue struct {
+	Value  interface{}
+	Length int
+}

internal/buffer/buffer.go

@@ -31,6 +31,10 @@ type Buffer struct {
 	validUntil int // exclusive upper bound of data that's already validated
 	mode       OutputMode
 	markerOpen bool
+
+	// TODO(kzh): describe these fields
+	runes int
+	limit int
 }
 
 // OutputMode determines how writes are processed in the Buffer.
@@ -129,6 +133,7 @@ func (b *Buffer) Cap() int {
 // buffer becomes too large, Write will panic with ErrTooLarge.
 func (b *Buffer) Write(p []byte) (n int, err error) {
 	b.startWrite()
+
 	m, ok := b.tryGrowByReslice(len(p))
 	if !ok {
 		m = b.grow(len(p))
@@ -140,6 +145,22 @@ func (b *Buffer) Write(p []byte) (n int, err error) {
 // 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) {
+	r := []rune(s)
+	if b.limit != 0 {
+		if b.runes >= b.limit {
+			b.runes = b.limit + 1
+			return
+		}
+
+		excess := (b.runes + len(r)) - b.limit
+		if excess > 0 {
+			r = r[:len(r)-excess]
+			b.runes++
+		}
+		b.runes += len(r)
+		s = string(r)
+	}
+
 	b.startWrite()
 	m, ok := b.tryGrowByReslice(len(s))
 	if !ok {
@@ -150,6 +171,16 @@ func (b *Buffer) WriteString(s string) (n int, err error) {
 
 // WriteByte emits a single byte.
 func (b *Buffer) WriteByte(s byte) error {
+	if b.limit != 0 {
+		if b.runes <= b.limit {
+			b.runes++
+		}
+
+		if b.Truncated() {
+			return nil
+		}
+	}
+
 	b.startWrite()
 	if b.mode == UnsafeEscaped &&
 		(s >= utf8.RuneSelf ||
@@ -168,6 +199,16 @@ func (b *Buffer) WriteByte(s byte) error {
 
 // WriteRune emits a single rune.
 func (b *Buffer) WriteRune(s rune) error {
+	if b.limit != 0 {
+		if b.runes <= b.limit {
+			b.runes++
+		}
+
+		if b.Truncated() {
+			return nil
+		}
+	}
+
 	b.startWrite()
 	l := utf8.RuneLen(s)
 	m, ok := b.tryGrowByReslice(l)
@@ -178,6 +219,43 @@ func (b *Buffer) WriteRune(s rune) error {
 	return nil
 }
 
+type TruncateState struct {
+	buf   *Buffer
+	limit int
+}
+
+// TODO(kzh): write here "limit rune count"
+func (b *Buffer) Truncate(limit int) TruncateState {
+	prev := TruncateState{b, b.limit}
+	if b.limit == 0 || b.runes+limit < b.limit {
+		b.limit = b.runes+limit
+	}
+	return prev
+}
+
+func (s TruncateState) Restore() {
+	if s.buf.Truncated() && (s.limit == 0 || s.limit > s.buf.limit) {
+		s.buf.runes--
+		s.buf.markTruncate()
+	}
+	s.buf.limit = s.limit
+}
+
+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)
+}
+
+func (b *Buffer) Truncated() bool {
+	return b.runes == b.limit+1
+}
+
 // 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›"},
+
+		// Truncate
+		{func(w p) { w.Print(Truncate(Safe("safe"), 3)) }, "saf..."},
+		{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("unsafe", 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..."},
+		// Recursive Truncate
+		{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 {