Persist skeleton group expansion state

frontend/javascripts/oxalis/store.ts

@@ -160,6 +160,7 @@ export type Tree = {
 export type TreeGroupTypeFlat = {
   readonly name: string;
   readonly groupId: number;
+  readonly isExpanded?: boolean;
 };
 export type TreeGroup = TreeGroupTypeFlat & {
   readonly children: TreeGroup[];

frontend/javascripts/oxalis/view/right-border-tabs/skeleton_tab_view.tsx

@@ -346,10 +346,11 @@ class SkeletonTabView extends React.PureComponent<Props, State> {
   getTreeAndTreeGroupList = memoizeOne(
     (trees: TreeMap, treeGroups: Array<TreeGroup>, sortBy: string): Array<TreeOrTreeGroup> => {
       const groupToTreesMap = createGroupToTreesMap(trees);
-      const rootGroup = {
+      const rootGroup: TreeGroup = {
         name: "Root",
         groupId: MISSING_GROUP_ID,
         children: treeGroups,
+        isExpanded: true,
       };
 
       const makeTree = (tree: Tree) => ({

frontend/javascripts/oxalis/view/right-border-tabs/tree_hierarchy_view.tsx

@@ -42,7 +42,6 @@ import {
   GroupTypeEnum,
   deepFlatFilter,
   getNodeKey,
-  getNodeKeyFromNode,
   moveGroupsHelper,
   findParentGroupNode,
 } from "oxalis/view/right-border-tabs/tree_hierarchy_view_helpers";
@@ -69,8 +68,8 @@ import { formatNumberToLength, formatLengthAsVx } from "libs/format_utils";
 import { api, Store } from "oxalis/singletons";
 import { ChangeColorMenuItemContent } from "components/color_picker";
 import { HideTreeEdgesIcon } from "./hide_tree_eges_icon";
-import { useEffectOnlyOnce } from "libs/react_hooks";
 import { ColoredDotIcon } from "./segments_tab/segment_list_item";
+import { mapGroups } from "oxalis/model/accessors/skeletontracing_accessor";
 
 type Props = {
   activeTreeId: number | null | undefined;
@@ -106,25 +105,20 @@ function TreeHierarchyView(props: Props) {
       name: "Root",
       groupId: MISSING_GROUP_ID,
       children: props.treeGroups,
+      isExpanded: true,
     };
 
     const generatedGroupTree = insertTreesAndTransform([rootGroup], groupToTreesMap, props.sortBy);
     setUITreeData(generatedGroupTree);
   }, [props.trees, props.sortBy, props.treeGroups]);
 
-  useEffectOnlyOnce(() => {
-    // set default expanded keys
-    // the defaults should include the root node and at the active tree's group if applicable
-    let defaultExpandedKeys: React.Key[] = [getNodeKey(GroupTypeEnum.GROUP, MISSING_GROUP_ID)];
-    if (props.activeTreeId) {
-      const activeTreesGroupId = props.trees[props.activeTreeId].groupId;
-      if (activeTreesGroupId) {
-        defaultExpandedKeys.push(getNodeKey(GroupTypeEnum.GROUP, activeTreesGroupId));
-      }
-
-      setExpandedNodeKeys(defaultExpandedKeys);
-    }
-  });
+  useEffect(() => {
+    const expandedKeys = deepFlatFilter(
+      UITreeData,
+      (node) => node.type === GroupTypeEnum.GROUP && node.expanded,
+    ).map((node) => node.key);
+    setExpandedNodeKeys(expandedKeys);
+  }, [UITreeData]);
 
   useEffect(() => {
     // scroll to active tree if it changes
@@ -148,17 +142,25 @@ function TreeHierarchyView(props: Props) {
 
   const onExpand: TreeProps<TreeNode>["onExpand"] = (expandedKeys, info) => {
     const clickedNode = info.node;
+    const expandedKeySet = new Set(expandedKeys);
 
     if (clickedNode.type === GroupTypeEnum.GROUP && info.expanded === false) {
       // when collapsing a group, we need to collapse all its sub-gropus
       const subGroupKeys = deepFlatFilter(
         [clickedNode],
         (node) => node.type === GroupTypeEnum.GROUP,
       ).map((node) => node.key);
-      expandedKeys = _.without(expandedKeys, ...subGroupKeys);
+      subGroupKeys.forEach((key) => expandedKeySet.delete(key));
     }
-
-    setExpandedNodeKeys(expandedKeys);
+    const newGroups = mapGroups(props.treeGroups, (group) => {
+      const shouldBeExpanded = expandedKeySet.has(getNodeKey(GroupTypeEnum.GROUP, group.groupId));
+      if (shouldBeExpanded !== group.isExpanded) {
+        return { ...group, isExpanded: shouldBeExpanded };
+      } else {
+        return group;
+      }
+    });
+    setUpdateTreeGroups(newGroups);
   };
 
   const onCheck: TreeProps<TreeNode>["onCheck"] = (_checkedKeysValue, info) => {
@@ -238,17 +240,30 @@ function TreeHierarchyView(props: Props) {
     }
   }
 
-  function setExpansionOfAllSubgroupsTo(group: TreeNode, expanded: boolean) {
-    const selectedGroupIdKey = getNodeKeyFromNode(group);
-    const subGroupKeys = deepFlatFilter([group], (node) => node.type === GroupTypeEnum.GROUP).map(
-      (node) => node.key,
-    );
-
-    if (expanded) {
-      setExpandedNodeKeys(_.uniq([...expandedNodeKeys, ...subGroupKeys]));
-    } else {
-      setExpandedNodeKeys(_.without(expandedNodeKeys, selectedGroupIdKey, ...subGroupKeys));
+  function setExpansionOfAllSubgroupsTo(parentGroup: TreeNode, expanded: boolean) {
+    if (parentGroup.id === MISSING_GROUP_ID) {
+      const newGroups = mapGroups(props.treeGroups, (group) => {
+        if (group.isExpanded !== expanded) {
+          return { ...group, isExpanded: expanded };
+        }
+        return group;
+      });
+      setUpdateTreeGroups(newGroups);
+      return;
     }
+    const subGroups = getGroupByIdWithSubgroups(props.treeGroups, parentGroup.id);
+    const subGroupsMap = new Set(subGroups);
+    // If the subgroups will should be collapsed, not collapse the group itself.
+    // Do expand the group if the subgroups are expanded though.
+    if (expanded === false) subGroupsMap.delete(parentGroup.id);
+    const newGroups = mapGroups(props.treeGroups, (group) => {
+      if (subGroupsMap.has(group.groupId) && expanded !== group.isExpanded) {
+        return { ...group, isExpanded: expanded };
+      } else {
+        return group;
+      }
+    });
+    setUpdateTreeGroups(newGroups);
   }
 
   function onMoveWithContextAction(targetParentNode: TreeNode) {
@@ -302,7 +317,14 @@ function TreeHierarchyView(props: Props) {
     }
 
     // in either case expand the parent group
-    setExpandedNodeKeys([...expandedNodeKeys, getNodeKey(GroupTypeEnum.GROUP, parentGroupId)]);
+    const newGroups = mapGroups(props.treeGroups, (group) => {
+      if (group.groupId === parentGroupId) {
+        return { ...group, isExpanded: true };
+      } else {
+        return group;
+      }
+    });
+    setUpdateTreeGroups(newGroups);
   }
 
   function createGroup(groupId: number) {

frontend/javascripts/oxalis/view/right-border-tabs/tree_hierarchy_view_helpers.ts

@@ -90,6 +90,7 @@ export function insertTreesAndTransform(
       timestamp: 0,
       children: _.orderBy(treeNodeChildren, ["name"], ["asc"]),
       disableCheckbox: treeNodeChildren.length === 0,
+      expanded: group.isExpanded == null || group.isExpanded,
       isChecked: treeNodeChildren.every(
         // Groups that don't contain any trees should not influence the state of their parents
         (groupOrTree) => groupOrTree.isChecked || !groupOrTree.containsTrees,