Streaming callback API, WASM chat mode, and C safety audit fixes

Add callback-based generate_response() in main.c for streaming token
output, refactor chat REPL to use it. Add 5-function chat state machine
to WASM API (init/reset/submit/next/end_turn) with loop detection and
auto-reset. Add chat UI with mode toggle, message history, and stop
button to the browser demo.

Safety audit fixes: validate tensor end-bounds in GGUF parser, add VLA
size guards, fix platform unload for is_mmap==0, consolidate model file
ownership in bn_model_free, replace atoi/atof with strtol/strtof, add
-Wshadow and make asan target, replace assert with runtime fallback in
batch matvec, add softmax size guard, add tokenizer malloc overflow
checks, fix stale test files to use bn_-prefixed API, and add new tests
for arena allocator, sampler repeat penalty/top-p, and tokenizer edge
cases.

Author: farkasmark

Makefile

@@ -6,9 +6,9 @@ LDFLAGS = -lm
 # -march=native on Apple clang misses __ARM_FEATURE_FP16_VECTOR_ARITHMETIC.
 UNAME_S := $(shell uname -s)
 ifeq ($(UNAME_S),Darwin)
-CFLAGS  = -O3 -mcpu=apple-m1 -Wall -Wextra -std=c11 -Iinclude
+CFLAGS  = -O3 -mcpu=apple-m1 -Wall -Wextra -Wshadow -std=c11 -Iinclude
 else
-CFLAGS  = -O3 -march=native -Wall -Wextra -std=c11 -Iinclude
+CFLAGS  = -O3 -march=native -Wall -Wextra -Wshadow -std=c11 -Iinclude
 endif
 
 # On Linux, enable GNU extensions for strdup, qsort_r, clock_gettime, etc.
@@ -30,14 +30,19 @@ bitnet: $(OBJS)
 debug: CFLAGS += -DDEBUG -g -O0
 debug: bitnet
 
+# Sanitizer build (ASan + UBSan)
+asan: CFLAGS += -DDEBUG -g -O0 -fsanitize=address,undefined -fno-omit-frame-pointer
+asan: LDFLAGS += -fsanitize=address,undefined
+asan: bitnet
+
 # Pattern rule for object files
 src/%.o: src/%.c
 	$(CC) $(CFLAGS) -c -o $@ $<
 
 # --- Tests ---
-.PHONY: debug test test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety test_prefill test_kv_f16 pgo clean
+.PHONY: debug asan test test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety test_arena test_prefill test_kv_f16 pgo clean
 
-test: test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety
+test: test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety test_arena
 
 test_gguf: test/test_gguf.c src/gguf.c src/platform.c src/sh_log.c
 	$(CC) $(CFLAGS) -o $@ $^ $(LDFLAGS) && ./$@
@@ -61,6 +66,9 @@ test_safety: test/test_safety.c src/platform.c src/gguf.c src/quant.c src/model.
              src/sh_arena.c src/sh_log.c
 	$(CC) $(CFLAGS) -o $@ $^ $(LDFLAGS) && ./$@
 
+test_arena: test/test_arena.c src/sh_arena.c
+	$(CC) $(CFLAGS) -o $@ $^ $(LDFLAGS) && ./$@
+
 test_e2e: test/test_e2e.c src/platform.c src/gguf.c src/quant.c src/model.c \
           src/transformer.c src/tokenizer.c src/sampler.c src/threadpool.c \
           src/sh_arena.c src/sh_log.c
@@ -93,4 +101,4 @@ pgo:
 	@echo "=== PGO build complete ==="
 
 clean:
