Object: edit-mode data support for object vertex coordinate access

Allows to access/transform/restore edit-mode
coordinates in a generic way.

5 files changed, 323 insertions, 32 deletions

diff --git a/source/blender/blenkernel/intern/editmesh.c b/source/blender/blenkernel/intern/editmesh.c
index d929c953b89..b135574a650 100644
--- a/source/blender/blenkernel/intern/editmesh.c
+++ b/source/blender/blenkernel/intern/editmesh.c
@@ -216,20 +216,7 @@ void BKE_editmesh_color_ensure(BMEditMesh *em, const char htype)
 
 float (*BKE_editmesh_vert_coords_alloc_orco(BMEditMesh *em, int *r_vert_len))[3]
 {
-  BMIter iter;
-  BMVert *eve;
-  float(*orco)[3];
-  int i;
-
-  orco = MEM_mallocN(em->bm->totvert * sizeof(*orco), __func__);
-
-  BM_ITER_MESH_INDEX (eve, &iter, em->bm, BM_VERTS_OF_MESH, i) {
-    copy_v3_v3(orco[i], eve->co);
-  }
-
-  *r_vert_len = em->bm->totvert;
-
-  return orco;
+  return BM_mesh_vert_coords_alloc(em->bm, r_vert_len);
 }
 
 void BKE_editmesh_lnorspace_update(BMEditMesh *em)
diff --git a/source/blender/bmesh/intern/bmesh_mesh.c b/source/blender/bmesh/intern/bmesh_mesh.c
index f44b758ef6b..4790c811e56 100644
--- a/source/blender/bmesh/intern/bmesh_mesh.c
+++ b/source/blender/bmesh/intern/bmesh_mesh.c
@@ -2708,3 +2708,49 @@ void BM_mesh_toolflags_set(BMesh *bm, bool use_toolflags)
 
   bm->use_toolflags = use_toolflags;
 }
+
+/* -------------------------------------------------------------------- */
+/** \name BMesh Coordinate Access
+ * \{ */
+
+void BM_mesh_vert_coords_get(BMesh *bm, float (*vert_coords)[3])
+{
+  BMIter iter;
+  BMVert *v;
+  int i;
+  BM_ITER_MESH_INDEX (v, &iter, bm, BM_VERTS_OF_MESH, i) {
+    copy_v3_v3(vert_coords[i], v->co);
+  }
+}
+
+float (*BM_mesh_vert_coords_alloc(BMesh *bm, int *r_vert_len))[3]
+{
+  float(*vert_coords)[3] = MEM_mallocN(bm->totvert * sizeof(*vert_coords), __func__);
+  BM_mesh_vert_coords_get(bm, vert_coords);
+  *r_vert_len = bm->totvert;
+  return vert_coords;
+}
+
+void BM_mesh_vert_coords_apply(BMesh *bm, const float (*vert_coords)[3])
+{
+  BMIter iter;
+  BMVert *v;
+  int i;
+  BM_ITER_MESH_INDEX (v, &iter, bm, BM_VERTS_OF_MESH, i) {
+    copy_v3_v3(v->co, vert_coords[i]);
+  }
+}
+
+void BM_mesh_vert_coords_apply_with_mat4(BMesh *bm,
+                                         const float (*vert_coords)[3],
+                                         const float mat[4][4])
+{
+  BMIter iter;
+  BMVert *v;
+  int i;
+  BM_ITER_MESH_INDEX (v, &iter, bm, BM_VERTS_OF_MESH, i) {
+    mul_v3_m4v3(v->co, mat, vert_coords[i]);
+  }
+}
+
+/** \} */
diff --git a/source/blender/bmesh/intern/bmesh_mesh.h b/source/blender/bmesh/intern/bmesh_mesh.h
index 83575552acc..604142445ae 100644
--- a/source/blender/bmesh/intern/bmesh_mesh.h
+++ b/source/blender/bmesh/intern/bmesh_mesh.h
@@ -159,4 +159,12 @@ extern const BMAllocTemplate bm_mesh_chunksize_default;
 #define BMALLOC_TEMPLATE_FROM_ME(...) \
   VA_NARGS_CALL_OVERLOAD(_VA_BMALLOC_TEMPLATE_FROM_ME_, __VA_ARGS__)
 
+/* Vertex coords access. */
+void BM_mesh_vert_coords_get(BMesh *bm, float (*orco)[3]);
+float (*BM_mesh_vert_coords_alloc(BMesh *bm, int *r_vert_len))[3];
+void BM_mesh_vert_coords_apply(BMesh *bm, const float (*orco)[3]);
+void BM_mesh_vert_coords_apply_with_mat4(BMesh *bm,
+                                         const float (*vert_coords)[3],
+                                         const float mat[4][4]);
+
 #endif /* __BMESH_MESH_H__ */
