Anim: bone collections, store expanded/collapsed state in DNA

Store the 'expanded/collapsed' state of the bone collection tree view in
the DNA data of the bone collections themselves. This way the tree state
is restored when loading the file.

This commit also adds some code to the abstract tree view classes, for
supporting synchronisation of the extended/collapsed state between it
and external data. It follows the same approach as the handling of the
active element.

RNA wrappers have been added to make it possible for Python code to
expand/collapse parts of the tree.

Library overrides are supported for this property, so the
expanded/collapsed state of linked armatures can be locally saved. If
there is no override, the `is_expanded` property is still editable;
changes will not be saved to file in that case, though.

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

source/blender/animrig/ANIM_bone_collections.hh

@@ -256,6 +256,11 @@ void ANIM_armature_refresh_solo_active(bArmature *armature);
 bool ANIM_armature_bonecoll_is_visible_effectively(const bArmature *armature,
                                                    const BoneCollection *bcoll);
 
+/**
+ * Expand or collapse a bone collection in the tree view.
+ */
+void ANIM_armature_bonecoll_is_expanded_set(BoneCollection *bcoll, bool is_expanded);
+
 /**
  * Assign the bone to the bone collection.
  *
@@ -381,6 +386,16 @@ bool armature_bonecoll_is_descendant_of(const bArmature *armature,
 
 bool bonecoll_has_children(const BoneCollection *bcoll);
 
+/**
+ * For each bone collection in the destination armature, copy its #BONE_COLLECTION_EXPANDED flag
+ * from the corresponding bone collection in the source armature.
+ *
+ * This is used in the handling of undo steps, to ensure that undo'ing does _not_
+ * modify this flag.
+ */
+void bonecolls_copy_expanded_flag(Span<BoneCollection *> bcolls_dest,
+                                  Span<const BoneCollection *> bcolls_source);
+
 /**
  * Move a bone collection from one parent to another.
  *

source/blender/animrig/intern/bone_collections.cc

@@ -711,9 +711,11 @@ void ANIM_armature_bonecoll_remove(bArmature *armature, BoneCollection *bcoll)
                                            armature_bonecoll_find_index(armature, bcoll));
 }
 
-BoneCollection *ANIM_armature_bonecoll_get_by_name(bArmature *armature, const char *name)
+template<typename MaybeConstBoneCollection>
+static MaybeConstBoneCollection *bonecolls_get_by_name(
+    blender::Span<MaybeConstBoneCollection *> bonecolls, const char *name)
 {
-  for (BoneCollection *bcoll : armature->collections_span()) {
+  for (MaybeConstBoneCollection *bcoll : bonecolls) {
     if (STREQ(bcoll->name, name)) {
       return bcoll;
     }
@@ -721,6 +723,11 @@ BoneCollection *ANIM_armature_bonecoll_get_by_name(bArmature *armature, const ch
   return nullptr;
 }
 
+BoneCollection *ANIM_armature_bonecoll_get_by_name(bArmature *armature, const char *name)
+{
+  return bonecolls_get_by_name(armature->collections_span(), name);
+}
+
 int ANIM_armature_bonecoll_get_index_by_name(bArmature *armature, const char *name)
 {
   for (int index = 0; index < armature->collection_array_num; index++) {
@@ -846,6 +853,16 @@ bool ANIM_armature_bonecoll_is_visible_effectively(const bArmature *armature,
   return bcoll->is_visible_with_ancestors();
 }
 
+void ANIM_armature_bonecoll_is_expanded_set(BoneCollection *bcoll, bool is_expanded)
+{
+  if (is_expanded) {
+    bcoll->flags |= BONE_COLLECTION_EXPANDED;
+  }
+  else {
+    bcoll->flags &= ~BONE_COLLECTION_EXPANDED;
+  }
+}
+
 /* Store the bone's membership on the collection. */
 static void add_membership(BoneCollection *bcoll, Bone *bone)
 {
@@ -1237,6 +1254,45 @@ bool bonecoll_has_children(const BoneCollection *bcoll)
   return bcoll->child_count > 0;
 }
 
