Refactor: Sculpt: Specialize IK chain FK face set calculation

Part of #118145.

* Removes usages of PBVHVertRef and flood_fill:execute for specific
  versions per PBVH type
* Adds stub method for BMesh face set retrieval to make future
  implementation easier.

Pull Request: https://projects.blender.org/blender/blender/pulls/127666

source/blender/editors/sculpt_paint/sculpt.cc

@@ -318,6 +318,11 @@ int vert_face_set_get(const GroupedSpan<int> vert_to_face_map,
   return face_set;
 }
 
+int vert_face_set_get(const int /*face_set_offset*/, const BMVert & /*vert*/)
+{
+  return SCULPT_FACE_SET_NONE;
+}
+
 bool vert_has_face_set(const GroupedSpan<int> vert_to_face_map,
                        const Span<int> face_sets,
                        const int vert,

source/blender/editors/sculpt_paint/sculpt_face_set.hh

@@ -27,6 +27,7 @@ namespace blender::ed::sculpt_paint::face_set {
 int active_face_set_get(const Object &object);
 int vert_face_set_get(const Object &object, PBVHVertRef vertex);
 int vert_face_set_get(GroupedSpan<int> vert_to_face_map, Span<int> face_sets, int vert);
+int vert_face_set_get(int face_set_offset, const BMVert &vert);
 
 bool vert_has_face_set(const Object &object, PBVHVertRef vertex, int face_set);
 bool vert_has_face_set(GroupedSpan<int> vert_to_face_map,

source/blender/editors/sculpt_paint/sculpt_pose.cc

@@ -894,12 +894,12 @@ static int brush_num_effective_segments(const Brush &brush)
   return brush.pose_ik_segments;
 }
 
-static std::unique_ptr<IKChain> pose_ik_chain_init_topology(const Depsgraph &depsgraph,
-                                                            Object &object,
-                                                            SculptSession &ss,
-                                                            const Brush &brush,
-                                                            const float3 &initial_location,
-                                                            const float radius)
+static std::unique_ptr<IKChain> ik_chain_init_topology(const Depsgraph &depsgraph,
+                                                       Object &object,
+                                                       SculptSession &ss,
+                                                       const Brush &brush,
+                                                       const float3 &initial_location,
+                                                       const float radius)
 {
 
   const float chain_segment_len = radius * (1.0f + brush.pose_offset);
@@ -1154,21 +1154,73 @@ static std::unique_ptr<IKChain> ik_chain_init_face_sets(const Depsgraph &depsgra
 
 static std::optional<float3> calc_average_face_set_center(const Depsgraph &depsgraph,
                                                           Object &object,
-                                                          const int totvert,
                                                           const Span<int> floodfill_step,
                                                           const int active_face_set,
                                                           const int target_face_set)
 {
   int count = 0;
   float3 sum(0.0f);
-  for (int i = 0; i < totvert; i++) {
-    const PBVHVertRef vertex = BKE_pbvh_index_to_vertex(object, i);
 
-    if (floodfill_step[i] != 0 && face_set::vert_has_face_set(object, vertex, active_face_set) &&
-        face_set::vert_has_face_set(object, vertex, target_face_set))
-    {
-      sum += SCULPT_vertex_co_get(depsgraph, object, vertex);
-      count++;
+  switch (bke::object::pbvh_get(object)->type()) {
+    case bke::pbvh::Type::Mesh: {
+      Mesh &mesh = *static_cast<Mesh *>(object.data);
+      Span<float3> vert_positions = bke::pbvh::vert_positions_eval(depsgraph, object);
+      const bke::AttributeAccessor attributes = mesh.attributes();
+      const VArraySpan face_sets = *attributes.lookup_or_default<int>(
+          ".sculpt_face_set", bke::AttrDomain::Face, 0);
+
+      for (const int vert : vert_positions.index_range()) {
+        if (floodfill_step[vert] != 0 &&
+            face_set::vert_has_face_set(
+                mesh.vert_to_face_map(), face_sets, vert, active_face_set) &&
+            face_set::vert_has_face_set(mesh.vert_to_face_map(), face_sets, vert, target_face_set))
+        {
+          sum += vert_positions[vert];
+          count++;
+        }
+      }
+      break;
+    }
+    case bke::pbvh::Type::Grids: {
+      SubdivCCG &subdiv_ccg = *object.sculpt->subdiv_ccg;
+      MutableSpan<float3> positions = subdiv_ccg.positions;
+
+      Mesh &mesh = *static_cast<Mesh *>(object.data);
+      const bke::AttributeAccessor attributes = mesh.attributes();
+      const VArraySpan face_sets = *attributes.lookup_or_default<int>(
+          ".sculpt_face_set", bke::AttrDomain::Face, 0);
+
+      const CCGKey key = BKE_subdiv_ccg_key_top_level(subdiv_ccg);
+      for (const int grid : IndexRange(subdiv_ccg.grids_num)) {
+        for (const int index : bke::ccg::grid_range(key, grid)) {
+          if (floodfill_step[index] != 0 &&
+              face_set::vert_has_face_set(subdiv_ccg, face_sets, grid, active_face_set) &&
+              face_set::vert_has_face_set(subdiv_ccg, face_sets, grid, target_face_set))
+          {
+            sum += positions[index];
+            count++;
+          }
+        }
+      }
+      break;
+    }
+    case bke::pbvh::Type::BMesh: {
+      SCULPT_vertex_random_access_ensure(object);
+      BMesh &bm = *object.sculpt->bm;
+      const int face_set_offset = CustomData_get_offset_named(
+          &bm.pdata, CD_PROP_INT32, ".sculpt_face_set");
+      for (const int vert : IndexRange(BM_mesh_elem_count(&bm, BM_VERT))) {
+        const BMVert *bm_vert = BM_vert_at_index(&bm, vert);
+        if (floodfill_step[vert] != 0 &&
+            face_set::vert_has_face_set(face_set_offset, *bm_vert, active_face_set) &&
+            face_set::vert_has_face_set(face_set_offset, *bm_vert, target_face_set))
+        {
+          sum += bm_vert->co;
+          count++;
+        }
+      }
+
+      break;
     }
   }
 
@@ -1179,35 +1231,121 @@ static std::optional<float3> calc_average_face_set_center(const Depsgraph &depsg
   return std::nullopt;
 }
 
-static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk(const Depsgraph &depsgraph,
-                                                           Object &object,
-                                                           SculptSession &ss,
-                                                           const float radius,
-                                                           const float3 &initial_location)
+static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk_mesh(const Depsgraph &depsgraph,
+                                                                Object &object,
+                                                                SculptSession &ss,
+                                                                const float radius,
+                                                                const float3 &initial_location)
 {
-  const int totvert = SCULPT_vertex_count_get(object);
+  Mesh &mesh = *static_cast<Mesh *>(object.data);
+  const bke::AttributeAccessor attributes = mesh.attributes();
+  const VArraySpan face_sets = *attributes.lookup_or_default<int>(
+      ".sculpt_face_set", bke::AttrDomain::Face, 0);
 
-  std::unique_ptr<IKChain> ik_chain = ik_chain_new(1, totvert);
+  std::unique_ptr<IKChain> ik_chain = ik_chain_new(1, mesh.verts_num);
 
-  const PBVHVertRef active_vertex = ss.active_vert_ref();
+  const int active_vert = std::get<int>(ss.active_vert());
   const bke::pbvh::Tree &pbvh = *bke::object::pbvh_get(object);
-  const int active_vertex_index = BKE_pbvh_vertex_to_index(pbvh, active_vertex);
 
   const int active_face_set = face_set::active_face_set_get(object);
 
   Set<int> visited_face_sets;
-  Array<int> floodfill_step(totvert);
-  floodfill_step[active_vertex_index] = 1;
+  Array<int> floodfill_step(mesh.verts_num);
+  floodfill_step[active_vert] = 1;
 
   int masked_face_set = SCULPT_FACE_SET_NONE;
   int target_face_set = SCULPT_FACE_SET_NONE;
   int masked_face_set_it = 0;
-  flood_fill::FillData step_floodfill = flood_fill::init_fill(object);
-  flood_fill::add_initial(step_floodfill, active_vertex);
-  flood_fill::execute(
-      object, step_floodfill, [&](PBVHVertRef from_v, PBVHVertRef to_v, bool is_duplicate) {
-        const int from_v_i = BKE_pbvh_vertex_to_index(pbvh, from_v);
-        const int to_v_i = BKE_pbvh_vertex_to_index(pbvh, to_v);
+  flood_fill::FillDataMesh step_floodfill(mesh.verts_num);
+  step_floodfill.add_initial(active_vert);
+  step_floodfill.execute(object, mesh.vert_to_face_map(), [&](int from_v, int to_v) {
+    floodfill_step[to_v] = floodfill_step[from_v] + 1;
+
+    const int to_face_set = face_set::vert_face_set_get(mesh.vert_to_face_map(), face_sets, to_v);
+    if (!visited_face_sets.contains(to_face_set)) {
+      if (face_set::vert_has_unique_face_set(mesh.vert_to_face_map(), face_sets, to_v) &&
+          !face_set::vert_has_unique_face_set(mesh.vert_to_face_map(), face_sets, from_v) &&
+          face_set::vert_has_face_set(mesh.vert_to_face_map(), face_sets, from_v, to_face_set))
+      {
+
+        visited_face_sets.add(to_face_set);
+
+        if (floodfill_step[to_v] >= masked_face_set_it) {
+          masked_face_set = to_face_set;
+          masked_face_set_it = floodfill_step[to_v];
+        }
+
+        if (target_face_set == SCULPT_FACE_SET_NONE) {
+          target_face_set = to_face_set;
+        }
+      }
+    }
+
+    return face_set::vert_has_face_set(mesh.vert_to_face_map(), face_sets, to_v, active_face_set);
+  });
+
+  const std::optional<float3> origin = calc_average_face_set_center(
+      depsgraph, object, floodfill_step, active_face_set, masked_face_set);
+  ik_chain->segments[0].orig = origin.value_or(float3(0));
+
+  std::optional<float3> head = std::nullopt;
+  if (target_face_set != masked_face_set) {
+    head = calc_average_face_set_center(
+        depsgraph, object, floodfill_step, active_face_set, target_face_set);
+  }
+
+  ik_chain->segments[0].head = head.value_or(initial_location);
+  ik_chain->grab_delta_offset = ik_chain->segments[0].head - initial_location;
+
+  flood_fill::FillDataMesh weight_floodfill(mesh.verts_num);
+  weight_floodfill.add_initial_with_symmetry(depsgraph, object, pbvh, active_vert, radius);
+  MutableSpan<float> fk_weights = ik_chain->segments[0].weights;
+  weight_floodfill.execute(object, mesh.vert_to_face_map(), [&](int /*from_v*/, int to_v) {
+    fk_weights[to_v] = 1.0f;
+    return !face_set::vert_has_face_set(mesh.vert_to_face_map(), face_sets, to_v, masked_face_set);
+  });
+
+  ik_chain_origin_heads_init(*ik_chain, ik_chain->segments[0].head);
+  return ik_chain;
+}
+
+static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk_grids(const Depsgraph &depsgraph,
+                                                                 Object &object,
+                                                                 SculptSession &ss,
+                                                                 const float radius,
+                                                                 const float3 &initial_location)
+{
+  const Mesh &mesh = *static_cast<const Mesh *>(object.data);
+  const bke::AttributeAccessor attributes = mesh.attributes();
+  const VArraySpan face_sets = *attributes.lookup_or_default<int>(
+      ".sculpt_face_set", bke::AttrDomain::Face, 0);
+
+  const SubdivCCG &subdiv_ccg = *ss.subdiv_ccg;
+  const CCGKey key = BKE_subdiv_ccg_key_top_level(subdiv_ccg);
+  const Span<int> grid_to_face_map = subdiv_ccg.grid_to_face_map;
+  const int grids_num = ss.subdiv_ccg->grids_num * key.grid_area;
+
+  std::unique_ptr<IKChain> ik_chain = ik_chain_new(1, grids_num);
+
+  const SubdivCCGCoord active_vert = std::get<SubdivCCGCoord>(ss.active_vert());
+  const bke::pbvh::Tree &pbvh = *bke::object::pbvh_get(object);
+  const int active_vert_index = ss.active_vert_index();
+
+  const int active_face_set = face_set::active_face_set_get(object);
+
+  Set<int> visited_face_sets;
+  Array<int> floodfill_step(grids_num);
+  floodfill_step[active_vert_index] = 1;
+
+  int masked_face_set = SCULPT_FACE_SET_NONE;
+  int target_face_set = SCULPT_FACE_SET_NONE;
+  int masked_face_set_it = 0;
+  flood_fill::FillDataGrids step_floodfill(grids_num);
+  step_floodfill.add_initial(active_vert);
+  step_floodfill.execute(
+      object, subdiv_ccg, [&](SubdivCCGCoord from_v, SubdivCCGCoord to_v, bool is_duplicate) {
+        const int from_v_i = from_v.to_index(key);
+        const int to_v_i = to_v.to_index(key);
 
         if (!is_duplicate) {
           floodfill_step[to_v_i] = floodfill_step[from_v_i] + 1;
@@ -1216,11 +1354,21 @@ static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk(const Depsgraph &deps
           floodfill_step[to_v_i] = floodfill_step[from_v_i];
         }
 
-        const int to_face_set = face_set::vert_face_set_get(object, to_v);
+        const int to_face_set = face_sets[grid_to_face_map[to_v.grid_index]];
         if (!visited_face_sets.contains(to_face_set)) {
-          if (face_set::vert_has_unique_face_set(object, to_v) &&
-              !face_set::vert_has_unique_face_set(object, from_v) &&
-              face_set::vert_has_face_set(object, from_v, to_face_set))
+          if (face_set::vert_has_unique_face_set(mesh.vert_to_face_map(),
+                                                 mesh.corner_verts(),
+                                                 mesh.faces(),
+                                                 face_sets,
+                                                 subdiv_ccg,
+                                                 to_v) &&
+              !face_set::vert_has_unique_face_set(mesh.vert_to_face_map(),
+                                                  mesh.corner_verts(),
+                                                  mesh.faces(),
+                                                  face_sets,
+                                                  subdiv_ccg,
+                                                  from_v) &&
+              face_set::vert_has_face_set(subdiv_ccg, face_sets, from_v.grid_index, to_face_set))
           {
 
             visited_face_sets.add(to_face_set);
@@ -1236,39 +1384,145 @@ static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk(const Depsgraph &deps
           }
         }
 
-        return face_set::vert_has_face_set(object, to_v, active_face_set);
+        return face_set::vert_has_face_set(
+            subdiv_ccg, face_sets, to_v.grid_index, active_face_set);
       });
 
   const std::optional<float3> origin = calc_average_face_set_center(
-      depsgraph, object, totvert, floodfill_step, active_face_set, masked_face_set);
+      depsgraph, object, floodfill_step, active_face_set, masked_face_set);
   ik_chain->segments[0].orig = origin.value_or(float3(0));
 
   std::optional<float3> head = std::nullopt;
   if (target_face_set != masked_face_set) {
     head = calc_average_face_set_center(
-        depsgraph, object, totvert, floodfill_step, active_face_set, target_face_set);
+        depsgraph, object, floodfill_step, active_face_set, target_face_set);
   }
 
   ik_chain->segments[0].head = head.value_or(initial_location);
   ik_chain->grab_delta_offset = ik_chain->segments[0].head - initial_location;
 
-  flood_fill::FillData weight_floodfill = flood_fill::init_fill(object);
-  flood_fill::add_initial_with_symmetry(
-      depsgraph, object, weight_floodfill, ss.active_vert_ref(), radius);
+  flood_fill::FillDataGrids weight_floodfill(grids_num);
+  weight_floodfill.add_initial_with_symmetry(object, pbvh, subdiv_ccg, active_vert, radius);
   MutableSpan<float> fk_weights = ik_chain->segments[0].weights;
-  flood_fill::execute(object,
-                      weight_floodfill,
-                      [&](PBVHVertRef /*from_v*/, PBVHVertRef to_v, bool /*is_duplicate*/) {
-                        int to_v_i = BKE_pbvh_vertex_to_index(pbvh, to_v);
+  weight_floodfill.execute(
+      object,
+      subdiv_ccg,
+      [&](SubdivCCGCoord /*from_v*/, SubdivCCGCoord to_v, bool /*is_duplicate*/) {
+        int to_v_i = to_v.to_index(key);
 
-                        fk_weights[to_v_i] = 1.0f;
-                        return !face_set::vert_has_face_set(object, to_v, masked_face_set);
-                      });
+        fk_weights[to_v_i] = 1.0f;
+        return !face_set::vert_has_face_set(
+            subdiv_ccg, face_sets, to_v.grid_index, masked_face_set);
+      });
 
   ik_chain_origin_heads_init(*ik_chain, ik_chain->segments[0].head);
   return ik_chain;
 }
 
+static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk_bmesh(const Depsgraph &depsgraph,
+                                                                 Object &object,
+                                                                 SculptSession &ss,
+                                                                 const float radius,
+                                                                 const float3 &initial_location)
+{
+  SCULPT_vertex_random_access_ensure(object);
+
+  BMesh &bm = *ss.bm;
+  const int face_set_offset = CustomData_get_offset_named(
+      &bm.pdata, CD_PROP_INT32, ".sculpt_face_set");
+  const int verts_num = BM_mesh_elem_count(&bm, BM_VERT);
+
+  std::unique_ptr<IKChain> ik_chain = ik_chain_new(1, verts_num);
+
+  BMVert *active_vert = std::get<BMVert *>(ss.active_vert());
+  const bke::pbvh::Tree &pbvh = *bke::object::pbvh_get(object);
+  const int active_vert_index = BM_elem_index_get(active_vert);
+
+  const int active_face_set = face_set::active_face_set_get(object);
+
+  Set<int> visited_face_sets;
+  Array<int> floodfill_step(verts_num);
+  floodfill_step[active_vert_index] = 1;
+
+  int masked_face_set = SCULPT_FACE_SET_NONE;
+  int target_face_set = SCULPT_FACE_SET_NONE;
+  int masked_face_set_it = 0;
+  flood_fill::FillDataBMesh step_floodfill(verts_num);
+  step_floodfill.add_initial(active_vert);
+  step_floodfill.execute(object, [&](BMVert *from_v, BMVert *to_v) {
+    const int from_v_i = BM_elem_index_get(from_v);
+    const int to_v_i = BM_elem_index_get(to_v);
+
+    floodfill_step[to_v_i] = floodfill_step[from_v_i] + 1;
+
+    const int to_face_set = face_set::vert_face_set_get(face_set_offset, *to_v);
+    if (!visited_face_sets.contains(to_face_set)) {
+      if (face_set::vert_has_unique_face_set(face_set_offset, *to_v) &&
+          !face_set::vert_has_unique_face_set(face_set_offset, *from_v) &&
+          face_set::vert_has_face_set(face_set_offset, *from_v, to_face_set))
+      {
+
+        visited_face_sets.add(to_face_set);
+
+        if (floodfill_step[to_v_i] >= masked_face_set_it) {
+          masked_face_set = to_face_set;
+          masked_face_set_it = floodfill_step[to_v_i];
+        }
+
+        if (target_face_set == SCULPT_FACE_SET_NONE) {
+          target_face_set = to_face_set;
+        }
+      }
+    }
+
+    return face_set::vert_has_face_set(face_set_offset, *to_v, active_face_set);
+  });
+
+  const std::optional<float3> origin = calc_average_face_set_center(
+      depsgraph, object, floodfill_step, active_face_set, masked_face_set);
+  ik_chain->segments[0].orig = origin.value_or(float3(0));
+
+  std::optional<float3> head = std::nullopt;
+  if (target_face_set != masked_face_set) {
+    head = calc_average_face_set_center(
+        depsgraph, object, floodfill_step, active_face_set, target_face_set);
+  }
+
+  ik_chain->segments[0].head = head.value_or(initial_location);
+  ik_chain->grab_delta_offset = ik_chain->segments[0].head - initial_location;
+
+  flood_fill::FillDataBMesh weight_floodfill(verts_num);
+  weight_floodfill.add_initial_with_symmetry(object, pbvh, active_vert, radius);
+  MutableSpan<float> fk_weights = ik_chain->segments[0].weights;
+  weight_floodfill.execute(object, [&](BMVert * /*from_v*/, BMVert *to_v) {
+    int to_v_i = BM_elem_index_get(to_v);
+
+    fk_weights[to_v_i] = 1.0f;
+    return !face_set::vert_has_face_set(face_set_offset, *to_v, masked_face_set);
+  });
+
+  ik_chain_origin_heads_init(*ik_chain, ik_chain->segments[0].head);
+  return ik_chain;
+}
+
+static std::unique_ptr<IKChain> ik_chain_init_face_sets_fk(const Depsgraph &depsgraph,
+                                                           Object &object,
+                                                           SculptSession &ss,
+                                                           const float radius,
+                                                           const float3 &initial_location)
+{
+  switch (bke::object::pbvh_get(object)->type()) {
+    case bke::pbvh::Type::Mesh:
+      return ik_chain_init_face_sets_fk_mesh(depsgraph, object, ss, radius, initial_location);
+    case bke::pbvh::Type::Grids:
+      return ik_chain_init_face_sets_fk_grids(depsgraph, object, ss, radius, initial_location);
+    case bke::pbvh::Type::BMesh:
+      return ik_chain_init_face_sets_fk_bmesh(depsgraph, object, ss, radius, initial_location);
+  }
+  BLI_assert_unreachable();
+  return nullptr;
+}
+
 static std::unique_ptr<IKChain> ik_chain_init(const Depsgraph &depsgraph,
                                               Object &ob,
                                               SculptSession &ss,
@@ -1287,7 +1541,7 @@ static std::unique_ptr<IKChain> ik_chain_init(const Depsgraph &depsgraph,
 
   switch (brush.pose_origin_type) {
     case BRUSH_POSE_ORIGIN_TOPOLOGY:
-      ik_chain = pose_ik_chain_init_topology(depsgraph, ob, ss, brush, initial_location, radius);
+      ik_chain = ik_chain_init_topology(depsgraph, ob, ss, brush, initial_location, radius);
       break;
     case BRUSH_POSE_ORIGIN_FACE_SETS:
       ik_chain = ik_chain_init_face_sets(depsgraph, ob, ss, brush, radius);