Cleanup: use template utility function to handle OpenVDB grid types in Cycles

2 files changed, 109 insertions, 120 deletions

diff --git a/intern/cycles/render/image_vdb.cpp b/intern/cycles/render/image_vdb.cpp
index fb6394e8917..13cdda552ba 100644
--- a/intern/cycles/render/image_vdb.cpp
+++ b/intern/cycles/render/image_vdb.cpp
@@ -16,8 +16,9 @@
 
 #include "render/image_vdb.h"
 
+#include "util/util_openvdb.h"
+
 #ifdef WITH_OPENVDB
-#  include <openvdb/openvdb.h>
 #  include <openvdb/tools/Dense.h>
 #endif
 #ifdef WITH_NANOVDB
@@ -26,6 +27,51 @@
 
 CCL_NAMESPACE_BEGIN
 
+#ifdef WITH_OPENVDB
+struct NumChannelsOp {
+  int num_channels = 0;
+
+  template<typename GridType, typename FloatGridType, typename FloatDataType, int channels>
+  bool operator()(const openvdb::GridBase::ConstPtr &grid)
+  {
+    num_channels = channels;
+    return true;
+  }
+};
+
+struct ToDenseOp {
+  openvdb::CoordBBox bbox;
+  void *pixels;
+
+  template<typename GridType, typename FloatGridType, typename FloatDataType, int channels>
+  bool operator()(const openvdb::GridBase::ConstPtr &grid)
+  {
+    openvdb::tools::Dense<FloatDataType, openvdb::tools::LayoutXYZ> dense(bbox,
+                                                                          (FloatDataType *)pixels);
+    openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<GridType>(grid), dense);
+    return true;
+  }
+};
+
+#  ifdef WITH_NANOVDB
+struct ToNanoOp {
+  nanovdb::GridHandle<> nanogrid;
+
+  template<typename GridType, typename FloatGridType, typename FloatDataType, int channels>
+  bool operator()(const openvdb::GridBase::ConstPtr &grid)
+  {
+    if constexpr (!std::is_same_v<GridType, openvdb::MaskGrid>) {
+      nanogrid = nanovdb::openToNanoVDB(FloatGridType(*openvdb::gridConstPtrCast<GridType>(grid)));
+      return true;
+    }
+    else {
+      return false;
+    }
+  }
+};
+#  endif
+#endif
+
 VDBImageLoader::VDBImageLoader(const string &grid_name) : grid_name(grid_name)
 {
 }
@@ -41,98 +87,40 @@ bool VDBImageLoader::load_metadata(const ImageDeviceFeatures &features, ImageMet
     return false;
   }
 
-  bbox = grid->evalActiveVoxelBoundingBox();
-  if (bbox.empty()) {
+  /* Get number of channels from type. */
+  NumChannelsOp op;
+  if (!openvdb::grid_type_operation(grid, op)) {
     return false;
   }
 
-  /* Set dimensions. */
-  openvdb::Coord dim = bbox.dim();
-  metadata.width = dim.x();
-  metadata.height = dim.y();
-  metadata.depth = dim.z();
+  metadata.channels = op.num_channels;
 
   /* Set data type. */
-  if (grid->isType<openvdb::FloatGrid>()) {
-    metadata.channels = 1;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(*openvdb::gridConstPtrCast<openvdb::FloatGrid>(grid));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::Vec3fGrid>()) {
-    metadata.channels = 3;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(*openvdb::gridConstPtrCast<openvdb::Vec3fGrid>(grid));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::BoolGrid>()) {
-    metadata.channels = 1;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::FloatGrid(*openvdb::gridConstPtrCast<openvdb::BoolGrid>(grid)));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::DoubleGrid>()) {
-    metadata.channels = 1;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::FloatGrid(*openvdb::gridConstPtrCast<openvdb::DoubleGrid>(grid)));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::Int32Grid>()) {
-    metadata.channels = 1;
 #  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::FloatGrid(*openvdb::gridConstPtrCast<openvdb::Int32Grid>(grid)));
+  if (features.has_nanovdb) {
+    /* NanoVDB expects no inactive leaf nodes. */
+    /*openvdb::FloatGrid &pruned_grid = *openvdb::gridPtrCast<openvdb::FloatGrid>(grid);
+    openvdb::tools::pruneInactive(pruned_grid.tree());
+    nanogrid = nanovdb::openToNanoVDB(pruned_grid);*/
+    ToNanoOp op;
+    if (!openvdb::grid_type_operation(grid, op)) {
+      return false;
     }
-#  endif
-  }
-  else if (grid->isType<openvdb::Int64Grid>()) {
-    metadata.channels = 1;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::FloatGrid(*openvdb::gridConstPtrCast<openvdb::Int64Grid>(grid)));
-    }
-#  endif
+    nanogrid = std::move(op.nanogrid);
   }
