From 96fdb9bd1607d034d5ec5c824ad9e6a58a53efd9 Mon Sep 17 00:00:00 2001
From: Neil Johari <njohari@roblox.com>
Date: Mon, 15 Sep 2025 23:58:45 -0700
Subject: [PATCH 1/3] Add debug logging and simple repro

---
 programs/dibio.c       | 20 +++++++++++++++++++
 test_allocation_bug.sh | 44 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)
 create mode 100755 test_allocation_bug.sh

diff --git a/programs/dibio.c b/programs/dibio.c
index 63c455a2216..f5f33f1bd5a 100644
--- a/programs/dibio.c
+++ b/programs/dibio.c
@@ -279,10 +279,16 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
 
     for (n=0; n<nbFiles; n++) {
       S64 const fileSize = DiB_getFileSize(fileNamesTable[n]);
+      DISPLAYLEVEL(1, "[DEBUG] File '%s': size=%lld\n", fileNamesTable[n], (long long)fileSize);
+      
       /* TODO: is there a minimum sample size? What if the file is 1-byte? */
       if (fileSize == 0) {
         DISPLAYLEVEL(3, "Sample file '%s' has zero size, skipping...\n", fileNamesTable[n]);
         continue;
+      } else if (fileSize < 0) {
+        /* BUG: This path is NOT skipped but should be! */
+        DISPLAYLEVEL(1, "[BUG] File '%s' has NEGATIVE size %lld but is NOT skipped!\n", 
+                     fileNamesTable[n], (long long)fileSize);
       }
 
       /* the case where we are breaking up files in sample chunks */
@@ -290,6 +296,8 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
         /* TODO: is there a minimum sample size? Can we have a 1-byte sample? */
         fs.nbSamples += (int)((fileSize + chunkSize-1) / chunkSize);
         fs.totalSizeToLoad += fileSize;
+        DISPLAYLEVEL(1, "[DEBUG] After chunked file: nbSamples=%d, totalSizeToLoad=%lld\n", 
+                     fs.nbSamples, (long long)fs.totalSizeToLoad);
       }
       else {
       /* the case where one file is one sample */
@@ -303,9 +311,14 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
         }
         fs.nbSamples += 1;
         fs.totalSizeToLoad += MIN(fileSize, SAMPLESIZE_MAX);
+        DISPLAYLEVEL(1, "[DEBUG] After single file: nbSamples=%d, totalSizeToLoad=%lld\n", 
+                     fs.nbSamples, (long long)fs.totalSizeToLoad);
       }
     }
     DISPLAYLEVEL(4, "Found training data %d files, %d KB, %d samples\n", nbFiles, (int)(fs.totalSizeToLoad / (1 KB)), fs.nbSamples);
+    DISPLAYLEVEL(1, "[DEBUG FINAL] fileStats: nbSamples=%d, totalSizeToLoad=%lld (%s)\n", 
+                 fs.nbSamples, (long long)fs.totalSizeToLoad,
+                 fs.totalSizeToLoad < 0 ? "NEGATIVE!" : "ok");
     return fs;
 }
 
@@ -344,11 +357,18 @@ int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
         /* Limit the size of the training data to 2GB */
         /* TODO: there is opportunity to stop DiB_fileStats() early when the data limit is reached */
         loadedSize = (size_t)MIN( MIN((S64)maxMem, fs.totalSizeToLoad), MAX_SAMPLES_SIZE );
+        DISPLAYLEVEL(1, "[DEBUG] Memory calc: totalSizeToLoad=%lld, maxMem=%zu, loadedSize=%zu (0x%zx)\n",
+                     (long long)fs.totalSizeToLoad, maxMem, loadedSize, loadedSize);
+        if (fs.totalSizeToLoad < 0) {
+            DISPLAYLEVEL(1, "[BUG] totalSizeToLoad is NEGATIVE! This will cause allocation issues!\n");
+        }
         if (memLimit != 0) {
             DISPLAYLEVEL(2, "!  Warning : setting manual memory limit for dictionary training data at %u MB \n",
                 (unsigned)(memLimit / (1 MB)));
             loadedSize = (size_t)MIN(loadedSize, memLimit);
         }
+        DISPLAYLEVEL(1, "[DEBUG] About to malloc: srcBuffer size=%zu, sampleSizes array size=%zu\n",
+                     loadedSize+NOISELENGTH, (size_t)(fs.nbSamples * sizeof(size_t)));
         srcBuffer = malloc(loadedSize+NOISELENGTH);
         sampleSizes = (size_t*)malloc(fs.nbSamples * sizeof(size_t));
     }
