Tree: How to keep opened states when tree updated

I have a hierarchical structure of data objects to display in a material tree. Each data object gets transformed into a TreeNodeModel containing the ID and the ParentID of the original data object. In the method that updates the data source the node's expanded state is saved/restored:

// Method that updates the data source
public updateDataSource(dataObjects: SomeDataObject) {
    
  // save node's expanded state
  const expandedNodesIds: string[] = [];
  if (this.treeControl.dataNodes) {
    this.treeControl.dataNodes.forEach((node: TreeNodeModel) => {
      if (this.treeControl.isExpandable(node) && this.treeControl.isExpanded(node)) {
        expandedNodesIds.push(node.id);
      }
    });
  }
    
  // update data source
  this.treeDataSource.data = dataObjects;
    
  // restore node's expanded state
  this.treeControl.dataNodes
    .filter(node => this.isActive(node) || expandedNodesIds.find(id => id === node.id))
    .forEach(nodeToExpand => {
      this.expandNode(nodeToExpand);
    });
}
    
// Method that expands the node (if not a leave) and its parent nodes (if any) using the TreeControl
private expandNode(treeNode: TreeNodeModel | undefined): void {
  if (!treeNode) {
    return;
  }
    
  if (this.treeControl.isExpandable(treeNode)) {
    this.treeControl.expand(treeNode);
  }
    
  const parentId = treeNode.parentId ? treeNode.parentId : undefined;
  this.expandNode(this.treeControl.dataNodes.find(node => node.id === parentId));
}

interface TreeNodeModel {
  expandable: boolean;
  level: number;
  id: string;
  parentId: string;
  // most probably some more custom data ;)
}

I added a boolean 'expanded' to my datamodel. I then use a function on (click) which inverts this, and a recursive loop to save that change to the actual data that is used for dataSource.data. So in reality I am not using the treecontrol anymore, even though I have still need it (the tree does not work without).

    <button mat-icon-button
    [attr.aria-label]="'toggle ' + node.name"
    (click)="changeState(node, myJson)"
    >
      <mat-icon class="mat-icon-rtl-mirror">
        {{node.expanded ? 'expand_more' : 'chevron_right'}}
      </mat-icon>
    </button>

--

  /** Changes expanded state for clicked tree-item, saves change to json data used by tree datasource */
  changeState(node, myJson) {
    node.expanded = !node.expanded;

    if (node.children && node.children.length > 0) {
      this.found = false;
      myJson.forEach(child => {
        if (!this.found) {
        this.saveStates(child, node);
        }
      });
    }
  }

  /** recursive loop-function used by this.changeState() to save tree-items expanded-state to the master array */
  saveStates(child, clickedChild) {
    if (child.id === clickedChild.id) {
      child.expanded = clickedChild.expanded;
      this.found = true;
      return;
    } else if (child.children && child.children.length > 0) {
      child.children.forEach(c => {
        this.saveStates(c, clickedGroup);
      });
    }
  }

-- And the standard functions from the tree-example I changed like this to work with my data:

  // checks if datasource for material tree has any children
  hasNestedChild = (_: number, nodeData: MyModel) => nodeData.children.length > 0;

  // returns children
  private _getChildren = (node: MyModel) => node.children;

this.treeControl.expansionModel.selected gives you your expanded nodes at the moment so you can save them this.prevExpansionModel = this.treeControl.expansionModel.selected and apply them,like this:

 this.prevExpansionModel.forEach(object => this.treeControl.expand(object));

Altough it seems to not be abke to close and open again after doing this so you have to figure that out but maybe this will get you closer to your goal