1package compressor
  2
  3import (
  4	"bytes"
  5	"io"
  6	"unicode/utf8"
  7
  8	"github.com/muesli/ansi"
  9)
 10
 11type Writer struct {
 12	Forward io.Writer
 13
 14	ansi         bool
 15	ansiseq      bytes.Buffer
 16	lastseq      bytes.Buffer
 17	prevlastseq  bytes.Buffer
 18	resetreq     bool
 19	runeBuf      []byte
 20	compressed   int
 21	uncompressed int
 22}
 23
 24// Bytes is shorthand for declaring a new default compressor instance,
 25// used to immediately compress a byte slice.
 26func Bytes(b []byte) []byte {
 27	var buf bytes.Buffer
 28	f := Writer{
 29		Forward: &buf,
 30	}
 31	_, _ = f.Write(b)
 32	_ = f.Close()
 33
 34	return buf.Bytes()
 35}
 36
 37// String is shorthand for declaring a new default compressor instance,
 38// used to immediately compress a string.
 39func String(s string) string {
 40	return string(Bytes([]byte(s)))
 41}
 42
 43// Write is used to write content to the ANSI buffer.
 44func (w *Writer) Write(b []byte) (int, error) {
 45	w.uncompressed += len(b)
 46
 47	for _, c := range string(b) {
 48		if c == ansi.Marker {
 49			// ANSI escape sequence
 50			w.ansi = true
 51			_, _ = w.ansiseq.WriteRune(c)
 52		} else if w.ansi {
 53			_, _ = w.ansiseq.WriteRune(c)
 54			if ansi.IsTerminator(c) {
 55				// ANSI sequence terminated
 56				w.ansi = false
 57
 58				terminated := false
 59				if bytes.HasSuffix(w.ansiseq.Bytes(), []byte("[0m")) {
 60					// reset sequence
 61					w.prevlastseq.Reset()
 62					w.prevlastseq.Write(w.lastseq.Bytes())
 63
 64					w.lastseq.Reset()
 65					terminated = true
 66					w.resetreq = true
 67				} else if c == 'm' {
 68					// color code
 69					_, _ = w.lastseq.Write(w.ansiseq.Bytes())
 70				}
 71
 72				if !terminated {
 73					// did we reset the sequence just to restore it again?
 74					if bytes.Equal(w.ansiseq.Bytes(), w.prevlastseq.Bytes()) {
 75						w.resetreq = false
 76						w.ansiseq.Reset()
 77					}
 78
 79					w.prevlastseq.Reset()
 80
 81					if w.resetreq {
 82						w.ResetAnsi()
 83					}
 84
 85					_, _ = w.Forward.Write(w.ansiseq.Bytes())
 86					w.compressed += w.ansiseq.Len()
 87				}
 88
 89				w.ansiseq.Reset()
 90			}
 91		} else {
 92			if w.resetreq {
 93				w.ResetAnsi()
 94			}
 95
 96			_, err := w.writeRune(c)
 97			if err != nil {
 98				return 0, err
 99			}
100		}
101	}
102
103	return len(b), nil
104}
105
106func (w *Writer) writeRune(r rune) (int, error) {
107	if w.runeBuf == nil {
108		w.runeBuf = make([]byte, utf8.UTFMax)
109	}
110	n := utf8.EncodeRune(w.runeBuf, r)
111	w.compressed += n
112	return w.Forward.Write(w.runeBuf[:n])
113}
114
115// Close finishes the compression operation. Always call it before trying to
116// retrieve the final result.
117func (w *Writer) Close() error {
118	if w.resetreq {
119		w.ResetAnsi()
120	}
121
122	// log.Println("Written uncompressed: ", w.uncompressed)
123	// log.Println("Written compressed: ", w.compressed)
124
125	return nil
126}
127
128func (w *Writer) ResetAnsi() {
129	w.prevlastseq.Reset()
130	_, _ = w.Forward.Write([]byte("\x1b[0m"))
131	w.resetreq = false
132}