From 4996dffa8825de10011925c85c951c73fc9a833d Mon Sep 17 00:00:00 2001
From: zargovv <zargovv@gmail.com>
Date: Sat, 13 Jan 2024 21:58:24 +0600
Subject: [PATCH 1/4] Assert `deflateInit` for potential errors

---
 src/knn.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/knn.c b/src/knn.c
index 5d120fe..3afa503 100644
--- a/src/knn.c
+++ b/src/knn.c
@@ -23,7 +23,7 @@ Nob_String_View deflate_sv(Arena *arena, Nob_String_View sv)
     defstream.avail_out = (uInt)output_size;
     defstream.next_out = (Bytef *)output;
 
-    deflateInit(&defstream, Z_BEST_COMPRESSION);
+    assert(deflateInit(&defstream, Z_BEST_COMPRESSION) == Z_OK, "Failed to initialize zlib deflate stream");
     int result = deflate(&defstream, Z_FINISH);
     assert(result == Z_STREAM_END && "Probably not enough output buffer was allocated");
     deflateEnd(&defstream);

From 4bb8612100d8ff4e1c86c3f29448ffeda0ee302a Mon Sep 17 00:00:00 2001
From: zargovv <zargovv@gmail.com>
Date: Sat, 13 Jan 2024 22:03:05 +0600
Subject: [PATCH 2/4] Allocating exact size using `deflateBound()`

---
 src/knn.c | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/knn.c b/src/knn.c
index 3afa503..381b44c 100644
--- a/src/knn.c
+++ b/src/knn.c
@@ -14,18 +14,17 @@
 // Stolen from https://gist.github.com/arq5x/5315739
 Nob_String_View deflate_sv(Arena *arena, Nob_String_View sv)
 {
-    size_t output_size = sv.count*2;
-    void *output = arena_alloc(arena, output_size);
-
     z_stream defstream = {0};
     defstream.avail_in = (uInt)sv.count;
     defstream.next_in = (Bytef *)sv.data;
+
+    assert(deflateInit(&defstream, Z_BEST_COMPRESSION) == Z_OK, "Failed to initialize zlib deflate stream");
+    size_t output_size = (size_t)deflateBound(&defstream, defstream.avail_in);
+    void *output = arena_alloc(arena, output_size);
     defstream.avail_out = (uInt)output_size;
     defstream.next_out = (Bytef *)output;
 
-    assert(deflateInit(&defstream, Z_BEST_COMPRESSION) == Z_OK, "Failed to initialize zlib deflate stream");
-    int result = deflate(&defstream, Z_FINISH);
-    assert(result == Z_STREAM_END && "Probably not enough output buffer was allocated");
+    assert(deflate(&defstream, Z_FINISH) == Z_STREAM_END && "Probably `avail_in` is zero");
     deflateEnd(&defstream);
 
     return nob_sv_from_parts(output, defstream.total_out);

From 20ee0b88d9cee0e083bf8cd49d23592ad6cfc539 Mon Sep 17 00:00:00 2001
From: zargovv <zargovv@gmail.com>
Date: Sun, 14 Jan 2024 02:12:35 +0600
Subject: [PATCH 3/4] Introduce `deflateReset`

---
 src/knn.c | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/src/knn.c b/src/knn.c
index 381b44c..1f21b49 100644
--- a/src/knn.c
+++ b/src/knn.c
@@ -11,25 +11,32 @@
 
 #define K 2
 
+static __thread z_stream defstream = {0};
+static __thread bool defstream_init = false;
+
 // Stolen from https://gist.github.com/arq5x/5315739
 Nob_String_View deflate_sv(Arena *arena, Nob_String_View sv)
 {
-    z_stream defstream = {0};
     defstream.avail_in = (uInt)sv.count;
     defstream.next_in = (Bytef *)sv.data;
 
-    assert(deflateInit(&defstream, Z_BEST_COMPRESSION) == Z_OK, "Failed to initialize zlib deflate stream");
+    int ret = defstream_init ? deflateReset(&defstream) : deflateInit(&defstream, Z_BEST_COMPRESSION);
+    defstream_init = true;
+    assert(ret == Z_OK && "Failed to initialize zlib deflate stream");
     size_t output_size = (size_t)deflateBound(&defstream, defstream.avail_in);
     void *output = arena_alloc(arena, output_size);
     defstream.avail_out = (uInt)output_size;
     defstream.next_out = (Bytef *)output;
 
     assert(deflate(&defstream, Z_FINISH) == Z_STREAM_END && "Probably `avail_in` is zero");
-    deflateEnd(&defstream);
 
     return nob_sv_from_parts(output, defstream.total_out);
 }
 
+void deflate_end() {
+    deflateEnd(&defstream);
+}
+
 typedef struct {
     size_t klass;
     Nob_String_View text;
@@ -114,6 +121,7 @@ void *klassify_thread(void *params)
             .klass = state->train[i].klass,
         }));
     }
+    deflate_end();
 
     return NULL;
 }
@@ -133,7 +141,7 @@ typedef struct {
 
 void klass_predictor_init(Klass_Predictor *kp, Samples train_samples)
 {
-    kp->nprocs = get_nprocs();
+    kp->nprocs = get_nproc();
     kp->chunk_size = train_samples.count/kp->nprocs;
     kp->chunk_rem = train_samples.count%kp->nprocs;
     kp->train_samples = train_samples;

From 141571abcdde053afce1930357acfe9cdbf79abe Mon Sep 17 00:00:00 2001
From: Max Zargov <42868763+zargovv@users.noreply.github.com>
Date: Sun, 14 Jan 2024 02:46:21 +0600
Subject: [PATCH 4/4] Fix typo in knn.c

---
 src/knn.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/knn.c b/src/knn.c
index 1f21b49..83aca1d 100644
--- a/src/knn.c
+++ b/src/knn.c
@@ -141,7 +141,7 @@ typedef struct {
 
 void klass_predictor_init(Klass_Predictor *kp, Samples train_samples)
 {
-    kp->nprocs = get_nproc();
+    kp->nprocs = get_nprocs();
     kp->chunk_size = train_samples.count/kp->nprocs;
     kp->chunk_rem = train_samples.count%kp->nprocs;
     kp->train_samples = train_samples;