-  else if (grid->isType<openvdb::Vec3IGrid>()) {
-    metadata.channels = 3;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::Vec3fGrid(*openvdb::gridConstPtrCast<openvdb::Vec3IGrid>(grid)));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::Vec3dGrid>()) {
-    metadata.channels = 3;
-#  ifdef WITH_NANOVDB
-    if (features.has_nanovdb) {
-      nanogrid = nanovdb::openToNanoVDB(
-          openvdb::Vec3fGrid(*openvdb::gridConstPtrCast<openvdb::Vec3dGrid>(grid)));
-    }
-#  endif
-  }
-  else if (grid->isType<openvdb::MaskGrid>()) {
-    metadata.channels = 1;
-#  ifdef WITH_NANOVDB
-    return false;  // Unsupported
 #  endif
-  }
-  else {
+
+  /* Set dimensions. */
+  bbox = grid->evalActiveVoxelBoundingBox();
+  if (bbox.empty()) {
     return false;
   }
 
+  openvdb::Coord dim = bbox.dim();
+  metadata.width = dim.x();
+  metadata.height = dim.y();
+  metadata.depth = dim.z();
+
 #  ifdef WITH_NANOVDB
   if (nanogrid) {
     metadata.byte_size = nanogrid.size();
@@ -200,45 +188,10 @@ bool VDBImageLoader::load_pixels(const ImageMetaData &, void *pixels, const size
   else
 #  endif
   {
-    if (grid->isType<openvdb::FloatGrid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::FloatGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::Vec3fGrid>()) {
-      openvdb::tools::Dense<openvdb::Vec3f, openvdb::tools::LayoutXYZ> dense(
-          bbox, (openvdb::Vec3f *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::Vec3fGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::BoolGrid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::BoolGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::DoubleGrid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::DoubleGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::Int32Grid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::Int32Grid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::Int64Grid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::Int64Grid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::Vec3IGrid>()) {
-      openvdb::tools::Dense<openvdb::Vec3f, openvdb::tools::LayoutXYZ> dense(
-          bbox, (openvdb::Vec3f *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::Vec3IGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::Vec3dGrid>()) {
-      openvdb::tools::Dense<openvdb::Vec3f, openvdb::tools::LayoutXYZ> dense(
-          bbox, (openvdb::Vec3f *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::Vec3dGrid>(grid), dense);
-    }
-    else if (grid->isType<openvdb::MaskGrid>()) {
-      openvdb::tools::Dense<float, openvdb::tools::LayoutXYZ> dense(bbox, (float *)pixels);
-      openvdb::tools::copyToDense(*openvdb::gridConstPtrCast<openvdb::MaskGrid>(grid), dense);
-    }
+    ToDenseOp op;
+    op.pixels = pixels;
+    op.bbox = bbox;
+    openvdb::grid_type_operation(grid, op);
   }
   return true;
 #else
diff --git a/intern/cycles/util/util_openvdb.h b/intern/cycles/util/util_openvdb.h
index a3ebb03e5a4..ae5326e3199 100644
--- a/intern/cycles/util/util_openvdb.h
+++ b/intern/cycles/util/util_openvdb.h
@@ -25,6 +25,42 @@ namespace openvdb {
 using Vec4fTree = tree::Tree4<Vec4f, 5, 4, 3>::Type;
 using Vec4fGrid = Grid<Vec4fTree>;
 
+/* Apply operation to known grid types. */
+template<typename OpType>
+bool grid_type_operation(const openvdb::GridBase::ConstPtr &grid, OpType &&op)
+{
+  if (grid->isType<openvdb::FloatGrid>()) {
+    return op.template operator()<openvdb::FloatGrid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else if (grid->isType<openvdb::Vec3fGrid>()) {
+    return op.template operator()<openvdb::Vec3fGrid, openvdb::Vec3fGrid, openvdb::Vec3f, 3>(grid);
+  }
+  else if (grid->isType<openvdb::BoolGrid>()) {
+    return op.template operator()<openvdb::BoolGrid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else if (grid->isType<openvdb::DoubleGrid>()) {
+    return op.template operator()<openvdb::DoubleGrid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else if (grid->isType<openvdb::Int32Grid>()) {
+    return op.template operator()<openvdb::Int32Grid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else if (grid->isType<openvdb::Int64Grid>()) {
+    return op.template operator()<openvdb::Int64Grid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else if (grid->isType<openvdb::Vec3IGrid>()) {
+    return op.template operator()<openvdb::Vec3IGrid, openvdb::Vec3fGrid, openvdb::Vec3f, 3>(grid);
+  }
+  else if (grid->isType<openvdb::Vec3dGrid>()) {
+    return op.template operator()<openvdb::Vec3dGrid, openvdb::Vec3fGrid, openvdb::Vec3f, 3>(grid);
+  }
+  else if (grid->isType<openvdb::MaskGrid>()) {
+    return op.template operator()<openvdb::MaskGrid, openvdb::FloatGrid, float, 1>(grid);
+  }
+  else {
+    return false;
+  }
+}
+
 };  // namespace openvdb
 
 #endif