2022-11-25 07:22:31PM

lib/compress/compress.go

@@ -0,0 +1,67 @@
+package compress
+
+import (
+	"bytes"
+	"compress/flate"
+	"encoding/base64"
+	"io"
+	"strings"
+
+	"oss.terrastruct.com/d2/d2graph"
+)
+
+var compressionDict = "->" +
+	"<-" +
+	"--" +
+	"<->"
+
+var compressionDictBytes []byte
+
+// Compress takes a D2 script and compresses it to a URL-encoded string
+func Compress(raw string) (string, error) {
+	var b bytes.Buffer
+
+	zw, err := flate.NewWriterDict(&b, flate.DefaultCompression, []byte(compressionDict))
+	if err != nil {
+		return "", err
+	}
+	if _, err := io.Copy(zw, strings.NewReader(raw)); err != nil {
+		return "", err
+	}
+	if err := zw.Close(); err != nil {
+		return "", err
+	}
+
+	encoded := base64.URLEncoding.EncodeToString(b.Bytes())
+	return encoded, nil
+}
+
+// Decompress takes a compressed, URL-encoded string and returns the decompressed D2 script
+func Decompress(encoded string) (string, error) {
+	b64Decoded, err := base64.URLEncoding.DecodeString(encoded)
+	if err != nil {
+		return "", err
+	}
+
+	zr := flate.NewReaderDict(bytes.NewReader(b64Decoded), []byte(compressionDict))
+	var b bytes.Buffer
+	if _, err := io.Copy(&b, zr); err != nil {
+		return "", err
+	}
+	if err := zr.Close(); err != nil {
+		return "", nil
+	}
+	return b.String(), nil
+}
+
+func init() {
+	for k := range d2graph.StyleKeywords {
+		compressionDict += k
+	}
+	for k := range d2graph.ReservedKeywords {
+		compressionDict += k
+	}
+	for k := range d2graph.ReservedKeywordHolders {
+		compressionDict += k
+	}
+}

lib/compress/compress_test.go

@@ -0,0 +1,27 @@
+package compress
+
+import (
+	"testing"
+
+	"oss.terrastruct.com/diff"
+)
+
+func TestCompression(t *testing.T) {
+	script := `x -> y
+I just forgot my whole philosophy of life!!!: {
+  s: TV is chewing gum for the eyes
+}
+`
+
+	encoded, err := Compress(script)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	decoded, err := Decompress(encoded)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	diff.AssertStringEq(t, script, decoded)
+}