diff --git a/lib/compress/compress.go b/lib/compress/compress.go new file mode 100644 index 000000000..513afef67 --- /dev/null +++ b/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 + } +} diff --git a/lib/compress/compress_test.go b/lib/compress/compress_test.go new file mode 100644 index 000000000..b1fe82e12 --- /dev/null +++ b/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) +}