Sculpt: Refactor for pose brush factors creation

Part of #118145.

source/blender/editors/sculpt_paint/sculpt_pose.cc

@@ -315,40 +315,166 @@ static void calc_bmesh(const Sculpt &sd,
 struct PoseGrowFactorData {
   float3 pos_avg;
   int pos_count;
+  static PoseGrowFactorData join(const PoseGrowFactorData &a, const PoseGrowFactorData &b)
+  {
+    PoseGrowFactorData joined;
+    joined.pos_avg = a.pos_avg + b.pos_avg;
+    joined.pos_count = a.pos_count + b.pos_count;
+    return joined;
+  }
 };
 
-static void pose_brush_grow_factor_task(const Object &ob,
-                                        const float3 &pose_initial_co,
-                                        const Span<float> prev_mask,
-                                        bke::pbvh::Node &node,
-                                        MutableSpan<float> pose_factor,
-                                        PoseGrowFactorData &gftd)
+struct GrowFactorLocalData {
+  Vector<int> vert_indices;
+  Vector<Vector<int>> vert_neighbors;
+};
+
+BLI_NOINLINE static void add_fake_neighbors(const Span<int> fake_neighbors,
+                                            const Span<int> verts,
+                                            const MutableSpan<Vector<int>> neighbors)
 {
-  const SculptSession &ss = *ob.sculpt;
-  const char symm = SCULPT_mesh_symmetry_xyz_get(ob);
-  PBVHVertexIter vd;
-  BKE_pbvh_vertex_iter_begin (*ss.pbvh, &node, vd, PBVH_ITER_UNIQUE) {
-    SculptVertexNeighborIter ni;
-    float max = 0.0f;
-
-    /* Grow the factor. */
-    SCULPT_VERTEX_NEIGHBORS_ITER_BEGIN (ss, vd.vertex, ni) {
-      float vmask_f = prev_mask[ni.index];
-      max = std::max(vmask_f, max);
+  for (const int i : verts.index_range()) {
+    if (fake_neighbors[verts[i]] != FAKE_NEIGHBOR_NONE) {
+      neighbors[i].append(fake_neighbors[verts[i]]);
     }
-    SCULPT_VERTEX_NEIGHBORS_ITER_END(ni);
+  }
+}
 
-    /* Keep the count of the vertices that where added to the factors in this grow iteration. */
-    if (max > prev_mask[vd.index]) {
-      pose_factor[vd.index] = max;
-      if (SCULPT_check_vertex_pivot_symmetry(vd.co, pose_initial_co, symm)) {
-        add_v3_v3(gftd.pos_avg, vd.co);
+static void grow_factors_mesh(const ePaintSymmetryFlags symm,
+                              const float3 &pose_initial_position,
+                              const Span<float3> vert_positions,
+                              const OffsetIndices<int> faces,
+                              const Span<int> corner_verts,
+                              const GroupedSpan<int> vert_to_face_map,
+                              const Span<bool> hide_vert,
+                              const Span<bool> hide_poly,
+                              const Span<int> fake_neighbors,
+                              const Span<float> prev_mask,
+                              const bke::pbvh::Node &node,
+                              GrowFactorLocalData &tls,
+                              const MutableSpan<float> pose_factor,
+                              PoseGrowFactorData &gftd)
+{
+  const Span<int> verts = hide::node_visible_verts(node, hide_vert, tls.vert_indices);
+
+  tls.vert_neighbors.resize(verts.size());
+  const MutableSpan<Vector<int>> neighbors = tls.vert_neighbors;
+  calc_vert_neighbors(faces, corner_verts, vert_to_face_map, hide_poly, verts, neighbors);
+
+  if (!fake_neighbors.is_empty()) {
+    add_fake_neighbors(fake_neighbors, verts, neighbors);
+  }
+
+  for (const int i : verts.index_range()) {
+    const int vert = verts[i];
+
+    float max = 0.0f;
+    for (const int neighbor : neighbors[i]) {
+      max = std::max(max, prev_mask[neighbor]);
+    }
+
+    if (max > prev_mask[vert]) {
+      const float3 &position = vert_positions[verts[i]];
+      pose_factor[vert] = max;
+      if (SCULPT_check_vertex_pivot_symmetry(position, pose_initial_position, symm)) {
+        gftd.pos_avg += position;
         gftd.pos_count++;
       }
     }
   }
+}
 
-  BKE_pbvh_vertex_iter_end;
+static void grow_factors_grids(const ePaintSymmetryFlags symm,
+                               const float3 &pose_initial_position,
+                               const SubdivCCG &subdiv_ccg,
+                               const Span<int> fake_neighbors,
+                               const Span<float> prev_mask,
+                               const bke::pbvh::Node &node,
+                               const MutableSpan<float> pose_factor,
+                               PoseGrowFactorData &gftd)
+{
+  const CCGKey key = BKE_subdiv_ccg_key_top_level(subdiv_ccg);
+  const Span<CCGElem *> elems = subdiv_ccg.grids;
+  const BitGroupVector<> &grid_hidden = subdiv_ccg.grid_hidden;
+  const Span<int> grids = bke::pbvh::node_grid_indices(node);
+
+  for (const int i : grids.index_range()) {
+    const int grid = grids[i];
+    CCGElem *elem = elems[grid];
+    const int start = key.grid_area * grid;
+    for (const short y : IndexRange(key.grid_size)) {
+      for (const short x : IndexRange(key.grid_size)) {
+        const int offset = CCG_grid_xy_to_index(key.grid_size, x, y);
+        if (!grid_hidden.is_empty() && grid_hidden[grid][offset]) {
+          continue;
+        }
+        const int vert = start + offset;
+
+        SubdivCCGNeighbors neighbors;
+        BKE_subdiv_ccg_neighbor_coords_get(
+            subdiv_ccg, SubdivCCGCoord{grids[i], x, y}, false, neighbors);
+
+        float max = 0.0f;
+        for (const SubdivCCGCoord neighbor : neighbors.coords) {
+          max = std::max(max, prev_mask[neighbor.to_index(key)]);
+        }
+        if (!fake_neighbors.is_empty()) {
+          if (fake_neighbors[vert] != FAKE_NEIGHBOR_NONE) {
+            max = std::max(max, prev_mask[fake_neighbors[vert]]);
+          }
+        }
+
+        if (max > prev_mask[vert]) {
+          const float3 &position = CCG_elem_offset_co(key, elem, offset);
+          pose_factor[vert] = max;
+          if (SCULPT_check_vertex_pivot_symmetry(position, pose_initial_position, symm)) {
+            gftd.pos_avg += position;
+            gftd.pos_count++;
+          }
+        }
+      }
+    }
+  }
+}
+
+static void grow_factors_bmesh(const ePaintSymmetryFlags symm,
+                               const float3 &pose_initial_position,
+                               const Span<int> fake_neighbors,
+                               const Span<float> prev_mask,
+                               bke::pbvh::Node &node,
+                               const MutableSpan<float> pose_factor,
+                               PoseGrowFactorData &gftd)
+{
+  const Set<BMVert *, 0> &verts = BKE_pbvh_bmesh_node_unique_verts(&node);
+
+  Vector<BMVert *, 64> neighbors;
+
+  int i = 0;
+  for (BMVert *bm_vert : verts) {
+    const int vert = BM_elem_index_get(bm_vert);
+
+    float max = 0.0f;
+    neighbors.clear();
+    for (const BMVert *neighbor : vert_neighbors_get_bmesh(*bm_vert, neighbors)) {
+      max = std::max(max, prev_mask[BM_elem_index_get(neighbor)]);
+    }
+    if (!fake_neighbors.is_empty()) {
+      if (fake_neighbors[vert] != FAKE_NEIGHBOR_NONE) {
+        max = std::max(max, prev_mask[fake_neighbors[vert]]);
+      }
+    }
+
+    if (max > prev_mask[vert]) {
+      const float3 &position = bm_vert->co;
+      pose_factor[vert] = max;
+      if (SCULPT_check_vertex_pivot_symmetry(position, pose_initial_position, symm)) {
+        gftd.pos_avg += position;
+        gftd.pos_count++;
+      }
+    }
+
+    i++;
+  }
 }
 
 /* Grow the factor until its boundary is near to the offset pose origin or outside the target
@@ -362,8 +488,10 @@ static void sculpt_pose_grow_pose_factor(Object &ob,
                                          MutableSpan<float> pose_factor)
 {
   bke::pbvh::Tree &pbvh = *ob.sculpt->pbvh;
+  const ePaintSymmetryFlags symm = SCULPT_mesh_symmetry_xyz_get(ob);
 
   Vector<bke::pbvh::Node *> nodes = bke::pbvh::search_gather(pbvh, {});
+  const Span<int> fake_neighbors = ss.fake_neighbors.fake_neighbor_index;
 
   bool grow_next_iteration = true;
   float prev_len = FLT_MAX;
@@ -371,22 +499,85 @@ static void sculpt_pose_grow_pose_factor(Object &ob,
   while (grow_next_iteration) {
     prev_mask.as_mutable_span().copy_from(pose_factor);
 
-    PoseGrowFactorData gftd = threading::parallel_reduce(
-        nodes.index_range(),
-        1,
-        PoseGrowFactorData{},
-        [&](const IndexRange range, PoseGrowFactorData gftd) {
-          for (const int i : range) {
-            pose_brush_grow_factor_task(ob, pose_target, prev_mask, *nodes[i], pose_factor, gftd);
-          }
-          return gftd;
-        },
-        [](const PoseGrowFactorData &a, const PoseGrowFactorData &b) {
-          PoseGrowFactorData joined;
-          joined.pos_avg = a.pos_avg + b.pos_avg;
-          joined.pos_count = a.pos_count + b.pos_count;
-          return joined;
-        });
+    PoseGrowFactorData gftd;
+    threading::EnumerableThreadSpecific<GrowFactorLocalData> all_tls;
+    switch (pbvh.type()) {
+      case bke::pbvh::Type::Mesh: {
+        const Mesh &mesh = *static_cast<const Mesh *>(ob.data);
+        const Span<float3> vert_positions = BKE_pbvh_get_vert_positions(pbvh);
+        const OffsetIndices faces = mesh.faces();
+        const Span<int> corner_verts = mesh.corner_verts();
+        const GroupedSpan<int> vert_to_face_map = mesh.vert_to_face_map();
+        const bke::AttributeAccessor attributes = mesh.attributes();
+        const VArraySpan<bool> hide_vert = *attributes.lookup<bool>(".hide_vert",
+                                                                    bke::AttrDomain::Point);
+        const VArraySpan<bool> hide_poly = *attributes.lookup<bool>(".hide_poly",
+                                                                    bke::AttrDomain::Face);
+        gftd = threading::parallel_reduce(
+            nodes.index_range(),
+            1,
+            PoseGrowFactorData{},
+            [&](const IndexRange range, PoseGrowFactorData gftd) {
+              GrowFactorLocalData &tls = all_tls.local();
+              for (const int i : range) {
+                grow_factors_mesh(symm,
+                                  pose_target,
+                                  vert_positions,
+                                  faces,
+                                  corner_verts,
+                                  vert_to_face_map,
+                                  hide_vert,
+                                  hide_poly,
+                                  fake_neighbors,
+                                  prev_mask,
+                                  *nodes[i],
+                                  tls,
+                                  pose_factor,
+                                  gftd);
+              }
+              return gftd;
+            },
+            PoseGrowFactorData::join);
+        break;
+      }
+      case bke::pbvh::Type::Grids: {
+        const SubdivCCG &subdiv_ccg = *ss.subdiv_ccg;
+        gftd = threading::parallel_reduce(
+            nodes.index_range(),
+            1,
+            PoseGrowFactorData{},
+            [&](const IndexRange range, PoseGrowFactorData gftd) {
+              for (const int i : range) {
+                grow_factors_grids(symm,
+                                   pose_target,
+                                   subdiv_ccg,
+                                   fake_neighbors,
+                                   prev_mask,
+                                   *nodes[i],
+                                   pose_factor,
+                                   gftd);
+              }
+              return gftd;
+            },
+            PoseGrowFactorData::join);
+        break;
+      }
+      case bke::pbvh::Type::BMesh: {
+        gftd = threading::parallel_reduce(
+            nodes.index_range(),
+            1,
+            PoseGrowFactorData{},
+            [&](const IndexRange range, PoseGrowFactorData gftd) {
+              for (const int i : range) {
+                grow_factors_bmesh(
+                    symm, pose_target, fake_neighbors, prev_mask, *nodes[i], pose_factor, gftd);
+              }
+              return gftd;
+            },
+            PoseGrowFactorData::join);
+        break;
+      }
+    }
 
     if (gftd.pos_count != 0) {
       gftd.pos_avg /= float(gftd.pos_count);
@@ -654,26 +845,110 @@ void calc_pose_data(Object &ob,
   }
 }
 
-static void pose_brush_init_task(SculptSession &ss,
-                                 MutableSpan<float> pose_factor,
-                                 bke::pbvh::Node *node)
+static void smooth_geometry_data_array(const Object &object,
+                                       const int iterations,
+                                       const MutableSpan<float> data)
 {
-  PBVHVertexIter vd;
-  BKE_pbvh_vertex_iter_begin (*ss.pbvh, node, vd, PBVH_ITER_UNIQUE) {
-    SculptVertexNeighborIter ni;
-    float avg = 0.0f;
-    int total = 0;
-    SCULPT_VERTEX_NEIGHBORS_ITER_BEGIN (ss, vd.vertex, ni) {
-      avg += pose_factor[ni.index];
-      total++;
-    }
-    SCULPT_VERTEX_NEIGHBORS_ITER_END(ni);
+  struct LocalData {
+    Vector<int> vert_indices;
+    Vector<Vector<int>> vert_neighbors;
+    Vector<float> new_factors;
+  };
+  const SculptSession &ss = *object.sculpt;
+  Vector<bke::pbvh::Node *> nodes = bke::pbvh::search_gather(*ss.pbvh, {});
 
-    if (total > 0) {
-      pose_factor[vd.index] = avg / total;
+  threading::EnumerableThreadSpecific<LocalData> all_tls;
+  switch (ss.pbvh->type()) {
+    case bke::pbvh::Type::Mesh: {
+      const Mesh &mesh = *static_cast<const Mesh *>(object.data);
+      const OffsetIndices faces = mesh.faces();
+      const Span<int> corner_verts = mesh.corner_verts();
+      const GroupedSpan<int> vert_to_face_map = mesh.vert_to_face_map();
+      const bke::AttributeAccessor attributes = mesh.attributes();
+      const VArraySpan<bool> hide_vert = *attributes.lookup<bool>(".hide_vert",
+                                                                  bke::AttrDomain::Point);
+      const VArraySpan<bool> hide_poly = *attributes.lookup<bool>(".hide_poly",
+                                                                  bke::AttrDomain::Face);
+      for ([[maybe_unused]] const int _ : IndexRange(iterations)) {
+        threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+          LocalData &tls = all_tls.local();
+          for (const int i : range) {
+            const Span<int> verts = hide::node_visible_verts(
+                *nodes[i], hide_vert, tls.vert_indices);
+
+            tls.vert_neighbors.resize(verts.size());
+            const MutableSpan<Vector<int>> neighbors = tls.vert_neighbors;
+            calc_vert_neighbors(
+                faces, corner_verts, vert_to_face_map, hide_poly, verts, neighbors);
+
+            tls.new_factors.resize(verts.size());
+            const MutableSpan<float> new_factors = tls.new_factors;
+            smooth::neighbor_data_average_mesh(data.as_span(), neighbors, new_factors);
+
+            scatter_data_mesh(new_factors.as_span(), verts, data);
+          }
+        });
+      }
+      break;
+    }
+    case bke::pbvh::Type::Grids: {
+      const SubdivCCG &subdiv_ccg = *ss.subdiv_ccg;
+      const CCGKey key = BKE_subdiv_ccg_key_top_level(subdiv_ccg);
+      const BitGroupVector<> &grid_hidden = subdiv_ccg.grid_hidden;
+      for ([[maybe_unused]] const int _ : IndexRange(iterations)) {
+        threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+          LocalData &tls = all_tls.local();
+          for (const int i : range) {
+            const Span<int> grids = bke::pbvh::node_grid_indices(*nodes[i]);
+            const int grid_verts_num = key.grid_area * grids.size();
+
+            tls.new_factors.resize(grid_verts_num);
+            const MutableSpan<float> new_factors = tls.new_factors;
+            smooth::average_data_grids(subdiv_ccg, data.as_span(), grids, new_factors);
+
+            if (grid_hidden.is_empty()) {
+              scatter_data_grids(subdiv_ccg, new_factors.as_span(), grids, data);
+            }
+            else {
+              for (const int i : grids.index_range()) {
+                const int node_start = i * key.grid_area;
+                BKE_subdiv_ccg_foreach_visible_grid_vert(
+                    key, grid_hidden, grids[i], [&](const int offset) {
+                      data[i] = new_factors[node_start + offset];
+                    });
+              }
+            }
+          }
+        });
+      }
+      break;
+    }
+    case bke::pbvh::Type::BMesh: {
+      for ([[maybe_unused]] const int _ : IndexRange(iterations)) {
+        threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+          LocalData &tls = all_tls.local();
+          for (const int node_index : range) {
+            const Set<BMVert *, 0> &verts = BKE_pbvh_bmesh_node_unique_verts(nodes[node_index]);
+
+            tls.new_factors.resize(verts.size());
+            const MutableSpan<float> new_factors = tls.new_factors;
+            smooth::average_data_bmesh(data.as_span(), verts, new_factors);
+
+            int i = 0;
+            for (const BMVert *vert : verts) {
+              if (BM_elem_flag_test(vert, BM_ELEM_HIDDEN)) {
+                i++;
+                continue;
+              }
+              data[BM_elem_index_get(vert)] = new_factors[i];
+              i++;
+            }
+          }
+        });
+      }
+      break;
     }
   }
-  BKE_pbvh_vertex_iter_end;
 }
 
 /* Init the IK chain with empty weights. */
@@ -1066,24 +1341,13 @@ static std::unique_ptr<SculptPoseIKChain> ik_chain_init(Object &ob,
 
 void pose_brush_init(Object &ob, SculptSession &ss, const Brush &brush)
 {
-  bke::pbvh::Tree &pbvh = *ob.sculpt->pbvh;
-
-  Vector<bke::pbvh::Node *> nodes = bke::pbvh::search_gather(pbvh, {});
-
   /* Init the IK chain that is going to be used to deform the vertices. */
   ss.cache->pose_ik_chain = ik_chain_init(
       ob, ss, brush, ss.cache->true_location, ss.cache->radius);
 
   /* Smooth the weights of each segment for cleaner deformation. */
   for (SculptPoseIKChainSegment &segment : ss.cache->pose_ik_chain->segments) {
-    MutableSpan<float> pose_factor = segment.weights;
-    for (int i = 0; i < brush.pose_smooth_iterations; i++) {
-      threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
-        for (const int i : range) {
-          pose_brush_init_task(ss, pose_factor, nodes[i]);
-        }
-      });
-    }
+    smooth_geometry_data_array(ob, brush.pose_smooth_iterations, segment.weights);
   }
 }