2 files changed, 6 insertions, 3 deletions
diff --git a/src/dft.c b/src/dft.c
index faaca9e..7b95518 100644
--- a/src/dft.c
+++ b/src/dft.c
@@ -1,4 +1,5 @@
 #include "dft.h"
+#include "common.h"
 #include "vec2.h"
 
 #include <string.h>
@@ -147,8 +148,8 @@ wav_data_t dft_data_to_wav(smrt_arena_t *arena, dft_data_t dft, u64 sample_rate,
 stft_data_t short_time_fourier_transform(smrt_arena_t *arena, u64 window_size, u64 hop_size, f64 *samples, u64 sample_count, u64 sample_rate) {
     assert(F64_EQ(round(log2(window_size)), log2(window_size), 1e-9) &&
            "STFT input window_size must be a power of 2");
-    assert(sample_count % window_size == 0 &&
-           "Sample count must evenly divide into window_size chunks");
+
+    sample_count = ALIGN_UP_POW2(sample_count, window_size);
 
     u64 segment_count = ((sample_count - window_size) / hop_size) + 1;
 
diff --git a/src/dft.h b/src/dft.h
index 74178aa..aaa2796 100644
--- a/src/dft.h
+++ b/src/dft.h
@@ -43,10 +43,12 @@ typedef struct {
 } stft_data_t;
 
 /// Run STFT algorithm on samples, sliding a window_size window by hop_size each step
+/// Warning: This function assumes that samples' allocation is large enough to fit
+/// ALIGN_UP_POW2(sample_count, window_size) elements.
 stft_data_t short_time_fourier_transform(smrt_arena_t *        arena ,
                                                     u64  window_size ,
                                                     u64     hop_size ,
-                                                  f64 *      samples ,
+                                                    f64 *    samples ,
                                                     u64 sample_count ,
                                                     u64  sample_rate);