Volumes: improve tree access token api

There was one functional issue with the previous API which was its
use in `VolumeGrid<T>::grid_for_write(tree_token)`. The issue was
that the tree token had to be received before the grid is accessed.
However, this `grid_for_write` method might create a copy of the
`VolumeGridData` internally and if it does, the passed in `tree_token`
corresponds to the wrong tree.

The solution is to output the token as part of the method. This has two
additional benefits:
* The API is more safe, because one can't pass an r-value into the methods
  anymore. This generally shouldn't be done, because the token should
  live at least as long as the OpenVDB tree is used and shouldn't be freed
  immediatly.
* The API is a bit simpler, because it's not necessary to call the
  `VolumeGrid.tree_access_token()` method anymore.

intern/cycles/blender/volume.cpp

@@ -231,7 +231,6 @@ class BlenderVolumeLoader : public VDBImageLoader {
       if (b_volume_grid.name() == grid_name) {
         const auto *volume_grid = static_cast<const blender::bke::VolumeGridData *>(
             b_volume_grid.ptr.data);
-        tree_access_token = volume_grid->tree_access_token();
         grid = volume_grid->grid_ptr(tree_access_token);
         break;
       }

source/blender/blenkernel/BKE_volume_grid.hh

@@ -139,12 +139,6 @@ class VolumeGridData : public ImplicitSharingMixin {
 
   ~VolumeGridData();
 
-  /**
-   * Get an access token for the underlying tree. This is necessary to be able to detect whether
-   * the grid is currently unused so that it can be safely unloaded.
-   */
-  VolumeTreeAccessToken tree_access_token() const;
-
   /**
    * Create a copy of the volume grid. This should generally only be done when the current grid is
    * shared and one owner wants to modify it.
@@ -158,22 +152,20 @@ class VolumeGridData : public ImplicitSharingMixin {
    * Get the underlying OpenVDB grid for read-only access. This may load the tree lazily if it's
    * not loaded already.
    */
-  const openvdb::GridBase &grid(const VolumeTreeAccessToken &tree_access_token) const;
+  const openvdb::GridBase &grid(VolumeTreeAccessToken &r_token) const;
   /**
    * Get the underlying OpenVDB grid for read and write access. This may load the tree lazily if
    * it's not loaded already. It may also make a copy of the tree if it's currently shared.
    */
-  openvdb::GridBase &grid_for_write(const VolumeTreeAccessToken &tree_access_token);
+  openvdb::GridBase &grid_for_write(VolumeTreeAccessToken &r_token);
 
   /**
    * Same as #grid and #grid_for_write but returns the grid as a `shared_ptr` so that it can be
    * used with APIs that only support grids wrapped into one. This method is not supposed to
    * actually transfer ownership of the grid.
    */
-  std::shared_ptr<const openvdb::GridBase> grid_ptr(
-      const VolumeTreeAccessToken &tree_access_token) const;
-  std::shared_ptr<openvdb::GridBase> grid_ptr_for_write(
-      const VolumeTreeAccessToken &tree_access_token);
+  std::shared_ptr<const openvdb::GridBase> grid_ptr(VolumeTreeAccessToken &r_token) const;
+  std::shared_ptr<openvdb::GridBase> grid_ptr_for_write(VolumeTreeAccessToken &r_token);
 
   /**
    * Get the name of the grid that's stored in the grid meta-data.
@@ -319,8 +311,8 @@ template<typename T> class VolumeGrid : public GVolumeGrid {
   /**
    * Wraps the same methods on #VolumeGridData but casts to the correct OpenVDB type.
    */
-  const OpenvdbGridType<T> &grid(const VolumeTreeAccessToken &tree_access_token) const;
-  OpenvdbGridType<T> &grid_for_write(const VolumeTreeAccessToken &tree_access_token);
+  const OpenvdbGridType<T> &grid(VolumeTreeAccessToken &r_token) const;
+  OpenvdbGridType<T> &grid_for_write(VolumeTreeAccessToken &r_token);
 
  private:
   void assert_correct_type() const;
@@ -381,18 +373,15 @@ inline VolumeGrid<T>::VolumeGrid(std::shared_ptr<OpenvdbGridType<T>> grid)
 }
 
 template<typename T>
-inline const OpenvdbGridType<T> &VolumeGrid<T>::grid(
-    const VolumeTreeAccessToken &tree_access_token) const
+inline const OpenvdbGridType<T> &VolumeGrid<T>::grid(VolumeTreeAccessToken &r_token) const
 {
-  return static_cast<const OpenvdbGridType<T> &>(data_->grid(tree_access_token));
+  return static_cast<const OpenvdbGridType<T> &>(data_->grid(r_token));
 }
 
 template<typename T>
-inline OpenvdbGridType<T> &VolumeGrid<T>::grid_for_write(
-    const VolumeTreeAccessToken &tree_access_token)
+inline OpenvdbGridType<T> &VolumeGrid<T>::grid_for_write(VolumeTreeAccessToken &r_token)
 {
-  return static_cast<OpenvdbGridType<T> &>(
-      this->get_for_write().grid_for_write(tree_access_token));
+  return static_cast<OpenvdbGridType<T> &>(this->get_for_write().grid_for_write(r_token));
 }
 
 template<typename T> inline void VolumeGrid<T>::assert_correct_type() const

source/blender/blenkernel/intern/volume.cc

@@ -556,11 +556,11 @@ bool BKE_volume_save(const Volume *volume,
   openvdb::GridCPtrVec vdb_grids;
 
   /* Tree users need to be kept alive for as long as the grids may be accessed. */
-  blender::Vector<blender::bke::VolumeTreeAccessToken> tree_users;
+  blender::Vector<blender::bke::VolumeTreeAccessToken> tree_tokens;
 
   for (const GVolumeGrid &grid : grids) {
-    tree_users.append(grid->tree_access_token());
-    vdb_grids.push_back(grid->grid_ptr(tree_users.last()));
+    tree_tokens.append_as();
+    vdb_grids.push_back(grid->grid_ptr(tree_tokens.last()));
   }
 
   try {
@@ -594,9 +594,9 @@ std::optional<blender::Bounds<blender::float3>> BKE_volume_min_max(const Volume
     std::optional<blender::Bounds<blender::float3>> result;
     for (const int i : IndexRange(BKE_volume_num_grids(volume))) {
       const blender::bke::VolumeGridData *volume_grid = BKE_volume_grid_get(volume, i);
-      blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
+      blender::bke::VolumeTreeAccessToken tree_token;
       result = blender::bounds::merge(result,
-                                      BKE_volume_grid_bounds(volume_grid->grid_ptr(access_token)));
+                                      BKE_volume_grid_bounds(volume_grid->grid_ptr(tree_token)));
     }
     return result;
   }

source/blender/blenkernel/intern/volume_grid.cc

@@ -93,41 +93,32 @@ void VolumeGridData::delete_self()
   MEM_delete(this);
 }
 
-VolumeTreeAccessToken VolumeGridData::tree_access_token() const
+const openvdb::GridBase &VolumeGridData::grid(VolumeTreeAccessToken &r_token) const
 {
-  VolumeTreeAccessToken user;
-  user.token_ = tree_access_token_;
-  return user;
+  return *this->grid_ptr(r_token);
 }
 
-const openvdb::GridBase &VolumeGridData::grid(const VolumeTreeAccessToken &access_token) const
+openvdb::GridBase &VolumeGridData::grid_for_write(VolumeTreeAccessToken &r_token)
 {
-  return *this->grid_ptr(access_token);
-}
-
-openvdb::GridBase &VolumeGridData::grid_for_write(const VolumeTreeAccessToken &access_token)
-{
-  return *this->grid_ptr_for_write(access_token);
+  return *this->grid_ptr_for_write(r_token);
 }
 
 std::shared_ptr<const openvdb::GridBase> VolumeGridData::grid_ptr(
-    const VolumeTreeAccessToken &access_token) const
+    VolumeTreeAccessToken &r_token) const
 {
-  BLI_assert(access_token.valid_for(*this));
-  UNUSED_VARS_NDEBUG(access_token);
   std::lock_guard lock{mutex_};
   this->ensure_grid_loaded();
+  r_token.token_ = tree_access_token_;
   return grid_;
 }
 
 std::shared_ptr<openvdb::GridBase> VolumeGridData::grid_ptr_for_write(
-    const VolumeTreeAccessToken &access_token)
+    VolumeTreeAccessToken &r_token)
 {
-  BLI_assert(access_token.valid_for(*this));
-  UNUSED_VARS_NDEBUG(access_token);
   BLI_assert(this->is_mutable());
   std::lock_guard lock{mutex_};
   this->ensure_grid_loaded();
+  r_token.token_ = tree_access_token_;
   if (tree_sharing_info_->is_mutable()) {
     tree_sharing_info_->tag_ensured_mutable();
   }
@@ -468,8 +459,8 @@ void set_transform_matrix(VolumeGridData &grid, const float4x4 &matrix)
 void clear_tree(VolumeGridData &grid)
 {
 #ifdef WITH_OPENVDB
-  VolumeTreeAccessToken access_token = grid.tree_access_token();
-  grid.grid_for_write(access_token).clear();
+  VolumeTreeAccessToken tree_token;
+  grid.grid_for_write(tree_token).clear();
 #else
   UNUSED_VARS(grid);
 #endif
@@ -488,9 +479,9 @@ bool is_loaded(const VolumeGridData &grid)
 void load(const VolumeGridData &grid)
 {
 #ifdef WITH_OPENVDB
-  VolumeTreeAccessToken access_token = grid.tree_access_token();
+  VolumeTreeAccessToken tree_token;
   /* Just "touch" the grid, so that it is loaded. */
-  grid.grid(access_token);
+  grid.grid(tree_token);
 #else
   UNUSED_VARS(grid);
 #endif

source/blender/blenkernel/intern/volume_grid_file_cache.cc

@@ -162,9 +162,9 @@ static GVolumeGrid get_cached_grid(const StringRef file_path,
     const GVolumeGrid main_grid = get_grid_from_file(file_path, grid_name, 0);
     const VolumeGridType grid_type = main_grid->grid_type();
     const float resolution_factor = 1.0f / (1 << simplify_level);
-    const VolumeTreeAccessToken access_token = main_grid->tree_access_token();
+    VolumeTreeAccessToken tree_token;
     return BKE_volume_grid_create_with_changed_resolution(
-        grid_type, main_grid->grid(access_token), resolution_factor);
+        grid_type, main_grid->grid(tree_token), resolution_factor);
   };
   /* This allows the returned grid to already contain meta-data and transforms, even if the tree is
    * not loaded yet. */

source/blender/blenkernel/intern/volume_render.cc

@@ -98,8 +98,8 @@ bool BKE_volume_grid_dense_floats(const Volume *volume,
 {
 #ifdef WITH_OPENVDB
   const VolumeGridType grid_type = volume_grid->grid_type();
-  blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  const openvdb::GridBase &grid = volume_grid->grid(access_token);
+  blender::bke::VolumeTreeAccessToken tree_token;
+  const openvdb::GridBase &grid = volume_grid->grid(tree_token);
 
   const openvdb::CoordBBox bbox = grid.evalActiveVoxelBoundingBox();
   if (bbox.empty()) {
@@ -333,8 +333,8 @@ void BKE_volume_grid_wireframe(const Volume *volume,
   }
 
 #ifdef WITH_OPENVDB
-  blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  const openvdb::GridBase &grid = volume_grid->grid(access_token);
+  blender::bke::VolumeTreeAccessToken tree_token;
+  const openvdb::GridBase &grid = volume_grid->grid(tree_token);
 
   if (volume->display.wireframe_type == VOLUME_WIREFRAME_BOUNDS) {
     /* Bounding box. */
@@ -412,8 +412,8 @@ void BKE_volume_grid_selection_surface(const Volume * /*volume*/,
                                        void *cb_userdata)
 {
 #ifdef WITH_OPENVDB
-  blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  const openvdb::GridBase &grid = volume_grid->grid(access_token);
+  blender::bke::VolumeTreeAccessToken tree_token;
+  const openvdb::GridBase &grid = volume_grid->grid(tree_token);
   blender::Vector<openvdb::CoordBBox> boxes = get_bounding_boxes(
       volume_grid->grid_type(), grid, true);
 

source/blender/blenkernel/intern/volume_test.cc

@@ -85,24 +85,23 @@ TEST_F(VolumeTest, lazy_load_grid)
   VolumeGrid<float> volume_grid{MEM_new<VolumeGridData>(__func__, load_grid)};
   EXPECT_EQ(load_counter, 0);
   EXPECT_FALSE(volume_grid->is_loaded());
-  VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  EXPECT_EQ(volume_grid.grid(access_token).background(), 10.0f);
+  VolumeTreeAccessToken tree_token;
+  EXPECT_EQ(volume_grid.grid(tree_token).background(), 10.0f);
   EXPECT_EQ(load_counter, 1);
   EXPECT_TRUE(volume_grid->is_loaded());
   EXPECT_TRUE(volume_grid->is_reloadable());
-  EXPECT_EQ(volume_grid.grid(access_token).background(), 10.0f);
+  EXPECT_EQ(volume_grid.grid(tree_token).background(), 10.0f);
   EXPECT_EQ(load_counter, 1);
   volume_grid->unload_tree_if_possible();
   EXPECT_TRUE(volume_grid->is_loaded());
-  access_token.reset();
+  tree_token.reset();
   volume_grid->unload_tree_if_possible();
   EXPECT_FALSE(volume_grid->is_loaded());
-  access_token = volume_grid->tree_access_token();
-  EXPECT_EQ(volume_grid.grid(access_token).background(), 10.0f);
+  EXPECT_EQ(volume_grid.grid(tree_token).background(), 10.0f);
   EXPECT_TRUE(volume_grid->is_loaded());
   EXPECT_EQ(load_counter, 2);
-  volume_grid.grid_for_write(access_token).getAccessor().setValue({0, 0, 0}, 1.0f);
-  EXPECT_EQ(volume_grid.grid(access_token).getAccessor().getValue({0, 0, 0}), 1.0f);
+  volume_grid.grid_for_write(tree_token).getAccessor().setValue({0, 0, 0}, 1.0f);
+  EXPECT_EQ(volume_grid.grid(tree_token).getAccessor().getValue({0, 0, 0}), 1.0f);
   EXPECT_FALSE(volume_grid->is_reloadable());
 }
 
@@ -121,11 +120,11 @@ TEST_F(VolumeTest, lazy_load_tree_only)
   volume_grid.get_for_write().set_name("Test");
   EXPECT_FALSE(load_run);
   EXPECT_EQ(volume_grid->name(), "Test");
-  VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  volume_grid.grid_for_write(access_token);
+  VolumeTreeAccessToken tree_token;
+  volume_grid.grid_for_write(tree_token);
   EXPECT_TRUE(load_run);
   EXPECT_EQ(volume_grid->name(), "Test");
-  EXPECT_EQ(volume_grid.grid(access_token).background(), 10.0f);
+  EXPECT_EQ(volume_grid.grid(tree_token).background(), 10.0f);
 }
 
 }  // namespace blender::bke::tests

source/blender/modifiers/intern/MOD_volume_displace.cc

@@ -289,8 +289,8 @@ static void displace_volume(ModifierData *md, const ModifierEvalContext *ctx, Vo
     blender::bke::VolumeGridData *volume_grid = BKE_volume_grid_get_for_write(volume, grid_index);
     BLI_assert(volume_grid);
 
-    blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-    openvdb::GridBase &grid = volume_grid->grid_for_write(access_token);
+    blender::bke::VolumeTreeAccessToken tree_token;
+    openvdb::GridBase &grid = volume_grid->grid_for_write(tree_token);
     VolumeGridType grid_type = volume_grid->grid_type();
 
     DisplaceGridOp displace_grid_op{grid, *vdmd, *ctx};

source/blender/modifiers/intern/MOD_volume_to_mesh.cc

@@ -158,8 +158,8 @@ static Mesh *modify_mesh(ModifierData *md, const ModifierEvalContext *ctx, Mesh
     return create_empty_mesh(input_mesh);
   }
 
-  const blender::bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-  const openvdb::GridBase &local_grid = volume_grid->grid(access_token);
+  blender::bke::VolumeTreeAccessToken tree_token;
+  const openvdb::GridBase &local_grid = volume_grid->grid(tree_token);
 
   openvdb::math::Transform::Ptr transform = local_grid.transform().copy();
   transform->postMult(openvdb::Mat4d((float *)vmmd->object->object_to_world));

source/blender/nodes/geometry/nodes/node_geo_distribute_points_in_volume.cc

@@ -225,8 +225,8 @@ static void node_geo_exec(GeoNodeExecParams params)
         continue;
       }
 
-      bke::VolumeTreeAccessToken access_token = volume_grid->tree_access_token();
-      const openvdb::GridBase &base_grid = volume_grid->grid(access_token);
+      bke::VolumeTreeAccessToken tree_token;
+      const openvdb::GridBase &base_grid = volume_grid->grid(tree_token);
 
       if (!base_grid.isType<openvdb::FloatGrid>()) {
         continue;

source/blender/nodes/geometry/nodes/node_geo_volume_to_mesh.cc

@@ -180,12 +180,12 @@ static Mesh *create_mesh_from_volume(GeometrySet &geometry_set, GeoNodeExecParam
   const Main *bmain = DEG_get_bmain(params.depsgraph());
   BKE_volume_load(volume, bmain);
 
-  Vector<bke::VolumeTreeAccessToken> access_tokens;
+  Vector<bke::VolumeTreeAccessToken> tree_tokens;
   Vector<const openvdb::GridBase *> grids;
   for (const int i : IndexRange(BKE_volume_num_grids(volume))) {
     const bke::VolumeGridData *volume_grid = BKE_volume_grid_get(volume, i);
-    access_tokens.append(volume_grid->tree_access_token());
-    grids.append(&volume_grid->grid(access_tokens.last()));
+    tree_tokens.append_as();
+    grids.append(&volume_grid->grid(tree_tokens.last()));
   }
 
   if (grids.is_empty()) {