From a1a679f473a9abf9170e31cd54098d3368c91f87 Mon Sep 17 00:00:00 2001
From: Andrew Bruce <me@andrewbruce.net>
Date: Sun, 20 Oct 2024 23:48:13 +0100
Subject: [PATCH] Only write compressed data when it's smaller

---
 lib/mudbrick.ex               |  9 ++++++++
 lib/mudbrick/stream.ex        | 41 +++++++++++++++++++++--------------
 test/mudbrick/stream_test.exs | 23 ++++++++++++++++++++
 3 files changed, 57 insertions(+), 16 deletions(-)

diff --git a/lib/mudbrick.ex b/lib/mudbrick.ex
index a9d71fa..8b76d9f 100644
--- a/lib/mudbrick.ex
+++ b/lib/mudbrick.ex
@@ -114,6 +114,15 @@ defmodule Mudbrick do
     Enum.map_join(list, separator, &Mudbrick.Object.from/1)
   end
 
+  def compress(data) do
+    z = :zlib.open()
+    :ok = :zlib.deflateInit(z)
+    deflated = :zlib.deflate(z, data, :finish)
+    :zlib.deflateEnd(z)
+    :zlib.close(z)
+    deflated
+  end
+
   defp contents({doc, page}) do
     import Document
 
diff --git a/lib/mudbrick/stream.ex b/lib/mudbrick/stream.ex
index 7d3be97..6e4f50a 100644
--- a/lib/mudbrick/stream.ex
+++ b/lib/mudbrick/stream.ex
@@ -5,30 +5,39 @@ defmodule Mudbrick.Stream do
             length: nil,
             filters: []
 
+  import :erlang, only: [iolist_size: 1]
+
   def new(opts) do
     opts =
-      if opts[:compress] do
-        opts
-        |> Keyword.update!(:data, fn data ->
-          z = :zlib.open()
-          :ok = :zlib.deflateInit(z)
-          deflated = :zlib.deflate(z, data, :finish)
-          :zlib.deflateEnd(z)
-          :zlib.close(z)
-          deflated
-        end)
-        |> then(fn opts ->
+      case decide_compression(opts) do
+        {:ok, data} ->
           Keyword.merge(
             opts,
-            length: :erlang.iolist_size(opts[:data]),
+            data: data,
+            length: iolist_size(data),
             filters: [:FlateDecode]
           )
-        end)
-      else
-        opts
+
+        :error ->
+          opts
       end
 
-    struct!(__MODULE__, Keyword.put(opts, :length, :erlang.iolist_size(opts[:data])))
+    struct!(__MODULE__, Keyword.put(opts, :length, iolist_size(opts[:data])))
+  end
+
+  defp decide_compression(opts) do
+    if opts[:compress] do
+      uncompressed = opts[:data]
+      compressed = Mudbrick.compress(uncompressed)
+
+      if iolist_size(compressed) < iolist_size(uncompressed) do
+        {:ok, compressed}
+      else
+        :error
+      end
+    else
+      :error
+    end
   end
 
   defimpl Mudbrick.Object do
diff --git a/test/mudbrick/stream_test.exs b/test/mudbrick/stream_test.exs
index 275f7c3..cadb193 100644
--- a/test/mudbrick/stream_test.exs
+++ b/test/mudbrick/stream_test.exs
@@ -1,8 +1,31 @@
 defmodule Mudbrick.StreamTest do
   use ExUnit.Case, async: true
+  use ExUnitProperties
 
   import TestHelper
 
+  property "compresses data when there's a saving" do
+    check all uncompressed <- string(:alphanumeric, min_length: 150), max_runs: 200 do
+      result =
+        Mudbrick.Stream.new(compress: true, data: uncompressed)
+        |> Mudbrick.Object.from()
+        |> :erlang.iolist_to_binary()
+
+      assert result =~ "FlateDecode"
+    end
+  end
+
+  property "doesn't compress data when there's no saving" do
+    check all uncompressed <- string(:alphanumeric, max_length: 120), max_runs: 200 do
+      result =
+        Mudbrick.Stream.new(compress: true, data: uncompressed)
+        |> Mudbrick.Object.from()
+        |> :erlang.iolist_to_binary()
+
+      refute result =~ "FlateDecode"
+    end
+  end
+
   test "includes length and stream markers when serialised" do
     serialised =
       Mudbrick.Stream.new(data: bodoni())