diff --git a/test_allocation_bug.sh b/test_allocation_bug.sh
new file mode 100755
index 00000000000..d8301e68186
--- /dev/null
+++ b/test_allocation_bug.sh
@@ -0,0 +1,44 @@
+#!/bin/bash
+echo "=== Test to show allocation bug with negative totalSizeToLoad ==="
+echo ""
+echo "We need at least 5 samples to pass the minimum check"
+echo ""
+
+# Clean up
+rm -rf alloc_test
+mkdir alloc_test
+
+# Create exactly 5 valid files (minimum to not exit early)
+echo "Creating 5 valid files (minimum required)..."
+for i in {1..5}; do
+    echo "data$i" > alloc_test/good_$i.txt
+done
+
+echo "Valid files created (about 6 bytes each = 30 bytes total)"
+echo ""
+
+# We need enough bad files to make totalSizeToLoad negative
+# 30 bytes positive, so we need at least 31 bad files
+echo "Adding 1000 non-existent files to make totalSizeToLoad very negative..."
+echo "Expected: totalSizeToLoad = 30 + (1000 * -1) = -970 bytes"
+echo ""
+
+# Build command
+CMD="./zstd --train alloc_test/good_*.txt"
+for i in {1..1000}; do
+    CMD="$CMD alloc_test/BAD_$i"
+done
+CMD="$CMD -o alloc_test/dict.zst --maxdict=65536 2>&1"
+
+echo "Running command..."
+echo "================="
+
+# Run and capture ALL debug output related to our issue
+eval $CMD | grep -E "\[DEBUG FINAL\]|\[DEBUG\] Memory calc|\[BUG\]|About to malloc|Error|not enough memory" 
+
+echo ""
+echo "Output should show something like the following:"
+echo "1. [DEBUG FINAL] fileStats: totalSizeToLoad=-970 (NEGATIVE!)"
+echo "2. [BUG] totalSizeToLoad is NEGATIVE!"
+echo "3. [DEBUG] Memory calc: showing huge loadedSize value"
+echo "4. Error about memory allocation"

From 85f4a7e84e266d9545a309ff3f25e71345d8be62 Mon Sep 17 00:00:00 2001
From: Neil Johari <njohari@roblox.com>
Date: Tue, 16 Sep 2025 00:02:04 -0700
Subject: [PATCH 2/3] Fix bug

---
 programs/dibio.c | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/programs/dibio.c b/programs/dibio.c
index f5f33f1bd5a..29a27d9dbb4 100644
--- a/programs/dibio.c
+++ b/programs/dibio.c
@@ -282,13 +282,16 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
       DISPLAYLEVEL(1, "[DEBUG] File '%s': size=%lld\n", fileNamesTable[n], (long long)fileSize);
       
       /* TODO: is there a minimum sample size? What if the file is 1-byte? */
-      if (fileSize == 0) {
-        DISPLAYLEVEL(3, "Sample file '%s' has zero size, skipping...\n", fileNamesTable[n]);
+      /* Skip empty or invalid files */
+      if (fileSize <= 0) {
+        if (fileSize < 0) {
+          DISPLAYLEVEL(3, "Sample file '%s' is unreadable or stat failed, skipping...\n",
+                       fileNamesTable[n]);
+        } else {
+          DISPLAYLEVEL(3, "Sample file '%s' has zero size, skipping...\n",
+                     fileNamesTable[n]);
+        }
         continue;
-      } else if (fileSize < 0) {
-        /* BUG: This path is NOT skipped but should be! */
-        DISPLAYLEVEL(1, "[BUG] File '%s' has NEGATIVE size %lld but is NOT skipped!\n", 
-                     fileNamesTable[n], (long long)fileSize);
       }
 
       /* the case where we are breaking up files in sample chunks */

From 236e44f00fed5cbcd02b4936b765f3bb97ad0e98 Mon Sep 17 00:00:00 2001
From: Neil Johari <njohari@roblox.com>
Date: Tue, 16 Sep 2025 00:03:08 -0700
Subject: [PATCH 3/3] Remove debug logging

---
 programs/dibio.c       | 16 ---------------
 test_allocation_bug.sh | 44 ------------------------------------------
 2 files changed, 60 deletions(-)
 delete mode 100755 test_allocation_bug.sh

diff --git a/programs/dibio.c b/programs/dibio.c
index 29a27d9dbb4..dc629103ba2 100644
--- a/programs/dibio.c
+++ b/programs/dibio.c
@@ -279,8 +279,6 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
 
     for (n=0; n<nbFiles; n++) {
       S64 const fileSize = DiB_getFileSize(fileNamesTable[n]);
-      DISPLAYLEVEL(1, "[DEBUG] File '%s': size=%lld\n", fileNamesTable[n], (long long)fileSize);
-      
       /* TODO: is there a minimum sample size? What if the file is 1-byte? */
       /* Skip empty or invalid files */
       if (fileSize <= 0) {
@@ -299,8 +297,6 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
         /* TODO: is there a minimum sample size? Can we have a 1-byte sample? */
         fs.nbSamples += (int)((fileSize + chunkSize-1) / chunkSize);
         fs.totalSizeToLoad += fileSize;
-        DISPLAYLEVEL(1, "[DEBUG] After chunked file: nbSamples=%d, totalSizeToLoad=%lld\n", 
-                     fs.nbSamples, (long long)fs.totalSizeToLoad);
       }
       else {
       /* the case where one file is one sample */
@@ -314,14 +310,9 @@ static fileStats DiB_fileStats(const char** fileNamesTable, int nbFiles, size_t
         }
         fs.nbSamples += 1;
         fs.totalSizeToLoad += MIN(fileSize, SAMPLESIZE_MAX);
-        DISPLAYLEVEL(1, "[DEBUG] After single file: nbSamples=%d, totalSizeToLoad=%lld\n", 
-                     fs.nbSamples, (long long)fs.totalSizeToLoad);
       }
     }
     DISPLAYLEVEL(4, "Found training data %d files, %d KB, %d samples\n", nbFiles, (int)(fs.totalSizeToLoad / (1 KB)), fs.nbSamples);
-    DISPLAYLEVEL(1, "[DEBUG FINAL] fileStats: nbSamples=%d, totalSizeToLoad=%lld (%s)\n", 
-                 fs.nbSamples, (long long)fs.totalSizeToLoad,
-                 fs.totalSizeToLoad < 0 ? "NEGATIVE!" : "ok");
     return fs;
 }
 
