From d048159798b59fcf7837329461e0f5f6bf0c9890 Mon Sep 17 00:00:00 2001
From: Grzegorz Kowal <grzegorz@amuncode.org>
Date: Thu, 13 Aug 2020 15:42:07 -0300
Subject: [PATCH] COMPRESSION: Add LZMA/XZ compression format.

Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
---
 CMakeLists.txt          | 10 +++++++++
 python/amunpy.py        | 31 +++++++++++++++++++++++++++
 sources/compression.F90 | 46 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 20ce6c5..2359e12 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -92,3 +92,13 @@ if(ENABLE_LZ4)
 		target_link_libraries(amun.x ${LZ4_LIBRARIES})
 	endif()
 endif()
+
+option(ENABLE_LZMA "Enable LZMA compression support" ON)
+if(ENABLE_LZMA)
+	include(FindPkgConfig)
+	pkg_search_module(LZMA QUIET liblzma)
+	if(LZMA_FOUND)
+		target_compile_definitions(amun.x PRIVATE LZMA)
+		target_link_libraries(amun.x ${LZMA_LIBRARIES})
+	endif()
+endif()
diff --git a/python/amunpy.py b/python/amunpy.py
index 3c63585..b6d17e8 100644
--- a/python/amunpy.py
+++ b/python/amunpy.py
@@ -60,6 +60,11 @@ try:
 	lz4_available = True
 except ImportError:
 	lz4_available = False
+try:
+	import lzma
+	lzma_available = True
+except ImportError:
+	lzma_available = False
 
 
 class AmunXML:
@@ -225,6 +230,19 @@ class AmunXML:
 					else:
 						print('LZ4 module unavailable. Please, install LZ4 for Python.')
 						return None
+				elif cformat == 'lzma':
+					if lzma_available:
+						with open(dfile, mode ='rb') as f:
+							compressed = f.read()
+						if 'compressed_digest' in self.binaries[aname]:
+							self.check_digest(dfile, compressed, self.binaries[aname]['compressed_digest'])
+						data = lzma.decompress(compressed)
+						if 'digest' in self.binaries[aname]:
+							self.check_digest(dfile, data, self.binaries[aname]['digest'])
+						return np.frombuffer(data, dtype = dtype)
+					else:
+						print('LZMA module unavailable. Please, install LZMA for Python.')
+						return None
 				else:
 					print('Compression format {} unavailable. Please, install corresponding Python module'.format(cformat))
 					return None
@@ -273,6 +291,19 @@ class AmunXML:
 					else:
 						print('LZ4 module unavailable. Please, install LZ4 for Python.')
 						return None
+				elif cformat == 'lzma':
+					if lzma_available:
+						with open(dfile, mode ='rb') as f:
+							compressed = f.read()
+						if 'compressed_digest' in self.chunks[n][var]:
+							self.check_digest(dfile, compressed, self.chunks[n][var]['compressed_digest'])
+						data = lzma.decompress(compressed)
+						if 'digest' in self.chunks[n][var]:
+							self.check_digest(dfile, data, self.chunks[n][var]['digest'])
+						return np.frombuffer(data, dtype = dtype)
+					else:
+						print('LZMA module unavailable. Please, install LZMA for Python.')
+						return None
 				else:
 					print('Compression format {} unavailable. Please, install corresponding Python module'.format(cformat))
 					return None
diff --git a/sources/compression.F90 b/sources/compression.F90
index 5c289eb..196afe2 100644
--- a/sources/compression.F90
+++ b/sources/compression.F90
@@ -79,6 +79,22 @@ module compression
 
   end interface
 #endif /* LZ4 */
+#ifdef LZMA
+  interface
+
+    integer(c_int) function lzma_compress(preset, check, allocator,            &
+                                          src, srcSize,                        &
+                                          dst, dstPos, dstCapacity)            &
+                            bind(C, name="lzma_easy_buffer_encode")
+      use iso_c_binding, only: c_int, c_size_t, c_ptr
+      implicit none
+      integer(c_int)        , value :: preset, check
+      integer(kind=c_size_t), value :: srcSize, dstCapacity
+      type(c_ptr)           , value :: allocator, src, dst, dstPos
+    end function lzma_compress
+
+  end interface
+#endif /* LZMA */
 
 ! compression parameters
 !
@@ -90,6 +106,7 @@ module compression
   integer, parameter :: compression_none = 0
   integer, parameter :: compression_zstd = 1
   integer, parameter :: compression_lz4  = 2
+  integer, parameter :: compression_lzma = 3
 
 ! by default everything is private
 !
@@ -151,6 +168,13 @@ module compression
       compression_level  = 0
       suffix             = ".lz4"
 #endif /* LZ4 */
+#ifdef LZMA
+    case("lzma", "LZMA", "xz", "XZ")
+      cformat            = "lzma"
+      compression_format = compression_lzma
+      compression_level  = max(0, min(9, clevel))
+      suffix             = ".xz"
+#endif /* LZMA */
     case default
       cformat            = "none"
       compression_format = compression_none
@@ -211,7 +235,13 @@ module compression
 !
     integer(kind=1), dimension(:), target, intent(in)  :: input
     integer(kind=1), dimension(:), target, intent(out) :: output
-    integer(kind=8)                      , intent(out) :: csize
+    integer(kind=8)              , target, intent(out) :: csize
+
+#ifdef LZMA
+! return value
+!
+    integer :: ret
+#endif /* LZMA */
 
 ! compression buffer
 !
@@ -245,6 +275,20 @@ module compression
       end if
       deallocate(buffer)
 #endif /* LZ4 */
+#ifdef LZMA
+    case(compression_lzma)
+      csize = 0
+      allocate(buffer(size(input)))
+      ret = lzma_compress(compression_level, 0, c_null_ptr,                    &
+                           c_loc(input), sizeof(input),                        &
+                           c_loc(buffer), c_loc(csize), sizeof(buffer))
+      if (ret == 0 .and. csize <= sizeof(output)) then
+        output(1:csize) = buffer(1:csize)
+      else
+        csize = -1
+      end if
+      deallocate(buffer)
+#endif /* LZMA */
     case default
       output(1:csize) = input(1:csize)
     end select