Fix use of references to buffers after calling dav1d_close()

9057d286 had the side effect of causing references to buffers allocated
using memory pools to no longer be valid after closing the decoder.

Restore this functionality by making buffer pools reference counted.

1 files changed, 58 insertions, 13 deletions

diff --git a/src/mem.c b/src/mem.c
index 8728247..558bc01 100644
--- a/src/mem.c
+++ b/src/mem.c
@@ -29,25 +29,40 @@
 
 #include <stdint.h>
 
-#include "src/mem.h"
-#include "src/thread.h"
+#include "src/internal.h"
+
+static COLD void mem_pool_destroy(Dav1dMemPool *const pool) {
+    pthread_mutex_destroy(&pool->lock);
+    free(pool);
+}
 
 void dav1d_mem_pool_push(Dav1dMemPool *const pool, Dav1dMemPoolBuffer *const buf) {
     pthread_mutex_lock(&pool->lock);
-    buf->next = pool->buf;
-    pool->buf = buf;
-    pthread_mutex_unlock(&pool->lock);
+    const int ref_cnt = --pool->ref_cnt;
+    if (!pool->end) {
+        buf->next = pool->buf;
+        pool->buf = buf;
+        pthread_mutex_unlock(&pool->lock);
+        assert(ref_cnt > 0);
+    } else {
+        pthread_mutex_unlock(&pool->lock);
+        dav1d_free_aligned(buf->data);
+        if (!ref_cnt) mem_pool_destroy(pool);
+    }
 }
 
 Dav1dMemPoolBuffer *dav1d_mem_pool_pop(Dav1dMemPool *const pool, const size_t size) {
+    assert(!(size & (sizeof(void*) - 1)));
     pthread_mutex_lock(&pool->lock);
     Dav1dMemPoolBuffer *buf = pool->buf;
+    pool->ref_cnt++;
     uint8_t *data;
     if (buf) {
         pool->buf = buf->next;
         pthread_mutex_unlock(&pool->lock);
         data = buf->data;
         if ((uintptr_t)buf - (uintptr_t)data != size) {
+            /* Reallocate if the size has changed */
             dav1d_free_aligned(data);
             goto alloc;
         }
@@ -55,7 +70,13 @@ Dav1dMemPoolBuffer *dav1d_mem_pool_pop(Dav1dMemPool *const pool, const size_t si
         pthread_mutex_unlock(&pool->lock);
 alloc:
         data = dav1d_alloc_aligned(size + sizeof(Dav1dMemPoolBuffer), 64);
-        if (!data) return NULL;
+        if (!data) {
+            pthread_mutex_lock(&pool->lock);
+            const int ref_cnt = --pool->ref_cnt;
+            pthread_mutex_unlock(&pool->lock);
+            if (!ref_cnt) mem_pool_destroy(pool);
+            return NULL;
+        }
         buf = (Dav1dMemPoolBuffer*)(data + size);
         buf->data = data;
     }
@@ -63,12 +84,36 @@ alloc:
     return buf;
 }
 
-COLD void dav1d_mem_pool_destroy(Dav1dMemPool *const pool) {
-    pthread_mutex_destroy(&pool->lock);
-    Dav1dMemPoolBuffer *buf = pool->buf;
-    while (buf) {
-        void *const data = buf->data;
-        buf = buf->next;
-        dav1d_free_aligned(data);
+COLD int dav1d_mem_pool_init(Dav1dMemPool **const ppool) {
+    Dav1dMemPool *const pool = malloc(sizeof(Dav1dMemPool));
+    if (pool) {
+        if (!pthread_mutex_init(&pool->lock, NULL)) {
+            pool->buf = NULL;
+            pool->ref_cnt = 1;
+            pool->end = 0;
+            *ppool = pool;
+            return 0;
+        }
+        free(pool);
+    }
+    *ppool = NULL;
+    return DAV1D_ERR(ENOMEM);
+}
+
+COLD void dav1d_mem_pool_end(Dav1dMemPool *const pool) {
+    if (pool) {
+        pthread_mutex_lock(&pool->lock);
+        Dav1dMemPoolBuffer *buf = pool->buf;
+        const int ref_cnt = --pool->ref_cnt;
+        pool->buf = NULL;
+        pool->end = 1;
+        pthread_mutex_unlock(&pool->lock);
+
+        while (buf) {
+            void *const data = buf->data;
+            buf = buf->next;
+            dav1d_free_aligned(data);
+        }
+        if (!ref_cnt) mem_pool_destroy(pool);
     }
 }