+void bonecolls_copy_expanded_flag(Span<BoneCollection *> bcolls_dest,
+                                  Span<const BoneCollection *> bcolls_source)
+{
+  /* Try to preserve the bone collection expanded/collapsed states. These are UI
+   * changes that shouldn't impact undo steps. Care has to be taken to match the
+   * old and the new bone collections, though, as they may have been reordered
+   * or renamed.
+   *
+   * Reordering is handled by looking up collections by name.
+   * Renames are handled by skipping those that cannot be found by name. */
+
+  auto find_old = [bcolls_source](const char *name, const int index) -> const BoneCollection * {
+    /* Only check index when it's valid in the old armature. */
+    if (index < bcolls_source.size()) {
+      const BoneCollection *bcoll = bcolls_source[index];
+      if (STREQ(bcoll->name, name)) {
+        /* Index and name matches, let's use */
+        return bcoll;
+      }
+    }
+
+    /* Try to find by name as a last resort. This function only works with
+     * non-const pointers, hence the const_cast.  */
+    const BoneCollection *bcoll = bonecolls_get_by_name(bcolls_source, name);
+    return bcoll;
+  };
+
+  for (int i = 0; i < bcolls_dest.size(); i++) {
+    BoneCollection *bcoll_new = bcolls_dest[i];
+
+    const BoneCollection *bcoll_old = find_old(bcoll_new->name, i);
+    if (!bcoll_old) {
+      continue;
+    }
+
+    ANIM_armature_bonecoll_is_expanded_set(bcoll_new, bcoll_old->is_expanded());
+  }
+}
+
 int armature_bonecoll_move_to_parent(bArmature *armature,
                                      const int from_bcoll_index,
                                      int to_child_num,

source/blender/blenkernel/intern/armature.cc

@@ -25,6 +25,7 @@
 #include "BLI_math_vector.h"
 #include "BLI_span.hh"
 #include "BLI_string.h"
+#include "BLI_string_ref.hh"
 #include "BLI_utildefines.h"
 #include "BLT_translation.h"
 
@@ -63,6 +64,8 @@
 
 #include "CLG_log.h"
 
+using namespace blender;
+
 /* -------------------------------------------------------------------- */
 /** \name Prototypes
  * \{ */
@@ -461,6 +464,14 @@ static void armature_blend_read_data(BlendDataReader *reader, ID *id)
   ANIM_armature_runtime_refresh(arm);
 }
 
+static void armature_undo_preserve(BlendLibReader * /*reader*/, ID *id_new, ID *id_old)
+{
+  bArmature *arm_new = (bArmature *)id_new;
+  bArmature *arm_old = (bArmature *)id_old;
+
+  animrig::bonecolls_copy_expanded_flag(arm_new->collections_span(), arm_old->collections_span());
+}
+
 IDTypeInfo IDType_ID_AR = {
     /*id_code*/ ID_AR,
     /*id_filter*/ FILTER_ID_AR,
@@ -485,7 +496,7 @@ IDTypeInfo IDType_ID_AR = {
     /*blend_read_data*/ armature_blend_read_data,
     /*blend_read_after_liblink*/ nullptr,
 
-    /*blend_read_undo_preserve*/ nullptr,
+    /*blend_read_undo_preserve*/ armature_undo_preserve,
 
     /*lib_override_apply_post*/ nullptr,
 };
@@ -3190,5 +3201,9 @@ bool BoneCollection::is_solo() const
 {
   return this->flags & BONE_COLLECTION_SOLO;
 }
+bool BoneCollection::is_expanded() const
+{
+  return this->flags & BONE_COLLECTION_EXPANDED;
+}
 
 /** \} */

source/blender/editors/armature/bone_collections.cc

@@ -135,6 +135,7 @@ static int bone_collection_add_exec(bContext *C, wmOperator * /*op*/)
   }
 
   ANIM_armature_bonecoll_active_set(armature, bcoll);
+  /* TODO: ensure the ancestors of the new bone collection are all expanded. */
 
   WM_event_add_notifier(C, NC_OBJECT | ND_POSE, ob);
   return OPERATOR_FINISHED;

