Use of custom TrackByFunction in MatTree does not preserve expansion state when nodes change

[takamori]

What is the expected behavior?

Use of trackBy should also ensure that the expansion state does not get lost when nodes are updated.

What is the current behavior?

Use of trackBy does not ensure that the expansion state does not get lost when nodes are updated. This is presumably since TreeControl uses an expansionModel which is a SelectionModel where the SelectionModel has no awareness of the TrackByFunction.

What are the steps to reproduce?

Relying upon the TreeControl expansionModel:
https://stackblitz.com/edit/angular-5wpr9o-kjnmsn?file=app%2Ftree-flat-overview-example.ts

Which versions of Angular, Material, OS, TypeScript, browsers are affected?

all

Is there anything else we should know?

Trying to store expansion state in the nodes themselves as a workaround doesn't work (unless you transfer the state):
https://stackblitz.com/edit/angular-5wpr9o-ahptjs?file=app%2Ftree-flat-overview-example.ts

Having a custom expansionModel that is aware of the trackBy state works to preserve state, since the TreeControl's expansion-related methods apparently aren't actually used internally by the tree component:
https://stackblitz.com/edit/angular-5wpr9o-pbmk1g?file=app%2Ftree-flat-overview-example.ts

[crisbeto]

I spent some time looking into it, but at least with the way things are set up at the moment, there's no nice way to solve it. We can pass the trackBy function to the tree control and have it use the function when tracking which nodes are expanded, but the problem is that the tree control doesn't know the indexes of the items so it can't call the trackBy function correctly. We do have access to the array of all the data nodes from which we can figure out the index, but we'd have to do it very frequently (multiple times per change detection) which would defeat the purpose of a trackBy a bit.

[inorganik]

Any updates on this one? I'd love to add a trackBy on our app's tree and improve performance. It's a big tree and trackBy could give some big performance gains.

Using the expansion model mostly worked but fails intermittently after calling this.treeControl.collapseAll() or this.treeControl.expandAll(). (Despite looping through treeControl.dataNodes and adding all node IDs to the expansionModel, or calling expansionModel.clear()).

[dlmoffett]

Oh man, I so wish trackBy worked with the expansion state! Without this I'm stuck with paying the performance cost of re-expanding the tree on every change to the dataSource. For even moderately sized trees this can be expensive:

    this.treeControl.expansionModel.select(
      ...this.treeControl.dataNodes.filter((node) => node.expanded)
    );

Has there been any update to this? I see the OP posted a workaround; I will reference that and see if it works me.

[wgibbons1]

Unfortunately this problem means I can't use Mat-Tree. Fortunately, Angular makes it very easy to make your own tree out of a ul and a recursive component.

[beckerjohannes]

TLDR; when using trackBy, the expansionKey function must also be set, which is not documented very good.

I tried to use the mat-tree together with ngxs and had very strange effects. When data changes were fired and trackBy was used (like all nodes collapsed and were not expandable anymore). After some time banging my head around this, I remembered, that I stumbled into a "Redux Example" some time ago and thought, that should not work, if it is not working for me... the example is nicely hidden in the documentation of the cdkTree and there you can find:

trackBy = (index: number, node: TransformedData) => this.expansionKey(node);
  expansionKey = (node: TransformedData) => node.raw.id;

<cdk-tree
      #tree
      [dataSource]="roots"
      [childrenAccessor]="getChildren"
      [trackBy]="trackBy"
      [expansionKey]="expansionKey">
    ...

Problem solved.

Please add a hint to this in the trackby documentation of the mat-tree here:
https://material.angular.io/components/tree/overview#trackby and for the cdktree here:
https://material.angular.io/cdk/tree/overview#trackby

[passee]

@beckerjohannes this is awesome dude! thank you, it also fixed my problem.

it would be great if this gets documented in the overview right below the trackBy!

[passee]

@beckerjohannes actually i found a bug when using node.id in as the expansionKey the tree doesnt update at all, when you try to add data, but if i use the object it works somehow
example
Edit: this seems to cause an endless loop when using (expandedChange) and providing the data via a store and updating the expansion in the subscription
Update: it looks like i got it to work:

//works with nested tree
//works with flat tree
trackByFn = (index: number, node: FoodNode) => node;
expansionKeyFn = (node: FoodNode) => node.id;

//doesnt work with nested tree -> no update
//works with flat tree
// trackByFn = (index: number, node: FoodNode) => this.expansionKeyFn(node);
// expansionKeyFn = (node: FoodNode) => node.id;

//doesnt work with nested tree -> update but endless loop (trackBy fails)
//doesnt work with flat tree -> endless loop
// trackByFn = (index: number, node: FoodNode) => this.expansionKeyFn(node);
// expansionKeyFn = (node: FoodNode) => node;