Merge pull request #26 from Lodin/fix/reset-height-after-recomputation

Fix resetting height after re-computation

Author: Lodin

__tests__/VariableSizeTree.spec.tsx

@@ -686,5 +686,21 @@ describe('VariableSizeTree', () => {
       );
       expect(foo3.height).toBe(100);
     });
+
+    it('properly resets heights after re-computation', async () => {
+      const listInstance: VariableSizeList = component
+        .find(VariableSizeList)
+        .instance() as VariableSizeList;
+
+      const resetAfterIndexSpy = jest.spyOn(listInstance, 'resetAfterIndex');
+      const foo1 = component.state('records')['foo-1']!;
+      const foo3 = component.state('records')['foo-3']!;
+
+      foo3.resize(100, true);
+      await foo1.toggle();
+
+      expect(resetAfterIndexSpy).toHaveBeenCalledWith(0, true);
+      expect(foo3.height).toBe(30);
+    });
   });
 });

src/FixedSizeTree.tsx

@@ -13,7 +13,7 @@ import {
   createRecord,
   shouldUpdateRecords,
   updateRecord,
-  updateRecordOnWalk,
+  updateRecordOnNewData,
 } from './utils';
 
 export type FixedSizeNodeData = NodeData;
@@ -50,7 +50,7 @@ const computeTree = createTreeComputer<
   createRecord,
   shouldUpdateRecords,
   updateRecord,
-  updateRecordOnWalk,
+  updateRecordOnNewData,
 });
 
 export class FixedSizeTree<T extends FixedSizeNodeData = NodeData> extends Tree<

src/Tree.tsx

@@ -136,7 +136,7 @@ export type TreeCreatorOptions<
     recordId: string | symbol,
     options: TUpdateOptions,
   ) => void;
-  updateRecordOnWalk: (record: TNodeRecord, options: TUpdateOptions) => void;
+  updateRecordOnNewData: (record: TNodeRecord, options: TUpdateOptions) => void;
 }>;
 
 export type TreeComputer<
@@ -173,7 +173,7 @@ export const createTreeComputer = <
   createRecord,
   shouldUpdateRecords,
   updateRecord,
-  updateRecordOnWalk,
+  updateRecordOnNewData,
 }: TreeCreatorOptions<
   TNodeComponentProps,
   TNodeRecord,
@@ -228,7 +228,7 @@ export const createTreeComputer = <
         records[id as string] = createRecord(value, state);
       } else {
         record.data = value;
-        updateRecordOnWalk(record, options);
+        updateRecordOnNewData(record, options);
       }
     }
 

src/VariableSizeTree.tsx

@@ -9,7 +9,11 @@ import Tree, {
   TreeState,
   UpdateOptions,
 } from './Tree';
-import {shouldUpdateRecords, updateRecord, updateRecordOnWalk} from './utils';
+import {
+  shouldUpdateRecords,
+  updateRecord,
+  updateRecordOnNewData,
+} from './utils';
 
 export type VariableSizeNodeData = Readonly<{
   /** Default node height. Can be used only with VariableSizeTree */
@@ -94,8 +98,8 @@ const computeTree = createTreeComputer<
     updateRecord(record, recordId, options);
   },
 
-  updateRecordOnWalk: (record, options) => {
-    updateRecordOnWalk(record, options);
+  updateRecordOnNewData: (record, options) => {
+    updateRecordOnNewData(record, options);
 
     if (options.useDefaultHeight) {
       record.height = record.data.defaultHeight;
@@ -132,6 +136,14 @@ export class VariableSizeTree<T extends VariableSizeNodeData> extends Tree<
     );
   }
 
+  public recomputeTree(options?: VariableSizeUpdateOptions): Promise<void> {
+    return super.recomputeTree(options).then(() => {
+      if (options?.useDefaultHeight) {
+        this.list.current?.resetAfterIndex(0, true);
+      }
+    });
+  }
+
   public render(): ReactNode {
     const {children, itemSize, rowComponent, treeWalker, ...rest} = this.props;
 

src/utils.ts

@@ -54,7 +54,7 @@ export const updateRecord: DefaultTreeCreatorOptions['updateRecord'] = (
     : opennessState?.[recordId as string] ?? record.isOpen;
 };
 
-export const updateRecordOnWalk: DefaultTreeCreatorOptions['updateRecordOnWalk'] = (
+export const updateRecordOnNewData: DefaultTreeCreatorOptions['updateRecordOnNewData'] = (
   record,
   {useDefaultOpenness = false},
 ) => {