Zstd initramfs decompression.zstd

Change "lib/decompress_zstd.c" to be more flexible depends on input params.

Change-Id: I6e2608bc4d0cdf0e1efb9ceaaefd62f321ca9361
Signed-off-by: Ivaylo Asenov <ivaylo.asenov@linaro.org>

5 files changed, 173 insertions, 35 deletions

diff --git a/arch/arm/boot/compressed/head.S b/arch/arm/boot/compressed/head.S
index 37270459dafe..5e1db4d14d5e 100644
--- a/arch/arm/boot/compressed/head.S
+++ b/arch/arm/boot/compressed/head.S
@@ -542,7 +542,7 @@ not_relocated:	mov	r0, #0
 		mov	r0, r4
 		mov	r1, sp			@ malloc space above stack
 #ifdef CONFIG_KERNEL_ZSTD
-		add	r2, sp, #0x30000	@ 196608M max
+		add	r2, sp, #0x30000	@ 196608K max
 #else
 		add	r2, sp, #0x10000	@ 64k max
 #endif
diff --git a/include/linux/decompress/zstd.h b/include/linux/decompress/zstd.h
index 965c9fc5b147..52dbe2b2ac90 100644
--- a/include/linux/decompress/zstd.h
+++ b/include/linux/decompress/zstd.h
@@ -5,7 +5,6 @@ int unzstd(unsigned char *inbuf, long len,
 	long (*fill)(void*, unsigned long),
 	long (*flush)(void*, unsigned long),
 	unsigned char *output,
-	long len_out,
 	long *pos,
 	void(*error)(char *x));
 #endif
diff --git a/lib/decompress_zstd.c b/lib/decompress_zstd.c
index 1119f771b237..714e102eb9b5 100644
--- a/lib/decompress_zstd.c
+++ b/lib/decompress_zstd.c
@@ -29,71 +29,182 @@
 #include <linux/decompress/mm.h>
 
 
