Reorganize zstream API and fix wrap problems

There were some confusing issues mixing up the number of bytes
written to the zstream output buffer with the number of bytes
consumed from the zstream input.  This reorganizes the zstream
API and makes it easier to deflate an arbitrarily large input
while still using a fixed size output.

3 files changed, 125 insertions, 59 deletions

diff --git a/src/pack-objects.c b/src/pack-objects.c
index 8475f64f2..8b65ac26a 100644
--- a/src/pack-objects.c
+++ b/src/pack-objects.c
@@ -291,7 +291,6 @@ static int write_object(
 	void *delta_data = NULL;
 	void *data;
 	size_t hdr_len, zbuf_len = COMPRESS_BUFLEN, data_len;
-	ssize_t written;
 	int error;
 
 	if (po->delta) {
@@ -337,19 +336,15 @@ static int write_object(
 		GITERR_CHECK_ALLOC(zbuf);
 
 		git_zstream_reset(&pb->zstream);
+		git_zstream_set_input(&pb->zstream, data, data_len);
 
-		while ((written = git_zstream_deflate(zbuf, zbuf_len, &pb->zstream, data, data_len)) > 0) {
-			if ((error = write_cb(zbuf, written, cb_data)) < 0 ||
-				(error = git_hash_update(&pb->ctx, zbuf, written)) < 0)
+		while (!git_zstream_done(&pb->zstream)) {
+			if ((error = git_zstream_get_output(zbuf, &zbuf_len, &pb->zstream)) < 0 ||
+				(error = write_cb(zbuf, zbuf_len, cb_data)) < 0 ||
+				(error = git_hash_update(&pb->ctx, zbuf, zbuf_len)) < 0)
 				goto done;
 
-			data = (char *)data + written;
-			data_len -= written;
-		}
-
-		if (written < 0) {
-			error = (int)written;
-			goto done;
+			zbuf_len = COMPRESS_BUFLEN; /* reuse buffer */
 		}
 
 		if (po->delta)
diff --git a/src/zstream.c b/src/zstream.c
index 0bca72ff3..c83602297 100644
--- a/src/zstream.c
+++ b/src/zstream.c
@@ -10,14 +10,18 @@
 #include "zstream.h"
 #include "buffer.h"
 
-#define BUFFER_SIZE (1024 * 1024)
+#define ZSTREAM_BUFFER_SIZE (1024 * 1024)
+#define ZSTREAM_BUFFER_MIN_EXTRA 8
 
-static int zstream_seterr(int zerr, git_zstream *zstream)
+static int zstream_seterr(git_zstream *zs)
 {
-	if (zerr == Z_MEM_ERROR)
+	if (zs->zerr == Z_OK || zs->zerr == Z_STREAM_END)
+		return 0;
+
+	if (zs->zerr == Z_MEM_ERROR)
 		giterr_set_oom();
-	else if (zstream->msg)
-		giterr_set(GITERR_ZLIB, zstream->msg);
+	else if (zs->z.msg)
+		giterr_set(GITERR_ZLIB, zs->z.msg);
 	else
 		giterr_set(GITERR_ZLIB, "Unknown compression error");
 
@@ -26,69 +30,123 @@ static int zstream_seterr(int zerr, git_zstream *zstream)
 
 int git_zstream_init(git_zstream *zstream)
 {
-	int zerr;
-
-	if ((zerr = deflateInit(zstream, Z_DEFAULT_COMPRESSION)) != Z_OK)
-		return zstream_seterr(zerr, zstream);
-
-	return 0;
+	zstream->zerr = deflateInit(&zstream->z, Z_DEFAULT_COMPRESSION);
+	return zstream_seterr(zstream);
 }
 
-ssize_t git_zstream_deflate(void *out, size_t out_len, git_zstream *zstream, const void *in, size_t in_len)
+void git_zstream_free(git_zstream *zstream)
 {
-	int zerr;
-
-	if ((ssize_t)out_len < 0)
-		out_len = INT_MAX;
+	deflateEnd(&zstream->z);
+}
 
-	zstream->next_in = (Bytef *)in;
-	zstream->avail_in = in_len;
-	zstream->next_out = out;
-	zstream->avail_out = out_len;
+void git_zstream_reset(git_zstream *zstream)
+{
+	deflateReset(&zstream->z);
+	zstream->in = NULL;
+	zstream->in_len = 0;
+	zstream->zerr = Z_STREAM_END;
+}
 
-	if ((zerr = deflate(zstream, Z_FINISH)) == Z_STREAM_ERROR)
-		return zstream_seterr(zerr, zstream);
+int git_zstream_set_input(git_zstream *zstream, const void *in, size_t in_len)
+{
+	zstream->in = in;
+	zstream->in_len = in_len;
+	zstream->zerr = Z_OK;
+	return 0;
+}
 
-	return (out_len - zstream->avail_out);
+bool git_zstream_done(git_zstream *zstream)
+{
+	return (!zstream->in_len && zstream->zerr == Z_STREAM_END);
 }
 
-void git_zstream_reset(git_zstream *zstream)
+size_t git_zstream_suggest_output_len(git_zstream *zstream)
 {
-	deflateReset(zstream);
+	if (zstream->in_len > ZSTREAM_BUFFER_SIZE)
+		return ZSTREAM_BUFFER_SIZE;
+	else if (zstream->in_len > ZSTREAM_BUFFER_MIN_EXTRA)
+		return zstream->in_len;
+	else
+		return ZSTREAM_BUFFER_MIN_EXTRA;
 }
 
-void git_zstream_free(git_zstream *zstream)
+int git_zstream_get_output(void *out, size_t *out_len, git_zstream *zstream)
 {
-	deflateEnd(zstream);
+	int zflush = Z_FINISH;
+	size_t out_remain = *out_len;
+
+	while (out_remain > 0 && zstream->zerr != Z_STREAM_END) {
+		size_t out_queued, in_queued, out_used, in_used;
+
+		/* set up in data */
+		zstream->z.next_in  = (Bytef *)zstream->in;
+		zstream->z.avail_in = (uInt)zstream->in_len;
+		if ((size_t)zstream->z.avail_in != zstream->in_len) {
+			zstream->z.avail_in = INT_MAX;
+			zflush = Z_NO_FLUSH;
+		} else {
+			zflush = Z_FINISH;
+		}
+		in_queued = (size_t)zstream->z.avail_in;
+
+		/* set up out data */
+		zstream->z.next_out = out;
+		zstream->z.avail_out = (uInt)out_remain;
+		if ((size_t)zstream->z.avail_out != out_remain)
+			zstream->z.avail_out = INT_MAX;
+		out_queued = (size_t)zstream->z.avail_out;
+
+		/* compress next chunk */
+		zstream->zerr = deflate(&zstream->z, zflush);
+		if (zstream->zerr == Z_STREAM_ERROR)
+			return zstream_seterr(zstream);
+
+		out_used = (out_queued - zstream->z.avail_out);
+		out_remain -= out_used;
+		out = ((char *)out) + out_used;
+
+		in_used = (in_queued - zstream->z.avail_in);
+		zstream->in_len -= in_used;
+		zstream->in += in_used;
+	}
+
+	/* either we finished the input or we did not flush the data */
+	assert(zstream->in_len > 0 || zflush == Z_FINISH);
+
+	/* set out_size to number of bytes actually written to output */
+	*out_len = *out_len - out_remain;
+
+	return 0;
 }
 
 int git_zstream_deflatebuf(git_buf *out, const void *in, size_t in_len)
 {
-	git_zstream zstream = GIT_ZSTREAM_INIT;
-	size_t out_len;
-	ssize_t written;
+	git_zstream zs = GIT_ZSTREAM_INIT;
 	int error = 0;
 
-	if ((error = git_zstream_init(&zstream)) < 0)
+	if ((error = git_zstream_init(&zs)) < 0)
 		return error;
 
-	do {
-		if (out->asize - out->size < BUFFER_SIZE)
-			git_buf_grow(out, out->asize + BUFFER_SIZE);
+	if ((error = git_zstream_set_input(&zs, in, in_len)) < 0)
+		goto done;
 
-		out_len = out->asize - out->size;
+	while (!git_zstream_done(&zs)) {
+		size_t step = git_zstream_suggest_output_len(&zs), written;
 
-		if ((written = git_zstream_deflate(out->ptr + out->size, out_len, &zstream, in, in_len)) <= 0)
-			break;
+		if ((error = git_buf_grow(out, out->asize + step)) < 0)
+			goto done;
 
-		in = (char *)in + written;
-		in_len -= written;
-		out->size += written;
-	} while (written > 0);
+		written = out->asize - out->size;
 
-	if (written < 0)
-		error = written;
+		if ((error = git_zstream_get_output(
+				out->ptr + out->size, &written, &zs)) < 0)
+			goto done;
+
+		out->size += written;
+		out->ptr[out->size] = '\0';
+	}
 
-	git_zstream_free(&zstream);
+done:
+	git_zstream_free(&zs);
 	return error;
 }
diff --git a/src/zstream.h b/src/zstream.h
index 9672903c0..9b5bf6ace 100644
--- a/src/zstream.h
+++ b/src/zstream.h
@@ -12,15 +12,28 @@
 #include "common.h"
 #include "buffer.h"
 
-#define git_zstream z_stream
+typedef struct {
+	z_stream z;
+	const char *in;
+	size_t in_len;
+	int zerr;
+} git_zstream;
 
-#define GIT_ZSTREAM_INIT {0}
+#define GIT_ZSTREAM_INIT {{0}}
 
 int git_zstream_init(git_zstream *zstream);
-ssize_t git_zstream_deflate(void *out, size_t out_len, git_zstream *zstream, const void *in, size_t in_len);
-void git_zstream_reset(git_zstream *zstream);
 void git_zstream_free(git_zstream *zstream);
 
+int git_zstream_set_input(git_zstream *zstream, const void *in, size_t in_len);
+
+size_t git_zstream_suggest_output_len(git_zstream *zstream);
+
+int git_zstream_get_output(void *out, size_t *out_len, git_zstream *zstream);
+
+bool git_zstream_done(git_zstream *zstream);
+
+void git_zstream_reset(git_zstream *zstream);
+
 int git_zstream_deflatebuf(git_buf *out, const void *in, size_t in_len);
 
 #endif /* INCLUDE_zstream_h__ */