DRW: InstanceData: Remove hacks of batch freeing callback

We instead use a handle reference counter on the GPUVertBufs used by
the instancing batches. This make sure that if an update happens on the
GPUVertBuf used to contruct the batch, they will never have the same
memory address than the previously allocated ones (since they are still
pending deletion thanks to the refcounter).

This avoid the linear search to update the GPUBatch in the case a
batch is deleted (which was even a bad option since they could be only
cleared)

3 files changed, 74 insertions, 85 deletions

diff --git a/source/blender/draw/intern/draw_instance_data.c b/source/blender/draw/intern/draw_instance_data.c
index a40fc1726d7..24f17f313c4 100644
--- a/source/blender/draw/intern/draw_instance_data.c
+++ b/source/blender/draw/intern/draw_instance_data.c
@@ -59,50 +59,50 @@ struct DRWInstanceDataList {
 };
 
 typedef struct DRWTempBufferHandle {
-  /** Must be first for casting. */
-  GPUVertBuf buf;
+  GPUVertBuf *buf;
   /** Format pointer for reuse. */
   GPUVertFormat *format;
   /** Touched vertex length for resize. */
   int *vert_len;
 } DRWTempBufferHandle;
 
-static ListBase g_idatalists = {NULL, NULL};
+typedef struct DRWTempInstancingHandle {
+  /** Copy of geom but with the per-instance attributes. */
+  GPUBatch *batch;
+  /** Batch containing instancing attributes. */
+  GPUBatch *instancer;
+  /** Callbuffer to be used instead of instancer . */
+  GPUVertBuf *buf;
+  /** Original non-instanced batch pointer. */
+  GPUBatch *geom;
+} DRWTempInstancingHandle;
 
-/* -------------------------------------------------------------------- */
-/** \name Instance Buffer Management
- * \{ */
+static ListBase g_idatalists = {NULL, NULL};
 
-static void instance_batch_free(GPUBatch *geom, void *UNUSED(user_data))
+static void instancing_batch_references_add(GPUBatch *batch)
 {
-  if (geom->verts[0] == NULL) {
-    /** XXX This is a false positive case.
-     * The batch has been requested but not init yet
-     * and there is a chance that it might become init.
-     */
-    return;
+  for (int i = 0; i < GPU_BATCH_VBO_MAX_LEN && batch->verts[i]; i++) {
+    GPU_vertbuf_handle_ref_add(batch->verts[i]);
+  }
+  for (int i = 0; i < GPU_BATCH_INST_VBO_MAX_LEN && batch->inst[i]; i++) {
+    GPU_vertbuf_handle_ref_add(batch->inst[i]);
   }
+}
 
-  /* Free all batches that use the same vbos before they are reused. */
-  /* TODO: Make it thread safe! Batch freeing can happen from another thread. */
-  /* FIXME: This is not really correct. The correct way would be to check based on
-   * the vertex buffers. We assume the batch containing the VBO is being when it should. */
-  /* PERF: This is doing a linear search. This can be very costly. */
-  LISTBASE_FOREACH (DRWInstanceDataList *, data_list, &g_idatalists) {
-    BLI_memblock *memblock = data_list->pool_instancing;
-    BLI_memblock_iter iter;
-    BLI_memblock_iternew(memblock, &iter);
-    GPUBatch **batch_ptr;
-    while ((batch_ptr = (GPUBatch **)BLI_memblock_iterstep(&iter))) {
-      GPUBatch *batch = *batch_ptr;
-      /* Only check verts[0] that's enough. */
-      if (batch->verts[0] == geom->verts[0]) {
-        GPU_batch_clear(batch);
-      }
-    }
+static void instancing_batch_references_remove(GPUBatch *batch)
+{
+  for (int i = 0; i < GPU_BATCH_VBO_MAX_LEN && batch->verts[i]; i++) {
+    GPU_vertbuf_handle_ref_remove(batch->verts[i]);
+  }
+  for (int i = 0; i < GPU_BATCH_INST_VBO_MAX_LEN && batch->inst[i]; i++) {
+    GPU_vertbuf_handle_ref_remove(batch->inst[i]);
   }
 }
 
+/* -------------------------------------------------------------------- */
+/** \name Instance Buffer Management
+ * \{ */
+
 /**
  * This manager allows to distribute existing batches for instancing
  * attributes. This reduce the number of batches creation.
@@ -119,20 +119,23 @@ GPUVertBuf *DRW_temp_buffer_request(DRWInstanceDataList *idatalist,
   BLI_assert(vert_len != NULL);
 
   DRWTempBufferHandle *handle = BLI_memblock_alloc(idatalist->pool_buffers);
-  GPUVertBuf *vert = &handle->buf;
-  handle->vert_len = vert_len;
 
   if (handle->format != format) {
     handle->format = format;
-    /* TODO/PERF: Save the allocated data from freeing to avoid reallocation. */
-    GPU_vertbuf_clear(vert);
+    GPU_VERTBUF_DISCARD_SAFE(handle->buf);
+
+    GPUVertBuf *vert = GPU_vertbuf_create(GPU_USAGE_DYNAMIC);
     GPU_vertbuf_init_with_format_ex(vert, format, GPU_USAGE_DYNAMIC);
     GPU_vertbuf_data_alloc(vert, DRW_BUFFER_VERTS_CHUNK);
+
+    handle->buf = vert;
   }
-  return vert;
+  handle->vert_len = vert_len;
+  return handle->buf;
 }
 