-STATIC inline int INIT unzstd(u8 *input, long in_len,
+STATIC int INIT __unzstd(u8 *input, long in_len,
 				long (*fill)(void *, unsigned long),
 				long (*flush)(void *, unsigned long),
 				u8 *output, long out_len, long *posp,
 				void (*error) (char *x))
 {
-	ZSTD_frameParams param;
 	ssize_t dst_size;
 	int ret = 0;
+	u8 *outp;
 
-	ret = ZSTD_getFrameParams(&param, input, in_len);
+	size_t cmpr_bytes;
+	size_t decmpr_bytes;
 
-	if(ret != 0) error("Decompress error: ZSTD_getFrameParams failed...");
+	//Try to get decompressed size
+	dst_size = ZSTD_findDecompressedSize(input, in_len);
 
-	dst_size = param.frameContentSize;
 
-	if(dst_size != out_len) {
+	if(dst_size != out_len && dst_size != ZSTD_CONTENTSIZE_ERROR && dst_size != ZSTD_CONTENTSIZE_UNKNOWN) {
 		  //If dst_size size calculated with zstd functions are different from out_len comes from unzstd function set out_len = dst_size
 		  out_len = dst_size;
-		 /*error("Decompress error: Decompression sizes are different...");*/
 	}
 
 
+	//Check input params
 	if(!input) error("Decompression error: ZSTD_decompress params are not ok...(input)");
-	if(dst_size<=0) error("Decompression error: ZSTD_decompress params are not ok...(dst_size)");
 	if(in_len<=0) error("Decompression error: ZSTD_decompress params are not ok...(in_len)");
+
+
+	//Streaming API (it is slower that Block API so stopping it for now and try to use only Block API for different cases)
+	/*note 1 : decompressed size is an optional field, that may not be present, especially in streaming mode.
+                When `return==0`, data to decompress could be any size.
+                In which case, it's necessary to use streaming mode to decompress data.
+                Optionally, application can still use ZSTD_decompress() while relying on implied limits.
+                (For example, data may be necessarily cut into blocks <= 16 KB).*/
+	/*
+	if((dst_size == ZSTD_CONTENTSIZE_ERROR || dst_size == ZSTD_CONTENTSIZE_UNKNOWN) && flush && 0){
+
+                size_t const outBuffSize = ZSTD_DStreamOutSize();
+
+		char* const outBuff = malloc(outBuffSize);
+		if(outBuff == NULL ) error("outBuff mem alloc failed");
+
+		size_t maxWindowSize = (1ULL << ZSTD_WINDOWLOG_MAX); // need to allocate enough memory: "vmalloc ="
+		size_t const workspaceSize = ZSTD_DStreamWorkspaceBound(maxWindowSize);
+
+		void *workspace = large_malloc(workspaceSize);
+		if(workspace == NULL) error("ZSTD_initDStream workspace mem alloc failed");
+
+		ZSTD_DStream *zds = ZSTD_initDStream(maxWindowSize, workspace, workspaceSize);
+		if(zds == NULL) error("zds mem alloc failed");
+
+		ZSTD_inBuffer in = {input, in_len, 0};
+
+                while (in.pos != in.size) {
+		      ZSTD_outBuffer out = {outBuff, outBuffSize, 0};
+		      ret = ZSTD_decompressStream(zds, &out, &in);
+		      if (ZSTD_isError(ret)) goto exit_errors;
+		      int l = out.pos;
+
+                      if(l != flush(outBuff, l)) error("write error");
+		}
+
+		//Report that all number of bytes read and decompressed
+                *posp =  in_len;
+
+                //Do some memory free
+                free(outBuff);
+                large_free(workspace);
+
+
+                return 0;
+	}
+	*/
+
+
+	if(output){
+		outp = output;
+
+	}else{
+		out_len = 0x800000; //(8 << 20)
+		outp = large_malloc(out_len);
+		if(!outp) error("Could not allocate outp buffer");
+	}
+
+	//In case of valid output but with no out_len
+	if(out_len<=0){
+		out_len = 0x800000; //(8 << 20)
+	}
+
+
+	//Check output params
+	if(!outp) error("Decompression error: ZSTD_decompress params are not ok...(outp)");
 	if(out_len<=0) error("Decompression error: ZSTD_decompress params are not ok...(out_len)");
-	if(!output) error("Decompression error: ZSTD_decompress params are not ok...(output)");
 
 	size_t const workspaceSize = ZSTD_DCtxWorkspaceBound();
 	void *workspace = malloc(workspaceSize);
-	if(workspace == NULL) error("workspace mem alloc failed");
+	if(workspace == NULL) error("ZSTD_initDCtx workspace mem alloc failed");
 
 	ZSTD_DCtx *dctx = ZSTD_initDCtx(workspace, workspaceSize);
 	if(dctx == NULL) error("dctx mem alloc failed");
 
-	ret = ZSTD_decompressDCtx(dctx ,output, out_len, input, in_len);
 
+	//When __unzstd function comes with valid output and dst_size is known let's use clear Block API(it is faster)
+	if(output &&  dst_size != ZSTD_CONTENTSIZE_ERROR && dst_size != ZSTD_CONTENTSIZE_UNKNOWN){
+		ret = ZSTD_decompressDCtx(dctx ,outp, out_len, input, in_len);
+		if (ZSTD_isError(ret)) goto exit_errors;
+		if(flush && ret != flush(outp, ret)) error("write error");
+
+		//Report that all number of bytes read and decompressed
+		if(posp) *posp = in_len;
+
+		goto exit_normal;
+
+	}
+
+
+	//ZSTD_nextSrcSizeToDecompress() tells how many bytes to provide as 'srcSize' to ZSTD_decompressContinue().
+	//A frame is fully decoded when ZSTD_nextSrcSizeToDecompress() returns zero.
+	while(cmpr_bytes = ZSTD_nextSrcSizeToDecompress(dctx))
+	{
+		decmpr_bytes = ZSTD_decompressContinue(dctx, outp, out_len, input, cmpr_bytes); //result of ZSTD_decompressContinue() is the number of bytes regenerated within 'dst'
+		if (ZSTD_isError(decmpr_bytes)) goto exit_errors;
+
+		if(flush && decmpr_bytes != flush(outp, decmpr_bytes)) error("write error");
+
+		input = input + cmpr_bytes;
+		if(output){outp = outp + decmpr_bytes;}
+	}
+	//Report that all number of bytes read and decompressed
+	if(posp) *posp = in_len;
+
+	goto exit_normal;
+
+
+
+
+
+
+
+exit_errors:
 	//Uncomment for debugging
 	// zstd errors can be found in zstd.h
-	if(ret == -1) error("Full decompression not yet implemented ...error1:");
-	if(ret == -2) error("Full decompression not yet implemented ...error2:");
-	if(ret == -3) error("Full decompression not yet implemented ...error3:");
-	if(ret == -4) error("Full decompression not yet implemented ...error4:");
-	if(ret == -5) error("Full decompression not yet implemented ...error5:");
-	if(ret == -6) error("Full decompression not yet implemented ...error6:");
-	if(ret == -7) error("Full decompression not yet implemented ...error7:");
-	if(ret == -8) error("Full decompression not yet implemented ...error8:");
-	if(ret == -9) error("Full decompression not yet implemented ...error9:");
-	if(ret == -10) error("Full decompression not yet implemented ...error10:");
-	if(ret == -11) error("Full decompression not yet implemented ...error11:");
-	if(ret == -12) error("Full decompression not yet implemented ...error12:");
-	if(ret == -13) error("Full decompression not yet implemented ...error13:");
-	if(ret == -14) error("Full decompression not yet implemented ...error14:");
-	if(ret == -15) error("Full decompression not yet implemented ...error15:");
-	if(ret == -16) error("Full decompression not yet implemented ...error16:");
-	if(ret == -17) error("Full decompression not yet implemented ...error17:");
-	if(ret == -18) error("Full decompression not yet implemented ...error18:");
-	if(ret == -19) error("Full decompression not yet implemented ...error19:");
-	if(ret == -20) error("Full decompression not yet implemented ...error20:");
-	if(ret == -21) error("Full decompression not yet implemented ...error21:");
+	if(ret == -1) error("Decompression error: ZSTD_error_GENERIC");
+	if(ret == -2) error("Decompression error: ZSTD_error_prefix_unknown");
+	if(ret == -3) error("Decompression error: ZSTD_error_version_unsupported");
+	if(ret == -4) error("Decompression error: ZSTD_error_parameter_unknown");
+	if(ret == -5) error("Decompression error: ZSTD_error_frameParameter_unsupported");
+	if(ret == -6) error("Decompression error: ZSTD_error_frameParameter_unsupportedBy32bits");
+	if(ret == -7) error("Decompression error: ZSTD_error_frameParameter_windowTooLarge");
+	if(ret == -8) error("Decompression error: ZSTD_error_compressionParameter_unsupported");
+	if(ret == -9) error("Decompression error: ZSTD_error_init_missing");
+	if(ret == -10) error("Decompression error: ZSTD_error_memory_allocation");
+	if(ret == -11) error("Decompression error: ZSTD_error_stage_wrong");
+	if(ret == -12) error("Decompression error: ZSTD_error_dstSize_tooSmall");
+	if(ret == -13) error("Decompression error: ZSTD_error_srcSize_wrong");
+	if(ret == -14) error("Decompression error: ZSTD_error_corruption_detected");
+	if(ret == -15) error("Decompression error: ZSTD_error_checksum_wrong");
+	if(ret == -16) error("Decompression error: ZSTD_error_tableLog_tooLarge");
+	if(ret == -17) error("Decompression error: ZSTD_error_maxSymbolValue_tooLarge");
+	if(ret == -18) error("Decompression error: ZSTD_error_maxSymbolValue_tooSmall");
+	if(ret == -19) error("Decompression error: ZSTD_error_dictionary_corrupted");
+	if(ret == -20) error("Decompression error: ZSTD_error_dictionary_wrong");
+	if(ret == -21) error("Decompression error: ZSTD_error_dictionaryCreation_failed");
+	if(ret == -22) error("Decompression error: ZSTD_error_maxCode");
 
 	if(ret <= 0) error("Decompression error: ZSTD_decompress failed...");
 	if(ret != dst_size) error("Decompression error!");
 
+exit_normal:
+
+	if (!output)
+		large_free(outp);
+
+	if(workspace)
+		free(workspace);
+
 	return 0;
 }
 
@@ -106,6 +217,16 @@ STATIC int INIT __decompress(unsigned char *buf, long in_len,
 			      void (*error)(char *x)
 	)
 {
-	return unzstd(buf, in_len, fill, flush, output, out_len, posp, error);
+	return __unzstd(buf, in_len, fill, flush, output, out_len, posp, error);
+}
+#else
+STATIC int INIT unzstd(unsigned char *buf, long in_len,
+		       long (*fill)(void*, unsigned long),
+		       long (*flush)(void*, unsigned long),
+		       unsigned char *outbuf,
+		       long *posp,
+		       void (*error)(char *x))
+{
+	return __unzstd(buf, in_len, fill, flush, outbuf, 0, posp, error);
 }
 #endif
diff --git a/scripts/gen_initramfs_list.sh b/scripts/gen_initramfs_list.sh
index 0055b07b03b6..ed52fe2090d8 100755
--- a/scripts/gen_initramfs_list.sh
+++ b/scripts/gen_initramfs_list.sh
@@ -261,6 +261,9 @@ case "$arg" in
 		echo "$output_file" | grep -q "\.lz4$" \
                 && [ -x "`which lz4 2> /dev/null`" ] \
                 && compr="lz4 -l -9 -f"
+	        echo "$output_file" | grep -q "\.zstd$" \
+                && [ -x "`which zstd 2> /dev/null`" ] \
+                && compr="zstd -f"
 		echo "$output_file" | grep -q "\.cpio$" && compr="cat"
 		shift
 		;;
diff --git a/usr/Kconfig b/usr/Kconfig
index 265b91e4dc17..0512b6c22d2f 100644
--- a/usr/Kconfig
+++ b/usr/Kconfig
@@ -216,6 +216,19 @@ config INITRAMFS_COMPRESSION_LZ4
 	  If you choose this, keep in mind that most distros don't provide lz4
 	  by default which could cause a build failure.
 
+
+config INITRAMFS_COMPRESSION_ZSTD
+	bool "ZSTD"
+	depends on RD_ZSTD
+	help
+	  It's compression ratio is the poorest amongst the choices. The kernel
+	  size is about 15% bigger than gzip; however its decompression speed
+	  is the fastest.
+
+	  If you choose this, keep in mind that you may need to install the zstd
+	  tool to be able to compress the initram.
+
+
 endchoice
 
 config INITRAMFS_COMPRESSION
@@ -227,10 +240,12 @@ config INITRAMFS_COMPRESSION
 	default ".xz"   if INITRAMFS_COMPRESSION_XZ
 	default ".lzo"  if INITRAMFS_COMPRESSION_LZO
 	default ".lz4"  if INITRAMFS_COMPRESSION_LZ4
+	default ".zstd"  if INITRAMFS_COMPRESSION_ZSTD
 	default ".gz"   if RD_GZIP
 	default ".lz4"  if RD_LZ4
 	default ".lzo"  if RD_LZO
 	default ".xz"   if RD_XZ
 	default ".lzma" if RD_LZMA
 	default ".bz2"  if RD_BZIP2
+	default ".zstd" if RD_ZSTD
 	default ""