@@ -360,18 +351,11 @@ int DiB_trainFromFiles(const char* dictFileName, size_t maxDictSize,
         /* Limit the size of the training data to 2GB */
         /* TODO: there is opportunity to stop DiB_fileStats() early when the data limit is reached */
         loadedSize = (size_t)MIN( MIN((S64)maxMem, fs.totalSizeToLoad), MAX_SAMPLES_SIZE );
-        DISPLAYLEVEL(1, "[DEBUG] Memory calc: totalSizeToLoad=%lld, maxMem=%zu, loadedSize=%zu (0x%zx)\n",
-                     (long long)fs.totalSizeToLoad, maxMem, loadedSize, loadedSize);
-        if (fs.totalSizeToLoad < 0) {
-            DISPLAYLEVEL(1, "[BUG] totalSizeToLoad is NEGATIVE! This will cause allocation issues!\n");
-        }
         if (memLimit != 0) {
             DISPLAYLEVEL(2, "!  Warning : setting manual memory limit for dictionary training data at %u MB \n",
                 (unsigned)(memLimit / (1 MB)));
             loadedSize = (size_t)MIN(loadedSize, memLimit);
         }
-        DISPLAYLEVEL(1, "[DEBUG] About to malloc: srcBuffer size=%zu, sampleSizes array size=%zu\n",
-                     loadedSize+NOISELENGTH, (size_t)(fs.nbSamples * sizeof(size_t)));
         srcBuffer = malloc(loadedSize+NOISELENGTH);
         sampleSizes = (size_t*)malloc(fs.nbSamples * sizeof(size_t));
     }
diff --git a/test_allocation_bug.sh b/test_allocation_bug.sh
deleted file mode 100755
index d8301e68186..00000000000
--- a/test_allocation_bug.sh
+++ /dev/null
@@ -1,44 +0,0 @@
-#!/bin/bash
-echo "=== Test to show allocation bug with negative totalSizeToLoad ==="
-echo ""
-echo "We need at least 5 samples to pass the minimum check"
-echo ""
-
-# Clean up
-rm -rf alloc_test
-mkdir alloc_test
-
-# Create exactly 5 valid files (minimum to not exit early)
-echo "Creating 5 valid files (minimum required)..."
-for i in {1..5}; do
-    echo "data$i" > alloc_test/good_$i.txt
-done
-
-echo "Valid files created (about 6 bytes each = 30 bytes total)"
-echo ""
-
-# We need enough bad files to make totalSizeToLoad negative
-# 30 bytes positive, so we need at least 31 bad files
-echo "Adding 1000 non-existent files to make totalSizeToLoad very negative..."
-echo "Expected: totalSizeToLoad = 30 + (1000 * -1) = -970 bytes"
-echo ""
-
-# Build command
-CMD="./zstd --train alloc_test/good_*.txt"
-for i in {1..1000}; do
-    CMD="$CMD alloc_test/BAD_$i"
-done
-CMD="$CMD -o alloc_test/dict.zst --maxdict=65536 2>&1"
-
-echo "Running command..."
-echo "================="
-
-# Run and capture ALL debug output related to our issue
-eval $CMD | grep -E "\[DEBUG FINAL\]|\[DEBUG\] Memory calc|\[BUG\]|About to malloc|Error|not enough memory" 
-
-echo ""
-echo "Output should show something like the following:"
-echo "1. [DEBUG FINAL] fileStats: totalSizeToLoad=-970 (NEGATIVE!)"
-echo "2. [BUG] totalSizeToLoad is NEGATIVE!"
-echo "3. [DEBUG] Memory calc: showing huge loadedSize value"
-echo "4. Error about memory allocation"