-	rm -f bitnet src/*.o test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety test_e2e test_prefill test_kv_f16 default.profraw default.profdata
+	rm -f bitnet src/*.o test_gguf test_quant test_tokenizer test_transformer test_threadpool test_safety test_arena test_e2e test_prefill test_kv_f16 default.profraw default.profdata

src/gguf.c

@@ -171,8 +171,8 @@ BnGGUFFile *bn_gguf_open(const uint8_t *buf, size_t size) {
 
     uint32_t magic = read_u32(&r);
     if (magic != BN_GGUF_MAGIC) {
-        char buf[16]; snprintf(buf, sizeof(buf), "0x%08x", magic);
-        SH_LOG_ERROR("Bad GGUF magic", "got", buf);
+        char hex[16]; snprintf(hex, sizeof(hex), "0x%08x", magic);
+        SH_LOG_ERROR("Bad GGUF magic", "got", hex);
         return NULL;
     }
 
@@ -186,8 +186,8 @@ BnGGUFFile *bn_gguf_open(const uint8_t *buf, size_t size) {
     f->alignment = BN_GGUF_DEFAULT_ALIGNMENT;
 
     if (f->version < 2 || f->version > 3) {
-        char buf[16]; snprintf(buf, sizeof(buf), "%u", f->version);
-        SH_LOG_ERROR("Unsupported GGUF version", "version", buf);
+        char ver[16]; snprintf(ver, sizeof(ver), "%u", f->version);
+        SH_LOG_ERROR("Unsupported GGUF version", "version", ver);
         free(f);
         return NULL;
     }
@@ -349,10 +349,39 @@ int bn_gguf_find_tensor(BnGGUFFile *f, const char *name) {
     return -1;
 }
 
-// #13: Validate that tensor data pointer falls within the mapped buffer
+// Compute byte size for a tensor given its type and element count.
+// Returns 0 for unknown types.
+static size_t tensor_type_size(uint32_t type, uint64_t nelements) {
+    switch (type) {
+        case BN_GGUF_TENSOR_F32:   return (size_t)nelements * 4;
+        case BN_GGUF_TENSOR_F16:   return (size_t)nelements * 2;
+        // I2_S: 2 bits per element + 4-byte per-tensor scale
+        case BN_GGUF_TENSOR_I2_S:  return (size_t)(nelements / 4) + 4;
+        // TQ1_0: 54 bytes per 256-element block
+        case BN_GGUF_TENSOR_TQ1_0: return (size_t)(nelements / 256) * 54;
+        // TQ2_0: 66 bytes per 256-element block
+        case BN_GGUF_TENSOR_TQ2_0: return (size_t)(nelements / 256) * 66;
+        default: return 0;
+    }
+}
+
+// #13: Validate that tensor data falls entirely within the mapped buffer
 void *bn_gguf_tensor_data(BnGGUFFile *f, int idx) {
     if (idx < 0 || (uint64_t)idx >= f->n_tensors) return NULL;
-    size_t offset = f->data_offset + f->tensors[idx].offset;
+    BnGGUFTensorInfo *t = &f->tensors[idx];
+    size_t offset = f->data_offset + t->offset;
     if (offset >= f->raw_size) return NULL;
+
+    // Compute total elements from dims
+    uint64_t nelements = 1;
+    for (uint32_t d = 0; d < t->n_dims; d++)
+        nelements *= t->dims[d];
+
+    size_t tsize = tensor_type_size(t->type, nelements);
+    if (tsize > 0 && offset + tsize > f->raw_size) {
+        SH_LOG_ERROR("Tensor data exceeds buffer", "tensor", t->name ? t->name : "?");
+        return NULL;
+    }
+
     return f->raw + offset;
 }

src/main.c

@@ -10,6 +10,10 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <limits.h>
+
+// Callback for streaming token output. Return non-zero to stop generation.
+typedef int (*bn_token_callback)(const char *piece, int token_id, void *user_data);
 
 #if defined(__APPLE__)
 #include <sys/sysctl.h>
@@ -49,6 +53,26 @@ static void print_usage(const char *prog) {
     fprintf(stderr, "  --no-prefill    Disable batch prompt prefill (compute logits for every token)\n");
 }
 
+static int parse_int(const char *s, const char *name) {
+    char *end;
+    long val = strtol(s, &end, 10);
+    if (*end != '\0' || val < INT_MIN || val > INT_MAX) {
+        fprintf(stderr, "Invalid value for %s: %s\n", name, s);
+        exit(1);
+    }
+    return (int)val;
+}
+
+static float parse_float(const char *s, const char *name) {
+    char *end;
+    float val = strtof(s, &end);
+    if (*end != '\0') {
+        fprintf(stderr, "Invalid value for %s: %s\n", name, s);
+        exit(1);
+    }
+    return val;
+}
+
 static CLIArgs parse_args(int argc, char **argv) {
     CLIArgs args = {0};
     args.prompt = "Hello";
@@ -69,16 +93,18 @@ static CLIArgs parse_args(int argc, char **argv) {
         if (strcmp(argv[i], "-p") == 0 && i + 1 < argc) {
             args.prompt = argv[++i];
         } else if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) {
-            args.n_tokens = atoi(argv[++i]);
+            args.n_tokens = parse_int(argv[++i], "-n");
         } else if (strcmp(argv[i], "--temp") == 0 && i + 1 < argc) {
-            args.temperature = (float)atof(argv[++i]);
+            args.temperature = parse_float(argv[++i], "--temp");
             args.temp_set = 1;
         } else if (strcmp(argv[i], "--topp") == 0 && i + 1 < argc) {
-            args.topp = (float)atof(argv[++i]);
+            args.topp = parse_float(argv[++i], "--topp");
         } else if (strcmp(argv[i], "--seed") == 0 && i + 1 < argc) {
-            args.seed = (uint64_t)atoll(argv[++i]);
+            char *end;
+            args.seed = (uint64_t)strtoull(argv[++i], &end, 10);
+            if (*end != '\0') { fprintf(stderr, "Invalid value for --seed: %s\n", argv[i]); exit(1); }
         } else if (strcmp(argv[i], "--maxseq") == 0 && i + 1 < argc) {
-            args.max_seq_len = atoi(argv[++i]);
+            args.max_seq_len = parse_int(argv[++i], "--maxseq");
         } else if (strcmp(argv[i], "--flash") == 0) {
             args.flash_attn = 1;
         } else if (strcmp(argv[i], "--chat") == 0) {
@@ -88,7 +114,7 @@ static CLIArgs parse_args(int argc, char **argv) {
         } else if (strcmp(argv[i], "--no-prefill") == 0) {
             args.no_prefill = 1;
         } else if (strcmp(argv[i], "--repeat-penalty") == 0 && i + 1 < argc) {
-            args.repeat_penalty = (float)atof(argv[++i]);
+            args.repeat_penalty = parse_float(argv[++i], "--repeat-penalty");
             args.repeat_set = 1;
         } else {
             fprintf(stderr, "Unknown option: %s\n", argv[i]);
@@ -100,6 +126,67 @@ static CLIArgs parse_args(int argc, char **argv) {
     return args;
 }
 
+// Loop detection constants
+#define LOOP_BUF_SIZE 32
+#define LOOP_NGRAM    4
+
+// Generate tokens with callback-based streaming.
+// Returns: n_generated, -1 on loop, -2 on error.
+static int generate_response(BnModel *model, BnTokenizer *tok, BnSampler *sampler,
+                              int max_tokens, int *pos,
+                              bn_token_callback cb, void *user_data) {
+    int loop_buf[LOOP_BUF_SIZE];
+    int loop_idx = 0, gen_count = 0;
+    memset(loop_buf, -1, sizeof(loop_buf));
+
+    float *logits = model->state.logits;
+    if (!logits) return -2;
+
+    for (int i = 0; i < max_tokens; i++) {
+        int next = bn_sampler_sample(sampler, logits);
+
+        if (next == tok->eot_id || next == tok->eos_id)
+            break;
+
+        // Ring buffer loop detection
+        loop_buf[loop_idx] = next;
+        loop_idx = (loop_idx + 1) % LOOP_BUF_SIZE;
+        gen_count++;
+
+        if (gen_count >= 2 * LOOP_NGRAM) {
+            int looping = 1;
+            for (int k = 0; k < LOOP_NGRAM; k++) {
+                int a = loop_buf[((loop_idx - 1 - k) % LOOP_BUF_SIZE + LOOP_BUF_SIZE) % LOOP_BUF_SIZE];
+                int b = loop_buf[((loop_idx - 1 - k - LOOP_NGRAM) % LOOP_BUF_SIZE + LOOP_BUF_SIZE) % LOOP_BUF_SIZE];
+                if (a != b) { looping = 0; break; }
+            }
+            if (looping) return -1;
+        }
+
+        bn_sampler_accept(sampler, next);
+
+        const char *piece = bn_tokenizer_decode(tok, next);
+        if (piece && cb) {
+            if (cb(piece, next, user_data))
+                break;
+        }
+
+        logits = bn_transformer_forward(model, next, *pos);
+        (*pos)++;
+        if (!logits) return -2;
+    }
+
+    return gen_count;
+}
+
+static int print_token(const char *piece, int token_id, void *user_data) {
+    (void)token_id;
+    (void)user_data;
+    printf("%s", piece);
+    fflush(stdout);
+    return 0;
+}
+
 int main(int argc, char **argv) {
     sh_log_init(NULL);
     CLIArgs args = parse_args(argc, argv);
@@ -187,7 +274,6 @@ int main(int argc, char **argv) {
         SH_LOG_ERROR("Failed to init tokenizer");
         bn_model_free(&model);
         bn_gguf_free(gf);
-        bn_platform_unload_file(&mf);
         return 1;
     }
     {
@@ -220,7 +306,6 @@ int main(int argc, char **argv) {
             bn_tokenizer_free(&tokenizer);
             bn_model_free(&model);
             bn_gguf_free(gf);
-            bn_platform_unload_file(&mf);
             return 1;
         }
 
@@ -287,48 +372,9 @@ int main(int argc, char **argv) {
                 break;
             }
 
-            // Generate until eot_id, eos_id, or seq_len
-            // Loop detector: ring buffer of recent tokens, check for repeating n-grams
-            #define LOOP_BUF_SIZE 32
-            #define LOOP_NGRAM    4
-            int loop_buf[LOOP_BUF_SIZE];
-            int loop_idx = 0, gen_count = 0;
-            memset(loop_buf, -1, sizeof(loop_buf));
-
-            for (int i = 0; i < args.n_tokens; i++) {
-                int next = bn_sampler_sample(&sampler, logits);
-
-                if (next == tokenizer.eot_id || next == tokenizer.eos_id)
-                    break;
-
-                // Record token in ring buffer and check for loops
-                loop_buf[loop_idx] = next;
-                loop_idx = (loop_idx + 1) % LOOP_BUF_SIZE;
-                gen_count++;
-
-                if (gen_count >= 2 * LOOP_NGRAM) {
-                    // Check if last LOOP_NGRAM tokens match the LOOP_NGRAM before them
-                    int looping = 1;
-                    for (int k = 0; k < LOOP_NGRAM; k++) {
-                        int a = loop_buf[((loop_idx - 1 - k) % LOOP_BUF_SIZE + LOOP_BUF_SIZE) % LOOP_BUF_SIZE];
-                        int b = loop_buf[((loop_idx - 1 - k - LOOP_NGRAM) % LOOP_BUF_SIZE + LOOP_BUF_SIZE) % LOOP_BUF_SIZE];
-                        if (a != b) { looping = 0; break; }
-                    }
-                    if (looping) { gen_count = -1; break; }
-                }
-
-                bn_sampler_accept(&sampler, next);
-
-                const char *piece = bn_tokenizer_decode(&tokenizer, next);
-                if (piece) {
-                    printf("%s", piece);
-                    fflush(stdout);
-                }
-
-                logits = bn_transformer_forward(&model, next, pos);
-                pos++;
-                if (!logits) break;
-            }
+            int gen_count = generate_response(&model, &tokenizer, &sampler,
+                                               args.n_tokens, &pos,
+                                               print_token, NULL);
 
             // Feed EOT into KV cache to close the assistant turn
             bn_transformer_forward(&model, tokenizer.eot_id, pos);
@@ -337,7 +383,7 @@ int main(int argc, char **argv) {
             printf("\n");
 
             turn_count++;
-            int should_reset = (turn_count >= 2) || (gen_count == -1);
+            int should_reset = (turn_count >= 2) || (gen_count < 0);
             if (should_reset) {
                 printf("[auto-reset: starting fresh]\n");
                 pos = 0;
@@ -360,7 +406,6 @@ int main(int argc, char **argv) {
             bn_tokenizer_free(&tokenizer);
             bn_model_free(&model);
             bn_gguf_free(gf);
-            bn_platform_unload_file(&mf);
             return 1;
         }
         int n_prompt = bn_tokenizer_encode(&tokenizer, args.prompt, 1, prompt_tokens,
@@ -442,9 +487,8 @@ int main(int argc, char **argv) {
     // Cleanup
     bn_sampler_free(&sampler);
     bn_tokenizer_free(&tokenizer);
-    bn_model_free(&model);
+    bn_model_free(&model);  // also unloads mmap'd file
     bn_gguf_free(gf);
-    bn_platform_unload_file(&mf);
 
     return 0;
 }

src/model.c

@@ -346,6 +346,7 @@ void bn_model_free(BnModel *m) {
     bn_tp_free(m->pool);
     free(m->weights.layers);
     sh_arena_free(m->arena);  // frees INT8 embeddings too
+    bn_platform_unload_file(&m->file);  // safe even if file.data is NULL
     memset(m, 0, sizeof(BnModel));
 }
 

src/platform.c

@@ -69,8 +69,10 @@ void bn_platform_unload_file(BnMappedFile *f) {
 #else
     if (f->is_mmap == 1) {
         munmap(f->data, f->size);
+    } else if (f->is_mmap == 0) {
+        free(f->data);
     }
-    // is_mmap == 0 or 2: don't free (external buffer or unused)
+    // is_mmap == 2: externally owned, don't free
 #endif
     f->data = NULL;
     f->size = 0;

src/quant.c

@@ -671,7 +671,12 @@ void bn_quant_matvec_batch(const BnMatvecTask *tasks, int n_tasks,
     }
 
     if (all_i2s) {
-        assert(n_tasks <= 4 && "bn_quant_matvec_batch: max 4 tasks");
+        // Fall back to individual matvecs if too many tasks for stack arrays
+        if (n_tasks > 4) {
+            for (int t = 0; t < n_tasks; t++)
+                bn_quant_matvec(tasks[t].out, tasks[t].W, x, x_q_buf, pool);
+            return;
+        }
 
         // Quantize x to int8 once, shared across all tasks
         float x_scale = bn_quant_x_to_i8(x, x_q_buf, cols);