diff --git a/source/blender/editors/include/ED_object.h b/source/blender/editors/include/ED_object.h
index 38d75aa57e9..a5f56689c95 100644
--- a/source/blender/editors/include/ED_object.h
+++ b/source/blender/editors/include/ED_object.h
@@ -410,11 +410,17 @@ void ED_object_facemap_face_add(struct Object *ob, struct bFaceMap *fmap, int fa
 void ED_object_facemap_face_remove(struct Object *ob, struct bFaceMap *fmap, int facenum);
 
 /* object_data_transform.c */
+struct XFormObjectData *ED_object_data_xform_create_ex(struct ID *id, bool is_edit_mode);
 struct XFormObjectData *ED_object_data_xform_create(struct ID *id);
+struct XFormObjectData *ED_object_data_xform_create_from_edit_mode(ID *id);
+
 void ED_object_data_xform_destroy(struct XFormObjectData *xod);
 
 void ED_object_data_xform_by_mat4(struct XFormObjectData *xod, const float mat[4][4]);
 
+void ED_object_data_xform_restore(struct XFormObjectData *xod);
+void ED_object_data_xform_tag_update(struct XFormObjectData *xod);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/source/blender/editors/object/object_data_transform.c b/source/blender/editors/object/object_data_transform.c
index ee86c79ead5..51d8fb8fd72 100644
--- a/source/blender/editors/object/object_data_transform.c
+++ b/source/blender/editors/object/object_data_transform.c
@@ -44,16 +44,21 @@
 #include "BLI_utildefines.h"
 
 #include "BKE_curve.h"
+#include "BKE_editmesh.h"
 #include "BKE_mball.h"
 #include "BKE_mesh.h"
 #include "BKE_armature.h"
 #include "BKE_lattice.h"
 
+#include "bmesh.h"
+
 #include "DEG_depsgraph.h"
 
 #include "WM_types.h"
 
 #include "ED_object.h"
+#include "ED_mesh.h"
+#include "ED_armature.h"
 
 #include "MEM_guardedalloc.h"
 
@@ -149,6 +154,70 @@ static void armature_coords_and_quats_apply_with_mat4(bArmature *arm,
   BKE_armature_transform(arm, mat, true);
 }
 
+static void armature_coords_and_quats_apply(bArmature *arm,
+                                            const struct ElemData_Armature *elem_array)
+{
+  /* Avoid code duplication by using a unit matrix. */
+  float mat[4][4];
+  unit_m4(mat);
+  armature_coords_and_quats_apply_with_mat4(arm, elem_array, mat);
+}
+
+/* Edit Armature */
+static void edit_armature_coords_and_quats_get(const bArmature *arm,
+                                               struct ElemData_Armature *elem_array)
+{
+  struct ElemData_Armature *elem = elem_array;
+  for (EditBone *ebone = arm->edbo->first; ebone; ebone = ebone->next, elem++) {
+
+#define COPY_PTR(member) memcpy(elem->member, ebone->member, sizeof(ebone->member))
+#define COPY_VAL(member) memcpy(&elem->member, &ebone->member, sizeof(ebone->member))
+    /* Unused for edit bones: arm_head, arm_tail, arm_roll */
+    COPY_PTR(head);
+    COPY_PTR(tail);
+    COPY_VAL(roll);
+    COPY_VAL(rad_tail);
+    COPY_VAL(rad_head);
+    COPY_VAL(dist);
+    COPY_VAL(xwidth);
+    COPY_VAL(zwidth);
+#undef COPY_PTR
+#undef COPY_VAL
+  }
+}
+
+static void edit_armature_coords_and_quats_apply_with_mat4(
+    bArmature *arm, const struct ElemData_Armature *elem_array, const float mat[4][4])
+{
+  const struct ElemData_Armature *elem = elem_array;
+  for (EditBone *ebone = arm->edbo->first; ebone; ebone = ebone->next, elem++) {
+
+#define COPY_PTR(member) memcpy(ebone->member, elem->member, sizeof(ebone->member))
+#define COPY_VAL(member) memcpy(&ebone->member, &elem->member, sizeof(ebone->member))
+    /* Unused for edit bones: arm_head, arm_tail, arm_roll */
+    COPY_PTR(head);
+    COPY_PTR(tail);
+    COPY_VAL(roll);
+    COPY_VAL(rad_tail);
+    COPY_VAL(rad_head);
+    COPY_VAL(dist);
+    COPY_VAL(xwidth);
+    COPY_VAL(zwidth);
+#undef COPY_PTR
+#undef COPY_VAL
+  }
+  ED_armature_edit_transform(arm, mat, true);
+}
+
+static void edit_armature_coords_and_quats_apply(bArmature *arm,
+                                                 const struct ElemData_Armature *elem_array)
+{
+  /* Avoid code duplication by using a unit matrix. */
+  float mat[4][4];
+  unit_m4(mat);
+  edit_armature_coords_and_quats_apply_with_mat4(arm, elem_array, mat);
+}
+
 /* MetaBall */
 
 struct ElemData_MetaBall {
@@ -183,6 +252,15 @@ static void metaball_coords_and_quats_apply_with_mat4(MetaBall *mb,
   BKE_mball_transform(mb, mat, true);
 }
 
+static void metaball_coords_and_quats_apply(MetaBall *mb,
+                                            const struct ElemData_MetaBall *elem_array)
+{
+  /* Avoid code duplication by using a unit matrix. */
+  float mat[4][4];
+  unit_m4(mat);
+  metaball_coords_and_quats_apply_with_mat4(mb, elem_array, mat);
+}
+
 /** \} */
 
 /* -------------------------------------------------------------------- */
@@ -195,6 +273,7 @@ static void metaball_coords_and_quats_apply_with_mat4(MetaBall *mb,
 
 struct XFormObjectData {
   ID *id;
+  bool is_edit_mode;
 };
 
 struct XFormObjectData_Mesh {
@@ -222,21 +301,32 @@ struct XFormObjectData_MetaBall {
   struct ElemData_MetaBall elem_array[0];
 };
 
-struct XFormObjectData *ED_object_data_xform_create(ID *id)
+struct XFormObjectData *ED_object_data_xform_create_ex(ID *id, bool is_edit_mode)
 {
   struct XFormObjectData *xod_base = NULL;
   switch (GS(id->name)) {
     case ID_ME: {
       Mesh *me = (Mesh *)id;
-      const int elem_array_len = me->totvert;
-      struct XFormObjectData_Mesh *xod = MEM_mallocN(
-          sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
-      BKE_mesh_vert_coords_get(me, xod->elem_array);
-      xod_base = &xod->base;
+      if (is_edit_mode) {
+        BMesh *bm = me->edit_mesh->bm;
+        const int elem_array_len = bm->totvert;
+        struct XFormObjectData_Mesh *xod = MEM_mallocN(
+            sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
+        BM_mesh_vert_coords_get(bm, xod->elem_array);
+        xod_base = &xod->base;
+      }
+      else {
+        const int elem_array_len = me->totvert;
+        struct XFormObjectData_Mesh *xod = MEM_mallocN(
+            sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
+        BKE_mesh_vert_coords_get(me, xod->elem_array);
+        xod_base = &xod->base;
+      }
       break;
     }
     case ID_LT: {
-      Lattice *lt = (Lattice *)id;
+      Lattice *lt_orig = (Lattice *)id;
+      Lattice *lt = is_edit_mode ? lt_orig->editlatt->latt : lt_orig;
       const int elem_array_len = lt->pntsu * lt->pntsv * lt->pntsw;
       struct XFormObjectData_Lattice *xod = MEM_mallocN(
           sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
@@ -251,23 +341,43 @@ struct XFormObjectData *ED_object_data_xform_create(ID *id)
         /* We could support translation. */
         break;
       }
-      const int elem_array_len = BKE_nurbList_verts_count(&cu->nurb);
+
+      ListBase *nurbs;
+      if (is_edit_mode) {
+        EditNurb *editnurb = cu->editnurb;
+        nurbs = &editnurb->nurbs;
+      }
+      else {
+        nurbs = &cu->nurb;
+      }
+
+      const int elem_array_len = BKE_nurbList_verts_count(nurbs);
       struct XFormObjectData_Curve *xod = MEM_mallocN(
           sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
-      BKE_curve_nurbs_vert_coords_get(&cu->nurb, xod->elem_array, elem_array_len);
+      BKE_curve_nurbs_vert_coords_get(nurbs, xod->elem_array, elem_array_len);
       xod_base = &xod->base;
       break;
     }
     case ID_AR: {
       bArmature *arm = (bArmature *)id;
-      const int elem_array_len = BKE_armature_bonelist_count(&arm->bonebase);
-      struct XFormObjectData_Armature *xod = MEM_mallocN(
-          sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
-      armature_coords_and_quats_get(arm, xod->elem_array);
-      xod_base = &xod->base;
+      if (is_edit_mode) {
+        const int elem_array_len = BLI_listbase_count(arm->edbo);
+        struct XFormObjectData_Armature *xod = MEM_mallocN(
+            sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
+        edit_armature_coords_and_quats_get(arm, xod->elem_array);
+        xod_base = &xod->base;
+      }
+      else {
+        const int elem_array_len = BKE_armature_bonelist_count(&arm->bonebase);
+        struct XFormObjectData_Armature *xod = MEM_mallocN(
+            sizeof(*xod) + (sizeof(*xod->elem_array) * elem_array_len), __func__);
+        armature_coords_and_quats_get(arm, xod->elem_array);
+        xod_base = &xod->base;
+      }
       break;
     }
     case ID_MB: {
+      /* Edit mode and object mode are shared. */
       MetaBall *mb = (MetaBall *)id;
       const int elem_array_len = BLI_listbase_count(&mb->elems);
       struct XFormObjectData_MetaBall *xod = MEM_mallocN(
@@ -282,10 +392,21 @@ struct XFormObjectData *ED_object_data_xform_create(ID *id)
   }
   if (xod_base) {
     xod_base->id = id;
+    xod_base->is_edit_mode = is_edit_mode;
   }
   return xod_base;
 }
 
+struct XFormObjectData *ED_object_data_xform_create(ID *id)
+{
+  return ED_object_data_xform_create_ex(id, false);
+}
+
+struct XFormObjectData *ED_object_data_xform_create_from_edit_mode(ID *id)
+{
+  return ED_object_data_xform_create_ex(id, true);
+}
+
 void ED_object_data_xform_destroy(struct XFormObjectData *xod)
 {
   MEM_freeN(xod);
@@ -297,28 +418,49 @@ void ED_object_data_xform_by_mat4(struct XFormObjectData *xod_base, const float
     case ID_ME: {
       Mesh *me = (Mesh *)xod_base->id;
       struct XFormObjectData_Mesh *xod = (struct XFormObjectData_Mesh *)xod_base;
-      BKE_mesh_vert_coords_apply_with_mat4(me, xod->elem_array, mat);
+      if (xod_base->is_edit_mode) {
+        BMesh *bm = me->edit_mesh->bm;
+        BM_mesh_vert_coords_apply_with_mat4(bm, xod->elem_array, mat);
+      }
+      else {
+        BKE_mesh_vert_coords_apply_with_mat4(me, xod->elem_array, mat);
+      }
       break;
     }
     case ID_LT: {
-      Lattice *lt = (Lattice *)xod_base->id;
+      Lattice *lt_orig = (Lattice *)xod_base->id;
+      Lattice *lt = xod_base->is_edit_mode ? lt_orig->editlatt->latt : lt_orig;
       struct XFormObjectData_Lattice *xod = (struct XFormObjectData_Lattice *)xod_base;
       BKE_lattice_vert_coords_apply_with_mat4(lt, xod->elem_array, mat);
       break;
     }
     case ID_CU: {
+      BLI_assert(xod_base->is_edit_mode == false); /* Not used currently. */
       Curve *cu = (Curve *)xod_base->id;
       struct XFormObjectData_Curve *xod = (struct XFormObjectData_Curve *)xod_base;
-      BKE_curve_nurbs_vert_coords_apply_with_mat4(&cu->nurb, xod->elem_array, mat, true);
+      if (xod_base->is_edit_mode) {
+        EditNurb *editnurb = cu->editnurb;
+        BKE_curve_nurbs_vert_coords_apply_with_mat4(&editnurb->nurbs, xod->elem_array, mat, true);
+      }
+      else {
+        BKE_curve_nurbs_vert_coords_apply_with_mat4(&cu->nurb, xod->elem_array, mat, true);
+      }
       break;
     }
     case ID_AR: {
+      BLI_assert(xod_base->is_edit_mode == false); /* Not used currently. */
       bArmature *arm = (bArmature *)xod_base->id;
       struct XFormObjectData_Armature *xod = (struct XFormObjectData_Armature *)xod_base;
-      armature_coords_and_quats_apply_with_mat4(arm, xod->elem_array, mat);
+      if (xod_base->is_edit_mode) {
+        edit_armature_coords_and_quats_apply_with_mat4(arm, xod->elem_array, mat);
+      }
+      else {
+        armature_coords_and_quats_apply_with_mat4(arm, xod->elem_array, mat);
+      }
       break;
     }
     case ID_MB: {
+      /* Metaballs are a special case, edit-mode and object mode data is shared. */
       MetaBall *mb = (MetaBall *)xod_base->id;
       struct XFormObjectData_MetaBall *xod = (struct XFormObjectData_MetaBall *)xod_base;
       metaball_coords_and_quats_apply_with_mat4(mb, xod->elem_array, mat);
@@ -330,4 +472,106 @@ void ED_object_data_xform_by_mat4(struct XFormObjectData *xod_base, const float
   }
 }
 
+void ED_object_data_xform_restore(struct XFormObjectData *xod_base)
+{
+  switch (GS(xod_base->id->name)) {
+    case ID_ME: {
+      Mesh *me = (Mesh *)xod_base->id;
+      struct XFormObjectData_Mesh *xod = (struct XFormObjectData_Mesh *)xod_base;
+      if (xod_base->is_edit_mode) {
+        BMesh *bm = me->edit_mesh->bm;
+        BM_mesh_vert_coords_apply(bm, xod->elem_array);
+      }
+      else {
+        BKE_mesh_vert_coords_apply(me, xod->elem_array);
+      }
+      break;
+    }
+    case ID_LT: {
+      Lattice *lt_orig = (Lattice *)xod_base->id;
+      Lattice *lt = xod_base->is_edit_mode ? lt_orig->editlatt->latt : lt_orig;
+      struct XFormObjectData_Lattice *xod = (struct XFormObjectData_Lattice *)xod_base;
+      BKE_lattice_vert_coords_apply(lt, xod->elem_array);
+      break;
+    }
+    case ID_CU: {
+      Curve *cu = (Curve *)xod_base->id;
+      struct XFormObjectData_Curve *xod = (struct XFormObjectData_Curve *)xod_base;
+      if (xod_base->is_edit_mode) {
+        EditNurb *editnurb = cu->editnurb;
+        BKE_curve_nurbs_vert_coords_apply(&editnurb->nurbs, xod->elem_array, true);
+      }
+      else {
+        BKE_curve_nurbs_vert_coords_apply(&cu->nurb, xod->elem_array, true);
+      }
+      break;
+    }
+    case ID_AR: {
+      bArmature *arm = (bArmature *)xod_base->id;
+      struct XFormObjectData_Armature *xod = (struct XFormObjectData_Armature *)xod_base;
+      if (xod_base->is_edit_mode) {
+        edit_armature_coords_and_quats_apply(arm, xod->elem_array);
+      }
+      else {
+        armature_coords_and_quats_apply(arm, xod->elem_array);
+      }
+      break;
+    }
+    case ID_MB: {
+      /* Metaballs are a special case, edit-mode and object mode data is shared. */
+      MetaBall *mb = (MetaBall *)xod_base->id;
+      struct XFormObjectData_MetaBall *xod = (struct XFormObjectData_MetaBall *)xod_base;
+      metaball_coords_and_quats_apply(mb, xod->elem_array);
+      break;
+    }
+    default: {
+      break;
+    }
+  }
+}
+
+void ED_object_data_xform_tag_update(struct XFormObjectData *xod_base)
+{
+  switch (GS(xod_base->id->name)) {
+    case ID_ME: {
+      Mesh *me = (Mesh *)xod_base->id;
+      if (xod_base->is_edit_mode) {
+        EDBM_update_generic(me->edit_mesh, true, false);
+        EDBM_mesh_normals_update(me->edit_mesh);
+      }
+      DEG_id_tag_update(&me->id, ID_RECALC_GEOMETRY);
+      break;
+    }
+    case ID_LT: {
+      /* Generic update. */
+      Lattice *lt = (Lattice *)xod_base->id;
+      DEG_id_tag_update(&lt->id, ID_RECALC_GEOMETRY);
+      break;
+    }
+    case ID_CU: {
+      /* Generic update. */
+      Curve *cu = (Curve *)xod_base->id;
+      DEG_id_tag_update(&cu->id, ID_RECALC_GEOMETRY);
+      break;
+    }
+    case ID_AR: {
+      /* Generic update. */
+      bArmature *arm = (bArmature *)xod_base->id;
+      /* XXX, zero is needed, no other flags properly update this. */
+      DEG_id_tag_update(&arm->id, 0);
+      break;
+    }
+    case ID_MB: {
+      /* Generic update. */
+      MetaBall *mb = (MetaBall *)xod_base->id;
+      DEG_id_tag_update(&mb->id, ID_RECALC_GEOMETRY | ID_RECALC_COPY_ON_WRITE);
+      break;
+    }
+
+    default: {
+      break;
+    }
+  }
+}
+
 /** \} */