-/* NOTE: Does not return a valid drawable batch until DRW_instance_buffer_finish has run. */
+/* NOTE: Does not return a valid drawable batch until DRW_instance_buffer_finish has run.
+ * Initialization is delayed because instancer or geom could still not be initialized. */
 GPUBatch *DRW_temp_batch_instance_request(DRWInstanceDataList *idatalist,
                                           GPUVertBuf *buf,
                                           GPUBatch *instancer,
@@ -143,15 +146,15 @@ GPUBatch *DRW_temp_batch_instance_request(DRWInstanceDataList *idatalist,
   /* Only call with one of them. */
   BLI_assert((instancer != NULL) != (buf != NULL));
 
-  GPUBatch **batch_ptr = BLI_memblock_alloc(idatalist->pool_instancing);
-  if (*batch_ptr == NULL) {
-    *batch_ptr = GPU_batch_calloc(1);
+  DRWTempInstancingHandle *handle = BLI_memblock_alloc(idatalist->pool_instancing);
+  if (handle->batch == NULL) {
+    handle->batch = GPU_batch_calloc(1);
   }
 
-  GPUBatch *batch = *batch_ptr;
+  GPUBatch *batch = handle->batch;
   bool instancer_compat = buf ? ((batch->inst[0] == buf) && (buf->vbo_id != 0)) :
-                                ((batch->inst[0] == instancer->inst[0]) &&
-                                 (batch->inst[1] == instancer->inst[1]));
+                                ((batch->inst[0] == instancer->verts[0]) &&
+                                 (batch->inst[1] == instancer->verts[1]));
   bool is_compatible = (batch->prim_type == geom->prim_type) && instancer_compat &&
                        (batch->phase == GPU_BATCH_READY_TO_DRAW) && (batch->elem == geom->elem);
   for (int i = 0; i < GPU_BATCH_VBO_MAX_LEN && is_compatible; i++) {
@@ -161,15 +164,13 @@ GPUBatch *DRW_temp_batch_instance_request(DRWInstanceDataList *idatalist,
   }
 
   if (!is_compatible) {
+    instancing_batch_references_remove(batch);
     GPU_batch_clear(batch);
-    /* Save args and init later */
-    batch->inst[0] = buf;
-    batch->inst[1] = (void *)instancer; /* HACK to save the pointer without other alloc. */
+    /* Save args and init later. */
     batch->phase = GPU_BATCH_READY_TO_BUILD;
-    batch->verts[0] = (void *)geom; /* HACK to save the pointer without other alloc. */
-
-    /* Make sure to free this batch if the instance geom gets free. */
-    GPU_batch_callback_free_set(geom, &instance_batch_free, NULL);
+    handle->buf = buf;
+    handle->instancer = instancer;
+    handle->geom = geom;
   }
   return batch;
 }
@@ -179,7 +180,7 @@ GPUBatch *DRW_temp_batch_request(DRWInstanceDataList *idatalist,
                                  GPUVertBuf *buf,
                                  GPUPrimType prim_type)
 {
-  GPUBatch **batch_ptr = BLI_memblock_alloc(idatalist->pool_instancing);
+  GPUBatch **batch_ptr = BLI_memblock_alloc(idatalist->pool_batching);
   if (*batch_ptr == NULL) {
     *batch_ptr = GPU_batch_calloc(1);
   }
@@ -197,7 +198,13 @@ GPUBatch *DRW_temp_batch_request(DRWInstanceDataList *idatalist,
 static void temp_buffer_handle_free(DRWTempBufferHandle *handle)
 {
   handle->format = NULL;
-  GPU_vertbuf_clear(&handle->buf);
+  GPU_VERTBUF_DISCARD_SAFE(handle->buf);
+}
+
+static void temp_instancing_handle_free(DRWTempInstancingHandle *handle)
+{
+  instancing_batch_references_remove(handle->batch);
+  GPU_BATCH_DISCARD_SAFE(handle->batch);
 }
 
 static void temp_batch_free(GPUBatch **batch)
@@ -215,23 +222,22 @@ void DRW_instance_buffer_finish(DRWInstanceDataList *idatalist)
     if (handle->vert_len != NULL) {
       uint vert_len = *(handle->vert_len);
       uint target_buf_size = ((vert_len / DRW_BUFFER_VERTS_CHUNK) + 1) * DRW_BUFFER_VERTS_CHUNK;
-      if (target_buf_size < handle->buf.vertex_alloc) {
-        GPU_vertbuf_data_resize(&handle->buf, target_buf_size);
+      if (target_buf_size < handle->buf->vertex_alloc) {
+        GPU_vertbuf_data_resize(handle->buf, target_buf_size);
       }
-      GPU_vertbuf_data_len_set(&handle->buf, vert_len);
-      GPU_vertbuf_use(&handle->buf); /* Send data. */
+      GPU_vertbuf_data_len_set(handle->buf, vert_len);
+      GPU_vertbuf_use(handle->buf); /* Send data. */
     }
   }
   /* Finish pending instancing batches. */
-  GPUBatch **batch_ptr;
+  DRWTempInstancingHandle *handle_inst;
   BLI_memblock_iternew(idatalist->pool_instancing, &iter);
-  while ((batch_ptr = BLI_memblock_iterstep(&iter))) {
-    GPUBatch *batch = *batch_ptr;
+  while ((handle_inst = BLI_memblock_iterstep(&iter))) {
+    GPUBatch *batch = handle_inst->batch;
     if (batch && batch->phase == GPU_BATCH_READY_TO_BUILD) {
-      GPUVertBuf *inst_buf = batch->inst[0];
-      /* HACK see DRW_temp_batch_instance_request. */
-      GPUBatch *inst_batch = (void *)batch->inst[1];
-      GPUBatch *geom = (void *)batch->verts[0];
+      GPUVertBuf *inst_buf = handle_inst->buf;
+      GPUBatch *inst_batch = handle_inst->instancer;
+      GPUBatch *geom = handle_inst->geom;
       GPU_batch_copy(batch, geom);
       if (inst_batch != NULL) {
         for (int i = 0; i < GPU_BATCH_INST_VBO_MAX_LEN && inst_batch->verts[i]; i++) {
@@ -241,11 +247,14 @@ void DRW_instance_buffer_finish(DRWInstanceDataList *idatalist)
       else {
         GPU_batch_instbuf_add_ex(batch, inst_buf, false);
       }
+      /* Add reference to avoid comparing pointers (in DRW_temp_batch_request) that could
+       * potentially be the same. This will delay the freeing of the GPUVertBuf itself. */
+      instancing_batch_references_add(batch);
     }
   }
   /* Resize pools and free unused. */
   BLI_memblock_clear(idatalist->pool_buffers, (MemblockValFreeFP)temp_buffer_handle_free);
-  BLI_memblock_clear(idatalist->pool_instancing, (MemblockValFreeFP)temp_batch_free);
+  BLI_memblock_clear(idatalist->pool_instancing, (MemblockValFreeFP)temp_instancing_handle_free);
   BLI_memblock_clear(idatalist->pool_batching, (MemblockValFreeFP)temp_batch_free);
 }
 
@@ -318,7 +327,7 @@ DRWInstanceDataList *DRW_instance_data_list_create(void)
   DRWInstanceDataList *idatalist = MEM_callocN(sizeof(DRWInstanceDataList), "DRWInstanceDataList");
 
   idatalist->pool_batching = BLI_memblock_create(sizeof(GPUBatch *));
-  idatalist->pool_instancing = BLI_memblock_create(sizeof(GPUBatch *));
+  idatalist->pool_instancing = BLI_memblock_create(sizeof(DRWTempInstancingHandle));
   idatalist->pool_buffers = BLI_memblock_create(sizeof(DRWTempBufferHandle));
 
   BLI_addtail(&g_idatalists, idatalist);
@@ -341,7 +350,7 @@ void DRW_instance_data_list_free(DRWInstanceDataList *idatalist)
   }
 
   BLI_memblock_destroy(idatalist->pool_buffers, (MemblockValFreeFP)temp_buffer_handle_free);
-  BLI_memblock_destroy(idatalist->pool_instancing, (MemblockValFreeFP)temp_batch_free);
+  BLI_memblock_destroy(idatalist->pool_instancing, (MemblockValFreeFP)temp_instancing_handle_free);
   BLI_memblock_destroy(idatalist->pool_batching, (MemblockValFreeFP)temp_batch_free);
 
   BLI_remlink(&g_idatalists, idatalist);
diff --git a/source/blender/gpu/GPU_batch.h b/source/blender/gpu/GPU_batch.h
index 7d18fe739d4..e2bf6a87a0f 100644
--- a/source/blender/gpu/GPU_batch.h
+++ b/source/blender/gpu/GPU_batch.h
@@ -89,11 +89,6 @@ typedef struct GPUBatch {
       uint32_t *vao_ids;
     } dynamic_vaos;
   };
-
-  /* XXX This is the only solution if we want to have some data structure using
-   * batches as key to identify nodes. We must destroy these nodes with this callback. */
-  void (*free_callback)(struct GPUBatch *, void *);
-  void *callback_data;
 } GPUBatch;
 
 enum {
@@ -118,8 +113,6 @@ void GPU_batch_discard(GPUBatch *); /* verts & elem are not discarded */
 
 void GPU_batch_vao_cache_clear(GPUBatch *);
 
-void GPU_batch_callback_free_set(GPUBatch *, void (*callback)(GPUBatch *, void *), void *);
-
 void GPU_batch_instbuf_set(GPUBatch *, GPUVertBuf *, bool own_vbo); /* Instancing */
 void GPU_batch_elembuf_set(GPUBatch *batch, GPUIndexBuf *elem, bool own_ibo);
 
diff --git a/source/blender/gpu/intern/gpu_batch.cc b/source/blender/gpu/intern/gpu_batch.cc
index e24b0122c79..fa2c8d728cb 100644
--- a/source/blender/gpu/intern/gpu_batch.cc
+++ b/source/blender/gpu/intern/gpu_batch.cc
@@ -120,7 +120,6 @@ void GPU_batch_init_ex(
   batch->phase = GPU_BATCH_READY_TO_DRAW;
   batch->is_dynamic_vao_count = false;
   batch->owns_flag = owns_flag;
-  batch->free_callback = NULL;
 }
 
 /* This will share the VBOs with the new batch. */
@@ -159,22 +158,10 @@ void GPU_batch_clear(GPUBatch *batch)
 
 void GPU_batch_discard(GPUBatch *batch)
 {
-  if (batch->free_callback) {
-    batch->free_callback(batch, batch->callback_data);
-  }
-
   GPU_batch_clear(batch);
   MEM_freeN(batch);
 }
 
-void GPU_batch_callback_free_set(GPUBatch *batch,
-                                 void (*callback)(GPUBatch *, void *),
-                                 void *user_data)
-{
-  batch->free_callback = callback;
-  batch->callback_data = user_data;
-}
-
 void GPU_batch_instbuf_set(GPUBatch *batch, GPUVertBuf *inst, bool own_vbo)
 {
 #if TRUST_NO_ONE