From 145514687c8b4abfc73c7d528610ad3c48fcc54a Mon Sep 17 00:00:00 2001
From: Bartosz Taudul <wolf.pld@gmail.com>
Date: Sat, 8 Feb 2020 15:43:01 +0100
Subject: [PATCH] Add zstd compression to FileWrite.

---
 server/TracyFileHeader.hpp |  3 ++-
 server/TracyFileWrite.hpp  | 37 +++++++++++++++++++++++++++++++------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/server/TracyFileHeader.hpp b/server/TracyFileHeader.hpp
index b49e0029..c5971019 100644
--- a/server/TracyFileHeader.hpp
+++ b/server/TracyFileHeader.hpp
@@ -6,7 +6,8 @@
 namespace tracy
 {
 
-static const char Lz4Header[4] = { 't', 'l', 'Z', 4 };
+static const char Lz4Header[4]  = { 't', 'l', 'Z', 4 };
+static const char ZstdHeader[4] = { 't', 'Z', 's', 't' };
 
 static constexpr tracy_force_inline int FileVersion( uint8_t h5, uint8_t h6, uint8_t h7 )
 {
diff --git a/server/TracyFileWrite.hpp b/server/TracyFileWrite.hpp
index d1d1d58b..099ab88a 100644
--- a/server/TracyFileWrite.hpp
+++ b/server/TracyFileWrite.hpp
@@ -15,6 +15,7 @@
 #include "../common/tracy_lz4.hpp"
 #include "../common/tracy_lz4hc.hpp"
 #include "../common/TracyForceInline.hpp"
+#include "../zstd/zstd.h"
 
 namespace tracy
 {
@@ -26,13 +27,14 @@ public:
     {
         Fast,
         Slow,
-        Extreme
+        Extreme,
+        Zstd
     };
 
-    static FileWrite* Open( const char* fn, Compression comp = Compression::Fast )
+    static FileWrite* Open( const char* fn, Compression comp = Compression::Fast, int level = 1 )
     {
         auto f = fopen( fn, "wb" );
-        return f ? new FileWrite( f, comp ) : nullptr;
+        return f ? new FileWrite( f, comp, level ) : nullptr;
     }
 
     ~FileWrite()
@@ -42,6 +44,7 @@ public:
 
         if( m_stream ) LZ4_freeStream( m_stream );
         if( m_streamHC ) LZ4_freeStreamHC( m_streamHC );
+        if( m_streamZstd ) ZSTD_freeCStream( m_streamZstd );
     }
 
     void Finish()
@@ -64,9 +67,10 @@ public:
     std::pair<size_t, size_t> GetCompressionStatistics() const { return std::make_pair( m_srcBytes, m_dstBytes ); }
 
 private:
-    FileWrite( FILE* f, Compression comp )
+    FileWrite( FILE* f, Compression comp, int level )
         : m_stream( nullptr )
         , m_streamHC( nullptr )
+        , m_streamZstd( nullptr )
         , m_file( f )
         , m_buf( m_bufData[0] )
         , m_second( m_bufData[1] )
@@ -86,12 +90,24 @@ private:
             m_streamHC = LZ4_createStreamHC();
             LZ4_resetStreamHC( m_streamHC, LZ4HC_CLEVEL_MAX );
             break;
+        case Compression::Zstd:
+            m_streamZstd = ZSTD_createCStream();
+            ZSTD_CCtx_setParameter( m_streamZstd, ZSTD_c_compressionLevel, level );
+            ZSTD_CCtx_setParameter( m_streamZstd, ZSTD_c_contentSizeFlag, 0 );
+            break;
         default:
             assert( false );
             break;
         }
 
-        fwrite( Lz4Header, 1, sizeof( Lz4Header ), m_file );
+        if( comp == Compression::Zstd )
+        {
+            fwrite( ZstdHeader, 1, sizeof( ZstdHeader ), m_file );
+        }
+        else
+        {
+            fwrite( Lz4Header, 1, sizeof( Lz4Header ), m_file );
+        }
     }
 
     tracy_force_inline void WriteSmall( const void* ptr, size_t size )
@@ -126,6 +142,14 @@ private:
         {
             sz = LZ4_compress_fast_continue( m_stream, m_buf, lz4, m_offset, LZ4Size, 1 );
         }
+        else if( m_streamZstd )
+        {
+            ZSTD_outBuffer out = { lz4, LZ4Size, 0 };
+            ZSTD_inBuffer in = { m_buf, m_offset, 0 };
+            const auto ret = ZSTD_compressStream2( m_streamZstd, &out, &in, ZSTD_e_flush );
+            assert( ret == 0 );
+            sz = out.pos;
+        }
         else
         {
             sz = LZ4_compress_HC_continue( m_streamHC, m_buf, lz4, m_offset, LZ4Size );
@@ -141,10 +165,11 @@ private:
     }
 
     enum { BufSize = 64 * 1024 };
-    enum { LZ4Size = LZ4_COMPRESSBOUND( BufSize ) };
+    enum { LZ4Size = std::max( LZ4_COMPRESSBOUND( BufSize ), ZSTD_COMPRESSBOUND( BufSize ) ) };
 
     LZ4_stream_t* m_stream;
     LZ4_streamHC_t* m_streamHC;
+    ZSTD_CStream* m_streamZstd;
     FILE* m_file;
     char m_bufData[2][BufSize];
     char* m_buf;