From da737c7ab89f61f1ea7c392299137f6ffe6f9733 Mon Sep 17 00:00:00 2001
From: Nick Terrell <terrelln@fb.com>
Date: Wed, 2 Mar 2022 11:04:04 -0800
Subject: [PATCH] [programs] Fix infinite loop when empty input is passed to
 trainer

When an empty input file was passed to the dictionary trainer, it would infinite loop.
The added test case exposes the bug, and is fixed with this PR.
---
 programs/dibio.c                                         | 5 ++++-
 tests/cli-tests/dict-builder/empty-input.sh              | 9 +++++++++
 tests/cli-tests/dict-builder/empty-input.sh.stderr.exact | 1 +
 3 files changed, 14 insertions(+), 1 deletion(-)
 create mode 100755 tests/cli-tests/dict-builder/empty-input.sh
 create mode 100644 tests/cli-tests/dict-builder/empty-input.sh.stderr.exact

diff --git a/programs/dibio.c b/programs/dibio.c
index fddbc9e5769..f6757dd3488 100644
--- a/programs/dibio.c
+++ b/programs/dibio.c
@@ -128,8 +128,11 @@ static int DiB_loadFiles(
     while ( nbSamplesLoaded < sstSize && fileIndex < nbFiles ) {
         size_t fileDataLoaded;
         S64 const fileSize = DiB_getFileSize(fileNamesTable[fileIndex]);
-        if (fileSize <= 0) /* skip if zero-size or file error */
+        if (fileSize <= 0) {
+            /* skip if zero-size or file error */
+            ++fileIndex;
             continue;
+        }
 
         f = fopen( fileNamesTable[fileIndex], "rb");
         if (f == NULL)
diff --git a/tests/cli-tests/dict-builder/empty-input.sh b/tests/cli-tests/dict-builder/empty-input.sh
new file mode 100755
index 00000000000..b500bfebda0
--- /dev/null
+++ b/tests/cli-tests/dict-builder/empty-input.sh
@@ -0,0 +1,9 @@
+#!/bin/sh
+set -e
+for i in $(seq 50); do
+    datagen -s$i > file$i
+done
+touch empty
+
+set -v
+zstd -q --train empty file*
diff --git a/tests/cli-tests/dict-builder/empty-input.sh.stderr.exact b/tests/cli-tests/dict-builder/empty-input.sh.stderr.exact
new file mode 100644
index 00000000000..2747e766cdf
--- /dev/null
+++ b/tests/cli-tests/dict-builder/empty-input.sh.stderr.exact
@@ -0,0 +1 @@
+zstd -q --train empty file*