Repository: adobe/react-spectrum
File: packages/react-aria-components/src/Tree.tsx (TypeScript)

Code before:
  props: TreeProps<T>,
  collection: ICollection<unknown>,
  treeRef: RefObject<HTMLDivElement | null>
}

function TreeInner<T extends object>({props, collection, treeRef: ref}: TreeInnerProps<T>) {
  const {dragAndDropHooks} = props;
  let {direction} = useLocale();
  let collator = useCollator({usage: 'search', sensitivity: 'base'});
  let hasDragHooks = !!dragAndDropHooks?.useDraggableCollectionState;
  let hasDropHooks = !!dragAndDropHooks?.useDroppableCollectionState;
  let dragHooksProvided = useRef(hasDragHooks);
  let dropHooksProvided = useRef(hasDropHooks);

  useEffect(() => {
    if (dragHooksProvided.current !== hasDragHooks) {
      console.warn('Drag hooks were provided during one render, but not another. This should be avoided as it may produce unexpected behavior.');
    }
    if (dropHooksProvided.current !== hasDropHooks) {
      console.warn('Drop hooks were provided during one render, but not another. This should be avoided as it may produce unexpected behavior.');
    }
  }, [hasDragHooks, hasDropHooks]);
  let {
    selectionMode = 'none',
    expandedKeys: propExpandedKeys,
    defaultExpandedKeys: propDefaultExpandedKeys,
    onExpandedChange,
    disabledBehavior = 'all'
  } = props;
  let {CollectionRoot, isVirtualized, layoutDelegate,  dropTargetDelegate: ctxDropTargetDelegate} = useContext(CollectionRendererContext);

  // Kinda annoying that we have to replicate this code here as well as in useTreeState, but don't want to add
  // flattenCollection stuff to useTreeState. Think about this later
  let [expandedKeys, setExpandedKeys] = useControlledState(
    propExpandedKeys ? new Set(propExpandedKeys) : undefined,
    propDefaultExpandedKeys ? new Set(propDefaultExpandedKeys) : new Set(),
    onExpandedChange
  );

  let [lastCollection, setLastCollection] = useState(collection);
  let [lastExpandedKeys, setLastExpandedKeys] = useState(expandedKeys);
  let [flattenedCollection, setFlattenedCollection] = useState(() => new TreeCollection<object>({collection, lastExpandedKeys: new Set(), expandedKeys}));


  // if the lastExpandedKeys is not the same as the currentExpandedKeys or the collection has changed, then run this
  if (!areSetsEqual(lastExpandedKeys, expandedKeys) || collection !== lastCollection) {
    setFlattenedCollection(new TreeCollection<object>({collection, lastExpandedKeys, expandedKeys}));
    setLastCollection(collection);
    setLastExpandedKeys(expandedKeys);
  }


Review comment (bug):
I was initially concerned about this since now the "flattened" collection isn't actually flattened any more and thus things like `itemCount` always include rows that are in collapsed parents, but so far it doesn't seem like there are any issues (it helps that most of the collection size checks are if the collection is empty). 

The only change in behavior I've found so far is with DnD where starting a keyboard drag on the first child of a parent and pressing ArrowUp actually loops you back to the drop position after the last child of the parent. You need to then traverse up through all of the children before escaping the parent, but the previous behavior was also broken but in a different way so we'll have to revisit 

Code after:
      console.warn('Drag hooks were provided during one render, but not another. This should be avoided as it may produce unexpected behavior.');
    }
    if (dropHooksProvided.current !== hasDropHooks) {
      console.warn('Drop hooks were provided during one render, but not another. This should be avoided as it may produce unexpected behavior.');
    }
  }, [hasDragHooks, hasDropHooks]);
  let {
    selectionMode = 'none',
    expandedKeys: propExpandedKeys,
    defaultExpandedKeys: propDefaultExpandedKeys,
    onExpandedChange,
    disabledBehavior = 'all'
  } = props;
  let {CollectionRoot, isVirtualized, layoutDelegate,  dropTargetDelegate: ctxDropTargetDelegate} = useContext(CollectionRendererContext);

  // Kinda annoying that we have to replicate this code here as well as in useTreeState, but don't want to add
  // flattenCollection stuff to useTreeState. Think about this later
  let [expandedKeys, setExpandedKeys] = useControlledState(
    propExpandedKeys ? new Set(propExpandedKeys) : undefined,
    propDefaultExpandedKeys ? new Set(propDefaultExpandedKeys) : new Set(),
    onExpandedChange
  );

  let [lastCollection, setLastCollection] = useState(collection);
  let [lastExpandedKeys, setLastExpandedKeys] = useState(expandedKeys);
  let [flattenedCollection, setFlattenedCollection] = useState(() => collection.withExpandedKeys(lastExpandedKeys, expandedKeys));


  // if the lastExpandedKeys is not the same as the currentExpandedKeys or the collection has changed, then run this
  if (!areSetsEqual(lastExpandedKeys, expandedKeys) || collection !== lastCollection) {
    setFlattenedCollection(collection.withExpandedKeys(lastExpandedKeys, expandedKeys));
    setLastCollection(collection);
    setLastExpandedKeys(expandedKeys);
  }

  let state = useTreeState({
    ...props,
    selectionMode,
    expandedKeys,
    onExpandedChange: setExpandedKeys,
    collection: flattenedCollection,
    children: undefined,
    disabledBehavior
  });

  let {gridProps} = useTree({
    ...props,
    isVirtualized,
    layoutDelegate
  }, state, ref);