From dfc662888b73583fa023cb4a3d275dfabb631a93 Mon Sep 17 00:00:00 2001
From: Robert Chisholm <robadob@robadob.org>
Date: Sat, 21 Sep 2024 22:23:55 +0100
Subject: [PATCH] BugFix: JSONStateWriter prettyprint output is nolonger
 minified.

rapidjson's PrettyWriter extends Writer, but methods are non-virtual, leading to former bug.

This fix is best solution of the bad options (heavy templating, dual internal writers).

Closes #1231
---
 include/flamegpu/io/JSONStateWriter.h |  6 +++++-
 src/flamegpu/io/JSONStateWriter.cu    | 30 ++++++++++++++++++++-------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/include/flamegpu/io/JSONStateWriter.h b/include/flamegpu/io/JSONStateWriter.h
index f34c64db9..0ead7850a 100644
--- a/include/flamegpu/io/JSONStateWriter.h
+++ b/include/flamegpu/io/JSONStateWriter.h
@@ -61,9 +61,13 @@ class JSONStateWriter : public StateWriter {
     bool environment_written = false;
     bool macro_environment_written = false;
     bool agents_written = false;
+    // Dirty workaround for PrettyWriter overloads not being virtual
+    bool newline_purge_required = false;
     std::string outputPath;
     rapidjson::StringBuffer buffer;
-    std::unique_ptr<rapidjson::Writer<rapidjson::StringBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, rapidjson::CrtAllocator, rapidjson::kWriteNanAndInfFlag>> writer;
+    // Typedef because the template is too long
+    typedef rapidjson::PrettyWriter<rapidjson::StringBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, rapidjson::CrtAllocator, rapidjson::kWriteNanAndInfFlag> PrettyWriter;
+    std::unique_ptr<PrettyWriter> writer;
 };
 }  // namespace io
 }  // namespace flamegpu
diff --git a/src/flamegpu/io/JSONStateWriter.cu b/src/flamegpu/io/JSONStateWriter.cu
index d4bfb71c6..ca6e79151 100644
--- a/src/flamegpu/io/JSONStateWriter.cu
+++ b/src/flamegpu/io/JSONStateWriter.cu
@@ -8,6 +8,8 @@
 #include <string>
 #include <numeric>
 #include <set>
+#include <algorithm>
+#include <regex>
 
 #include "flamegpu/exception/FLAMEGPUException.h"
 #include "flamegpu/model/AgentDescription.h"
@@ -27,13 +29,11 @@ void JSONStateWriter::beginWrite(const std::string &output_file, bool pretty_pri
         THROW exception::UnknownInternalError("Writing already active, in JSONStateWriter::beginWrite()");
     }
     buffer = rapidjson::StringBuffer();
-    if (pretty_print) {
-        auto t_writer = std::make_unique<rapidjson::PrettyWriter<rapidjson::StringBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, rapidjson::CrtAllocator, rapidjson::kWriteNanAndInfFlag>>(buffer);
-        t_writer->SetIndent('\t', 1);
-        writer = std::move(t_writer);
-    } else {
-        writer = std::make_unique<rapidjson::Writer<rapidjson::StringBuffer, rapidjson::UTF8<>, rapidjson::UTF8<>, rapidjson::CrtAllocator, rapidjson::kWriteNanAndInfFlag>>(buffer);
-    }
+    writer = std::make_unique<PrettyWriter>(buffer);
+    // PrettyWriter overloads Writer, but methods aren't virtual so pointer casting doesn't work as intended
+    // This is the best of the bad options for this particular implementation
+    writer->SetIndent('\t', pretty_print ? 1 : 0);
+    newline_purge_required = !pretty_print;
     // Begin Json file
     writer->StartObject();
 
@@ -53,7 +53,21 @@ void JSONStateWriter::endWrite() {
     writer->EndObject();
 
     std::ofstream out(outputPath, std::ofstream::trunc);
-    out << buffer.GetString();
+    if (newline_purge_required) {
+        // Minify the output
+        std::string t_buffer = buffer.GetString();
+        // Replace all spaces outside of quotes with \n
+        bool in_string = false;
+        for (auto it = t_buffer.begin(); it != t_buffer.end(); ++it) {
+            if (*it == '"') in_string = !in_string;
+            if (*it == ' ' && !in_string) *it = '\n';
+        }
+        // Remove newlines
+        t_buffer.erase(std::remove(t_buffer.begin(), t_buffer.end(), '\n'), t_buffer.end());
+        out << t_buffer;
+    } else {
+        out << buffer.GetString();
+    }
     out.close();
 
     writer.reset();