source/blender/editors/armature/editarmature_undo.cc

@@ -97,6 +97,11 @@ static void undoarm_to_editarm(UndoArmature *uarm, bArmature *arm)
 
   ED_armature_ebone_listbase_temp_clear(arm->edbo);
 
+  /* Before freeing the old bone collections, copy their 'expanded' flag. This
+   * flag is not supposed to be restored with any undo steps. */
+  bonecolls_copy_expanded_flag(blender::Span(uarm->collection_array, uarm->collection_array_num),
+                               arm->collections_span());
+
   /* Copy bone collections. */
   ANIM_bonecoll_array_free(&arm->collection_array, &arm->collection_array_num, true);
   auto bcoll_map = ANIM_bonecoll_array_copy_no_membership(&arm->collection_array,

source/blender/editors/include/UI_tree_view.hh

@@ -198,8 +198,22 @@ class AbstractTreeViewItem : public AbstractViewItem, public TreeViewItemContain
   std::optional<rctf> get_win_rect(const ARegion &region) const;
 
   void begin_renaming();
-  void toggle_collapsed();
-  void set_collapsed(bool collapsed);
+
+  /**
+   * Toggle the expanded/collapsed state.
+   *
+   * \note this does not call #on_collapse_change().
+   * \returns true when the collapsed state was changed, false otherwise.
+   */
+  bool toggle_collapsed();
+  /**
+   * Expand or collapse this tree view item.
+   *
+   * \note this does not call #on_collapse_change().
+   * \returns true when the collapsed state was changed, false otherwise.
+   */
+  virtual bool set_collapsed(bool collapsed);
+
   /**
    * Requires the tree to have completed reconstruction, see #is_reconstructed(). Otherwise we
    * can't be sure about the item state.
@@ -207,6 +221,19 @@ class AbstractTreeViewItem : public AbstractViewItem, public TreeViewItemContain
   bool is_collapsed() const;
   bool is_collapsible() const;
 
+  /**
+   * Called when the view changes an item's state from expanded to collapsed, or vice versa. Will
+   * only be called if the state change is triggered through the view, not through external
+   * changes. E.g. a click on an item calls it, a change in the value returned by
+   * #should_be_collapsed() to reflect an external state change does not.
+   */
+  virtual void on_collapse_change(bContext &C, bool is_collapsed);
+  /**
+   * If the result is not empty, it controls whether the item should be collapsed or not, usually
+   * depending on the data that the view represents.
+   */
+  virtual std::optional<bool> should_be_collapsed() const;
+
  protected:
   /** See AbstractViewItem::get_rename_string(). */
   /* virtual */ StringRef get_rename_string() const override;
@@ -219,6 +246,13 @@ class AbstractTreeViewItem : public AbstractViewItem, public TreeViewItemContain
    */
   virtual bool supports_collapsing() const;
 
+  /**
+   * Toggle the collapsed/expanded state, and call on_collapse_change() if it changed.
+   */
+  void toggle_collapsed_from_view(bContext &C);
+
+  void change_state_delayed() override;
+
   /** See #AbstractViewItem::matches(). */
   /* virtual */ bool matches(const AbstractViewItem &other) const override;
 

source/blender/editors/interface/interface_template_bone_collection_tree.cc

@@ -38,6 +38,8 @@ class BoneCollectionTreeView : public AbstractTreeView {
   explicit BoneCollectionTreeView(bArmature &armature);
   void build_tree() override;
 
+  bool listen(const wmNotifier &notifier) const override;
+
  private:
   void build_tree_node_recursive(TreeViewItemContainer &parent, const int bcoll_index);
 
@@ -281,6 +283,37 @@ class BoneCollectionItem : public AbstractTreeViewItem {
     ED_undo_push(&const_cast<bContext &>(C), "Change Armature's Active Bone Collection");
   }
 
+  std::optional<bool> should_be_collapsed() const override
+  {
+    const bool is_collapsed = !bone_collection_.is_expanded();
+    return is_collapsed;
+  }
+
+  bool set_collapsed(const bool collapsed) override
+  {
+    if (!AbstractTreeViewItem::set_collapsed(collapsed)) {
+      return false;
+    }
+
+    /* Ensure that the flag in DNA is set. */
+    ANIM_armature_bonecoll_is_expanded_set(&bone_collection_, !collapsed);
+    return true;
+  }
+
+  void on_collapse_change(bContext &C, const bool is_collapsed) override
+  {
+    const bool is_expanded = !is_collapsed;
+
+    /* Let RNA handle the property change. This makes sure all the notifiers and DEG
+     * update calls are properly called. */
+    PointerRNA bcoll_ptr = RNA_pointer_create(
+        &armature_.id, &RNA_BoneCollection, &bone_collection_);
+    PropertyRNA *prop = RNA_struct_find_property(&bcoll_ptr, "is_expanded");
+
+    RNA_property_boolean_set(&bcoll_ptr, prop, is_expanded);
+    RNA_property_update(&C, &bcoll_ptr, prop);
+  }
+
   bool supports_renaming() const override
   {
     return ANIM_armature_bonecoll_is_editable(&armature_, &bone_collection_);
@@ -351,8 +384,6 @@ void BoneCollectionTreeView::build_tree_node_recursive(TreeViewItemContainer &pa
   const bool has_any_selected_bones = bcolls_with_selected_bones_.contains(bcoll);
   BoneCollectionItem &bcoll_tree_item = parent.add_tree_item<BoneCollectionItem>(
       armature_, bcoll_index, has_any_selected_bones);
-  bcoll_tree_item.set_collapsed(false);
-
   for (int child_index = bcoll->child_index; child_index < bcoll->child_index + bcoll->child_count;
        child_index++)
   {
@@ -360,6 +391,11 @@ void BoneCollectionTreeView::build_tree_node_recursive(TreeViewItemContainer &pa
   }
 }
 
+bool BoneCollectionTreeView::listen(const wmNotifier &notifier) const
+{
+  return notifier.data == ND_BONE_COLLECTION;
+}
+
 void BoneCollectionTreeView::build_bcolls_with_selected_bones()
 {
   bcolls_with_selected_bones_.clear();

source/blender/editors/interface/views/tree_view.cc

@@ -354,7 +354,7 @@ void AbstractTreeViewItem::collapse_chevron_click_fn(bContext *C,
   AbstractTreeViewItem *hovered_item = from_item_handle<AbstractTreeViewItem>(hovered_item_handle);
   BLI_assert(hovered_item != nullptr);
 
-  hovered_item->toggle_collapsed();
+  hovered_item->toggle_collapsed_from_view(*C);
   /* When collapsing an item with an active child, make this collapsed item active instead so the
    * active item stays visible. */
   if (hovered_item->has_active_child()) {
@@ -504,24 +504,63 @@ bool AbstractTreeViewItem::is_collapsed() const
   return is_collapsible() && !is_open_;
 }
 
-void AbstractTreeViewItem::toggle_collapsed()
+bool AbstractTreeViewItem::toggle_collapsed()
 {
-  is_open_ = !is_open_;
+  return set_collapsed(is_open_);
 }
 
-void AbstractTreeViewItem::set_collapsed(const bool collapsed)
+bool AbstractTreeViewItem::set_collapsed(const bool collapsed)
 {
+  if (!is_collapsible()) {
+    return false;
+  }
+  if (collapsed == is_collapsed()) {
+    return false;
+  }
+
   is_open_ = !collapsed;
+  return true;
 }
 
 bool AbstractTreeViewItem::is_collapsible() const
 {
+  BLI_assert_msg(get_tree_view().is_reconstructed(),
+                 "State can't be queried until reconstruction is completed");
   if (children_.is_empty()) {
     return false;
   }
   return this->supports_collapsing();
 }
 
+void AbstractTreeViewItem::on_collapse_change(bContext & /*C*/, const bool /*is_collapsed*/)
+{
+  /* Do nothing by default. */
+}
+
+std::optional<bool> AbstractTreeViewItem::should_be_collapsed() const
+{
+  return std::nullopt;
+}
+
+void AbstractTreeViewItem::toggle_collapsed_from_view(bContext &C)
+{
+  if (toggle_collapsed()) {
+    on_collapse_change(C, is_collapsed());
+  }
+}
+
+void AbstractTreeViewItem::change_state_delayed()
+{
+  AbstractViewItem::change_state_delayed();
+
+  const std::optional<bool> should_be_collapsed = this->should_be_collapsed();
+  if (should_be_collapsed.has_value()) {
+    /* This reflects an external state change and therefore shouldn't call #on_collapse_change().
+     */
+    set_collapsed(*should_be_collapsed);
+  }
+}
+
 void AbstractTreeViewItem::ensure_parents_uncollapsed()
 {
   for (AbstractTreeViewItem *parent = parent_; parent; parent = parent->parent_) {

source/blender/makesdna/DNA_armature_types.h

@@ -306,6 +306,12 @@ typedef struct BoneCollection {
    * Return whether this collection is marked as 'solo'.
    */
   bool is_solo() const;
+  /**
+   * Whether or not this bone collection is expanded in the tree view.
+   *
+   * This corresponds to the #BONE_COLLECTION_EXPANDED flag.
+   */
+  bool is_expanded() const;
 #endif
 } BoneCollection;
 
@@ -534,8 +540,10 @@ typedef enum eBoneCollection_Flag {
    * \see eArmature_Flag::ARM_BCOLL_SOLO_ACTIVE
    */
   BONE_COLLECTION_SOLO = (1 << 4),
+  
+  BONE_COLLECTION_EXPANDED = (1 << 5), /* Expanded in the tree view. */
 } eBoneCollection_Flag;
-ENUM_OPERATORS(eBoneCollection_Flag, BONE_COLLECTION_SOLO)
+ENUM_OPERATORS(eBoneCollection_Flag, BONE_COLLECTION_EXPANDED)
 
 #ifdef __cplusplus
 

source/blender/makesrna/intern/rna_armature.cc

@@ -399,6 +399,12 @@ static void rna_BoneCollection_is_solo_set(PointerRNA *ptr, const bool is_solo)
   ANIM_armature_bonecoll_solo_set(arm, bcoll, is_solo);
 }
 
+static void rna_BoneCollection_is_expanded_set(PointerRNA *ptr, const bool is_expanded)
+{
+  BoneCollection *bcoll = (BoneCollection *)ptr->data;
+  ANIM_armature_bonecoll_is_expanded_set(bcoll, is_expanded);
+}
+
 static std::optional<std::string> rna_BoneCollection_path(const PointerRNA *ptr)
 {
   const BoneCollection *bcoll = (const BoneCollection *)ptr->data;
@@ -2306,6 +2312,15 @@ static void rna_def_bonecollection(BlenderRNA *brna)
   RNA_def_property_string_funcs(prop, nullptr, nullptr, "rna_BoneCollection_name_set");
   RNA_def_property_update(prop, NC_OBJECT | ND_BONE_COLLECTION, nullptr);
 
+  prop = RNA_def_property(srna, "is_expanded", PROP_BOOLEAN, PROP_NONE);
+  RNA_def_property_boolean_sdna(prop, nullptr, "flags", BONE_COLLECTION_EXPANDED);
+  RNA_def_property_ui_text(
+      prop, "Expanded", "This bone collection is expanded in the bone collections tree view");
+  RNA_def_property_flag(prop, PROP_LIB_EXCEPTION);
+  RNA_def_property_override_flag(prop, PROPOVERRIDE_OVERRIDABLE_LIBRARY);
+  RNA_def_property_boolean_funcs(prop, nullptr, "rna_BoneCollection_is_expanded_set");
+  RNA_def_property_update(prop, NC_OBJECT | ND_BONE_COLLECTION, nullptr);
+
   prop = RNA_def_property(srna, "is_visible", PROP_BOOLEAN, PROP_NONE);
   RNA_def_property_boolean_sdna(prop, nullptr, "flags", BONE_COLLECTION_VISIBLE);
   RNA_def_property_ui_text(