From bdbccaaff826d0378e3d1af4d5c1c4197dfcc9ac Mon Sep 17 00:00:00 2001
From: Cole Bemis <colebemis@github.com>
Date: Thu, 10 Nov 2022 06:32:19 -0800
Subject: [PATCH] TreeView: Preserve expanded state of nested items (#2538)

* Add failing test

* Require id prop on tree items

* Create expanded cache

* Update expanded state cache

* Remove console.log

* Add comment about expanded cache

* Update comment

* Create sour-readers-rescue.md

* Update src/TreeView/TreeView.tsx

Co-authored-by: Josh Black <joshblack@github.com>

* Fix lint errors

Co-authored-by: Josh Black <joshblack@github.com>
---
 .changeset/sour-readers-rescue.md          |  16 ++
 docs/content/TreeView.mdx                  |  46 ++--
 src/TreeView/TreeView.features.stories.tsx |  84 ++++---
 src/TreeView/TreeView.stories.tsx          |   8 +-
 src/TreeView/TreeView.test.tsx             | 272 +++++++++++++--------
 src/TreeView/TreeView.tsx                  |  79 +++---
 src/hooks/useControllableState.ts          |   2 +-
 7 files changed, 310 insertions(+), 197 deletions(-)
 create mode 100644 .changeset/sour-readers-rescue.md

diff --git a/.changeset/sour-readers-rescue.md b/.changeset/sour-readers-rescue.md
new file mode 100644
index 00000000000..b70dcaf0f45
--- /dev/null
+++ b/.changeset/sour-readers-rescue.md
@@ -0,0 +1,16 @@
+---
+"@primer/react": patch
+---
+
+TreeView: Preserve expanded state of nested items when parent item is collapsed.
+
+**Breaking change**
+
+`TreeView.Item` and `TreeView.LinkItem` now require an `id` prop:
+
+```diff
+- <TreeView.Item>...</TreeView.Item>
++ <TreeView.Item id="unqiue-id">...</TreeView.Item>
+```
+
+This is not in a major release because TreeView is still a `draft` component.
diff --git a/docs/content/TreeView.mdx b/docs/content/TreeView.mdx
index 16dd99d9221..f32b9ab1782 100644
--- a/docs/content/TreeView.mdx
+++ b/docs/content/TreeView.mdx
@@ -13,13 +13,13 @@ description: A hierarchical list of items where nested items can be expanded and
 <Box sx={{maxWidth: 400}}>
   <nav aria-label="Files">
     <TreeView aria-label="Files">
-      <TreeView.Item>
+      <TreeView.Item id="src">
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#">
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -28,7 +28,7 @@ description: A hierarchical list of items where nested items can be expanded and
               <StyledOcticon icon={DiffAddedIcon} color="success.fg" aria-label="added" />
             </TreeView.TrailingVisual>
           </TreeView.LinkItem>
-          <TreeView.LinkItem href="#" current>
+          <TreeView.LinkItem id="src/Button.tsx" href="#" current>
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -39,7 +39,7 @@ description: A hierarchical list of items where nested items can be expanded and
           </TreeView.LinkItem>
         </TreeView.SubTree>
       </TreeView.Item>
-      <TreeView.LinkItem href="#">
+      <TreeView.LinkItem id="package.json" href="#">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
@@ -59,31 +59,31 @@ description: A hierarchical list of items where nested items can be expanded and
 <Box sx={{maxWidth: 400}}>
   <nav aria-label="Files">
     <TreeView aria-label="Files">
-      <TreeView.LinkItem href="#">
+      <TreeView.LinkItem id="src" href="#">
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#">
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
             Avatar.tsx
           </TreeView.LinkItem>
-          <TreeView.LinkItem href="#" current>
+          <TreeView.LinkItem id="src/Button" href="#" current>
             <TreeView.LeadingVisual>
               <TreeView.DirectoryIcon />
             </TreeView.LeadingVisual>
             Button
             <TreeView.SubTree>
-              <TreeView.LinkItem href="#">
+              <TreeView.LinkItem id="src/Button/Button.tsx" href="#">
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
                 Button.tsx
               </TreeView.LinkItem>
-              <TreeView.LinkItem href="#">
+              <TreeView.LinkItem id="src/Button/Button.test.tsx" href="#">
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
@@ -93,7 +93,7 @@ description: A hierarchical list of items where nested items can be expanded and
           </TreeView.LinkItem>
         </TreeView.SubTree>
       </TreeView.LinkItem>
-      <TreeView.LinkItem href="#">
+      <TreeView.LinkItem id="package.json" href="#">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
@@ -115,11 +115,13 @@ function ControlledTreeView() {
       <Button onClick={() => setExpanded(!expanded)}>{expanded ? 'Collapse' : 'Expand'}</Button>
       <nav aria-label="Files">
         <TreeView aria-label="Files">
-          <TreeView.Item expanded={expanded} onExpandedChange={setExpanded}>
+          <TreeView.Item id="src" expanded={expanded} onExpandedChange={setExpanded}>
             src
             <TreeView.SubTree>
-              <TreeView.LinkItem href="#">Avatar.tsx</TreeView.LinkItem>
-              <TreeView.LinkItem href="#" current>
+              <TreeView.LinkItem id="src/Avatar.tsx" href="#">
+                Avatar.tsx
+              </TreeView.LinkItem>
+              <TreeView.LinkItem id="src/Button.tsx" href="#" current>
                 Button.tsx
               </TreeView.LinkItem>
             </TreeView.SubTree>
@@ -141,14 +143,16 @@ To render stateful visuals, pass a render function to `TreeView.LeadingVisual` o
 <Box sx={{maxWidth: 400}}>
   <nav aria-label="Files">
     <TreeView aria-label="Files">
-      <TreeView.Item>
+      <TreeView.Item id="src">
         <TreeView.LeadingVisual>
           {({isExpanded}) => (isExpanded ? <FileDirectoryOpenFillIcon /> : <FileDirectoryFillIcon />)}
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#">Avatar.tsx</TreeView.LinkItem>
-          <TreeView.LinkItem href="#" current>
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#">
+            Avatar.tsx
+          </TreeView.LinkItem>
+          <TreeView.LinkItem id="src/Button.tsx" href="#" current>
             Button.tsx
           </TreeView.LinkItem>
         </TreeView.SubTree>
@@ -164,14 +168,16 @@ Since stateful directory icons are a common use case for TreeView, we provide a
 <Box sx={{maxWidth: 400}}>
   <nav aria-label="Files">
     <TreeView aria-label="Files">
-      <TreeView.Item>
+      <TreeView.Item id="src">
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#">Avatar.tsx</TreeView.LinkItem>
-          <TreeView.LinkItem href="#" current>
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#">
+            Avatar.tsx
+          </TreeView.LinkItem>
+          <TreeView.LinkItem id="src/Button.tsx" href="#" current>
             Button.tsx
           </TreeView.LinkItem>
         </TreeView.SubTree>
@@ -197,6 +203,7 @@ See [Storybook](https://primer.style/react/storybook?path=/story/components-tree
 ### TreeView.Item
 
 <PropsTable>
+  <PropsTableRow name="id" type="string" required description="Unique identifier for the item." />
   <PropsTableRow name="children" type="React.ReactNode" required />
   <PropsTableRow
     name="current"
@@ -230,6 +237,7 @@ See [Storybook](https://primer.style/react/storybook?path=/story/components-tree
 ### TreeView.LinkItem
 
 <PropsTable>
+  <PropsTableRow name="id" type="string" required description="Unique identifier for the item." />
   <PropsTableRow name="children" type="React.ReactNode" required />
   <PropsTableRow
     name="href"
diff --git a/src/TreeView/TreeView.features.stories.tsx b/src/TreeView/TreeView.features.stories.tsx
index 9601121d14a..d55373d40fb 100644
--- a/src/TreeView/TreeView.features.stories.tsx
+++ b/src/TreeView/TreeView.features.stories.tsx
@@ -23,31 +23,31 @@ const meta: Meta = {
 export const Files: Story = () => (
   <nav aria-label="Files">
     <TreeView aria-label="Files">
-      <TreeView.LinkItem href="#src" defaultExpanded>
+      <TreeView.LinkItem id="src" href="#src" defaultExpanded>
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#avatar-tsx">
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#avatar-tsx">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
             Avatar.tsx
           </TreeView.LinkItem>
-          <TreeView.LinkItem href="#button" current>
+          <TreeView.LinkItem id="src/Button" href="#button" current>
             <TreeView.LeadingVisual>
               <TreeView.DirectoryIcon />
             </TreeView.LeadingVisual>
             Button
             <TreeView.SubTree>
-              <TreeView.LinkItem href="#button-tsx">
+              <TreeView.LinkItem id="src/Button/Button.tsx" href="#button-tsx">
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
                 Button.tsx
               </TreeView.LinkItem>
-              <TreeView.LinkItem href="#button-test-tsx">
+              <TreeView.LinkItem id="src/Button/Button.test.tsx" href="#button-test-tsx">
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
@@ -55,7 +55,7 @@ export const Files: Story = () => (
               </TreeView.LinkItem>
             </TreeView.SubTree>
           </TreeView.LinkItem>
-          <TreeView.Item>
+          <TreeView.Item id="src/ReallyLongFileNameThatShouldBeTruncated.tsx">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -64,6 +64,7 @@ export const Files: Story = () => (
         </TreeView.SubTree>
       </TreeView.LinkItem>
       <TreeView.LinkItem
+        id="public"
         href="#public"
         // eslint-disable-next-line no-console
         onExpandedChange={isExpanded => console.log(`${isExpanded ? 'Expanded' : 'Collapsed'} "public" folder `)}
@@ -73,13 +74,13 @@ export const Files: Story = () => (
         </TreeView.LeadingVisual>
         public
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#index-html">
+          <TreeView.LinkItem id="public/index.html" href="#index-html">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
             index.html
           </TreeView.LinkItem>
-          <TreeView.LinkItem href="#favicon-ico">
+          <TreeView.LinkItem id="public/favicon.ico" href="#favicon-ico">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -87,7 +88,7 @@ export const Files: Story = () => (
           </TreeView.LinkItem>
         </TreeView.SubTree>
       </TreeView.LinkItem>
-      <TreeView.LinkItem href="#package-json">
+      <TreeView.LinkItem id="package.json" href="#package-json">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
@@ -101,13 +102,13 @@ export const FilesChanged: Story = () => {
   return (
     <nav aria-label="Files">
       <TreeView aria-label="Files">
-        <TreeView.Item defaultExpanded>
+        <TreeView.Item id="src" defaultExpanded>
           <TreeView.LeadingVisual>
             <TreeView.DirectoryIcon />
           </TreeView.LeadingVisual>
           src
           <TreeView.SubTree>
-            <TreeView.LinkItem href="#avatar-tsx">
+            <TreeView.LinkItem id="src/Avatar.tsx" href="#avatar-tsx">
               <TreeView.LeadingVisual>
                 <FileIcon />
               </TreeView.LeadingVisual>
@@ -116,13 +117,13 @@ export const FilesChanged: Story = () => {
                 <StyledOcticon icon={DiffAddedIcon} color="success.fg" />
               </TreeView.TrailingVisual>
             </TreeView.LinkItem>
-            <TreeView.Item defaultExpanded>
+            <TreeView.Item id="src/Button" defaultExpanded>
               <TreeView.LeadingVisual>
                 <TreeView.DirectoryIcon />
               </TreeView.LeadingVisual>
               Button
               <TreeView.SubTree>
-                <TreeView.LinkItem href="#button-tsx" current>
+                <TreeView.LinkItem id="src/Button/Button.tsx" href="#button-tsx" current>
                   <TreeView.LeadingVisual>
                     <FileIcon />
                   </TreeView.LeadingVisual>
@@ -131,7 +132,7 @@ export const FilesChanged: Story = () => {
                     <StyledOcticon icon={DiffModifiedIcon} color="attention.fg" />
                   </TreeView.TrailingVisual>
                 </TreeView.LinkItem>
-                <TreeView.LinkItem href="#button-test-tsx">
+                <TreeView.LinkItem id="src/Button/Button.test.tsx" href="#button-test-tsx">
                   <TreeView.LeadingVisual>
                     <FileIcon />
                   </TreeView.LeadingVisual>
@@ -142,7 +143,7 @@ export const FilesChanged: Story = () => {
                 </TreeView.LinkItem>
               </TreeView.SubTree>
             </TreeView.Item>
-            <TreeView.Item>
+            <TreeView.Item id="src/ReallyLongFileNameThatShouldBeTruncated.tsx">
               <TreeView.LeadingVisual>
                 <FileIcon />
               </TreeView.LeadingVisual>
@@ -153,13 +154,13 @@ export const FilesChanged: Story = () => {
             </TreeView.Item>
           </TreeView.SubTree>
         </TreeView.Item>
-        <TreeView.Item defaultExpanded>
+        <TreeView.Item id="public" defaultExpanded>
           <TreeView.LeadingVisual>
             <TreeView.DirectoryIcon />
           </TreeView.LeadingVisual>
           public
           <TreeView.SubTree>
-            <TreeView.LinkItem href="#index-html">
+            <TreeView.LinkItem id="public/index.html" href="#index-html">
               <TreeView.LeadingVisual>
                 <FileIcon />
               </TreeView.LeadingVisual>
@@ -168,7 +169,7 @@ export const FilesChanged: Story = () => {
                 <StyledOcticon icon={DiffRenamedIcon} />
               </TreeView.TrailingVisual>
             </TreeView.LinkItem>
-            <TreeView.LinkItem href="#favicon-ico">
+            <TreeView.LinkItem id="public/favicon.ico" href="#favicon-ico">
               <TreeView.LeadingVisual>
                 <FileIcon />
               </TreeView.LeadingVisual>
@@ -301,6 +302,7 @@ function TreeItem({
   const {currentPath, setCurrentPath} = React.useContext(CurrentPathContext)
   return (
     <TreeView.Item
+      id={path.join('/')}
       current={currentPath.join('/') === path.join('/')}
       expanded={item.data.expanded}
       onExpandedChange={expanded => onExpandedChange(path, expanded)}
@@ -347,13 +349,14 @@ export const AsyncSuccess: Story = args => {
   return (
     <nav aria-label="Files">
       <TreeView aria-label="Files">
-        <TreeView.Item>
+        <TreeView.Item id="file-1">
           <TreeView.LeadingVisual>
             <FileIcon />
           </TreeView.LeadingVisual>
           Some file
         </TreeView.Item>
         <TreeView.Item
+          id="async-directory"
           onExpandedChange={async isExpanded => {
             if (asyncItems.length === 0 && isExpanded) {
               setIsLoading(true)
@@ -372,7 +375,7 @@ export const AsyncSuccess: Story = args => {
           Directory with async items
           <TreeView.SubTree state={state}>
             {asyncItems.map(item => (
-              <TreeView.Item key={item}>
+              <TreeView.Item id={`item-${item}`} key={item}>
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
@@ -381,7 +384,7 @@ export const AsyncSuccess: Story = args => {
             ))}
           </TreeView.SubTree>
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="another-file">
           <TreeView.LeadingVisual>
             <FileIcon />
           </TreeView.LeadingVisual>
@@ -411,13 +414,14 @@ export const AsyncWithCount: Story = args => {
   return (
     <nav aria-label="Files">
       <TreeView aria-label="Files">
-        <TreeView.Item>
+        <TreeView.Item id="some-file">
           <TreeView.LeadingVisual>
             <FileIcon />
           </TreeView.LeadingVisual>
           Some file
         </TreeView.Item>
         <TreeView.Item
+          id="async-directory"
           onExpandedChange={async isExpanded => {
             if (asyncItems.length === 0 && isExpanded) {
               setIsLoading(true)
@@ -436,7 +440,7 @@ export const AsyncWithCount: Story = args => {
           Directory with async items
           <TreeView.SubTree state={state} count={args.count}>
             {asyncItems.map(item => (
-              <TreeView.Item key={item}>
+              <TreeView.Item key={item} id={`item-${item}`}>
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
@@ -445,7 +449,7 @@ export const AsyncWithCount: Story = args => {
             ))}
           </TreeView.SubTree>
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="another-file">
           <TreeView.LeadingVisual>
             <FileIcon />
           </TreeView.LeadingVisual>
@@ -506,13 +510,14 @@ export const AsyncError: Story = args => {
 
   return (
     <TreeView aria-label="Files">
-      <TreeView.Item>
+      <TreeView.Item id="some-file">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
         Some file
       </TreeView.Item>
       <TreeView.Item
+        id="async-directory"
         onExpandedChange={isExpanded => {
           if (isExpanded) {
             loadItems()
@@ -538,7 +543,7 @@ export const AsyncError: Story = args => {
             </TreeView.ErrorDialog>
           ) : null}
           {asyncItems.map(item => (
-            <TreeView.Item key={item}>
+            <TreeView.Item key={item} id={`item-${item}`}>
               <TreeView.LeadingVisual>
                 <FileIcon />
               </TreeView.LeadingVisual>
@@ -547,7 +552,7 @@ export const AsyncError: Story = args => {
           ))}
         </TreeView.SubTree>
       </TreeView.Item>
-      <TreeView.Item>
+      <TreeView.Item id="another-file">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
@@ -577,6 +582,7 @@ export const EmptyDirectories: Story = () => {
   return (
     <TreeView aria-label="Files">
       <TreeView.Item
+        id="src"
         onExpandedChange={expanded => {
           if (expanded) {
             timeoutId.current = setTimeout(() => {
@@ -592,7 +598,7 @@ export const EmptyDirectories: Story = () => {
         src
         <TreeView.SubTree state={state} />
       </TreeView.Item>
-      <TreeView.Item>
+      <TreeView.Item id=".github">
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
@@ -608,15 +614,19 @@ export const NestedScrollContainer: Story = () => {
     <Box sx={{maxHeight: '50vh', overflow: 'auto'}}>
       <TreeView aria-label="Files">
         {Array.from({length: 100}).map((_, i) => (
-          <TreeView.Item key={i}>
+          <TreeView.Item key={i} id={`directory-${i}`}>
             <TreeView.LeadingVisual>
               <TreeView.DirectoryIcon />
             </TreeView.LeadingVisual>
             Directory {i}
             <TreeView.SubTree>
               {Array.from({length: 10}).map((_, j) => (
-                // eslint-disable-next-line no-console
-                <TreeView.Item key={j} onSelect={() => console.log(`Directory ${i}/File ${j}`)}>
+                <TreeView.Item
+                  key={j}
+                  id={`directory-${i}/file-${j}`}
+                  // eslint-disable-next-line no-console
+                  onSelect={() => console.log(`Directory ${i}/File ${j}`)}
+                >
                   <TreeView.LeadingVisual>
                     <FileIcon />
                   </TreeView.LeadingVisual>
@@ -634,19 +644,19 @@ export const NestedScrollContainer: Story = () => {
 export const StressTest: Story = () => {
   return (
     <TreeView aria-label="Files">
-      {Array.from({length: 1000}).map((_, index) => (
-        <TreeView.Item key={index}>
+      {Array.from({length: 1000}).map((_, i) => (
+        <TreeView.Item key={i} id={`directory-${i}`}>
           <TreeView.LeadingVisual>
             <TreeView.DirectoryIcon />
           </TreeView.LeadingVisual>
-          Directory {index}
+          Directory {i}
           <TreeView.SubTree>
-            {Array.from({length: 100}).map((_, index) => (
-              <TreeView.Item key={index}>
+            {Array.from({length: 100}).map((_, j) => (
+              <TreeView.Item key={i} id={`directory-${i}/file-${j}`}>
                 <TreeView.LeadingVisual>
                   <FileIcon />
                 </TreeView.LeadingVisual>
-                File {index}
+                File {j}
               </TreeView.Item>
             ))}
           </TreeView.SubTree>
diff --git a/src/TreeView/TreeView.stories.tsx b/src/TreeView/TreeView.stories.tsx
index 3e342a11365..1dccded87a6 100644
--- a/src/TreeView/TreeView.stories.tsx
+++ b/src/TreeView/TreeView.stories.tsx
@@ -22,13 +22,13 @@ const meta: Meta = {
 export const Default: Story = () => (
   <nav aria-label="Files changed">
     <TreeView aria-label="Files changed">
-      <TreeView.Item defaultExpanded>
+      <TreeView.Item id="src" defaultExpanded>
         <TreeView.LeadingVisual>
           <TreeView.DirectoryIcon />
         </TreeView.LeadingVisual>
         src
         <TreeView.SubTree>
-          <TreeView.LinkItem href="#">
+          <TreeView.LinkItem id="src/Avatar.tsx" href="#">
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -37,7 +37,7 @@ export const Default: Story = () => (
               <StyledOcticon icon={DiffAddedIcon} color="success.fg" aria-label="added" />
             </TreeView.TrailingVisual>
           </TreeView.LinkItem>
-          <TreeView.LinkItem href="#" current>
+          <TreeView.LinkItem id="src/Button.tsx" href="#" current>
             <TreeView.LeadingVisual>
               <FileIcon />
             </TreeView.LeadingVisual>
@@ -48,7 +48,7 @@ export const Default: Story = () => (
           </TreeView.LinkItem>
         </TreeView.SubTree>
       </TreeView.Item>
-      <TreeView.LinkItem href="#">
+      <TreeView.LinkItem id="package.json" href="#">
         <TreeView.LeadingVisual>
           <FileIcon />
         </TreeView.LeadingVisual>
diff --git a/src/TreeView/TreeView.test.tsx b/src/TreeView/TreeView.test.tsx
index aceb5c1303e..65ada1fe8c5 100644
--- a/src/TreeView/TreeView.test.tsx
+++ b/src/TreeView/TreeView.test.tsx
@@ -22,9 +22,9 @@ describe('Markup', () => {
   it('uses tree role', () => {
     const {queryByRole} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>Item 1</TreeView.Item>
-        <TreeView.Item>Item 2</TreeView.Item>
-        <TreeView.Item>Item 3</TreeView.Item>
+        <TreeView.Item id="item-1">Item 1</TreeView.Item>
+        <TreeView.Item id="item-2">Item 2</TreeView.Item>
+        <TreeView.Item id="item-3">Item 3</TreeView.Item>
       </TreeView>
     )
 
@@ -36,9 +36,9 @@ describe('Markup', () => {
   it('uses treeitem role', () => {
     const {queryAllByRole} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>Item 1</TreeView.Item>
-        <TreeView.Item>Item 2</TreeView.Item>
-        <TreeView.Item>Item 3</TreeView.Item>
+        <TreeView.Item id="item-1">Item 1</TreeView.Item>
+        <TreeView.Item id="item-2">Item 2</TreeView.Item>
+        <TreeView.Item id="item-3">Item 3</TreeView.Item>
       </TreeView>
     )
 
@@ -50,10 +50,10 @@ describe('Markup', () => {
   it('hides subtrees by default', () => {
     const {queryByRole} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>
+        <TreeView.Item id="parent">
           Parent
           <TreeView.SubTree>
-            <TreeView.Item>Child</TreeView.Item>
+            <TreeView.Item id="child">Child</TreeView.Item>
           </TreeView.SubTree>
         </TreeView.Item>
       </TreeView>
@@ -69,9 +69,11 @@ describe('Markup', () => {
   it('uses aria-current', () => {
     const {getByRole} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>Item 1</TreeView.Item>
-        <TreeView.Item current>Item 2</TreeView.Item>
-        <TreeView.Item>Item 3</TreeView.Item>
+        <TreeView.Item id="item-1">Item 1</TreeView.Item>
+        <TreeView.Item id="item-2" current>
+          Item 2
+        </TreeView.Item>
+        <TreeView.Item id="item-3">Item 3</TreeView.Item>
       </TreeView>
     )
 
@@ -83,13 +85,13 @@ describe('Markup', () => {
   it('should be described by leading visuals', () => {
     const {getByLabelText} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>
+        <TreeView.Item id="item-1">
           <TreeView.LeadingVisual label="leading">
             <svg aria-hidden={true} />
           </TreeView.LeadingVisual>
           Item 1
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="item-2">
           <TreeView.LeadingVisual>
             <svg aria-hidden={true} />
           </TreeView.LeadingVisual>
@@ -107,13 +109,13 @@ describe('Markup', () => {
   it('should be described by trailing visuals', () => {
     const {getByLabelText} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>
+        <TreeView.Item id="item-1">
           Item 1
           <TreeView.TrailingVisual label="trailing">
             <svg aria-hidden={true} />
           </TreeView.TrailingVisual>
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="item-2">
           Item 2
           <TreeView.TrailingVisual>
             <svg aria-hidden={true} />
@@ -131,7 +133,7 @@ describe('Markup', () => {
   it('should be described by leading and trailing visuals', () => {
     const {getByLabelText} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>
+        <TreeView.Item id="item-1">
           <TreeView.LeadingVisual label="leading">
             <svg aria-hidden={true} />
           </TreeView.LeadingVisual>
@@ -140,7 +142,7 @@ describe('Markup', () => {
             <svg aria-hidden={true} />
           </TreeView.TrailingVisual>
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="item-2">
           <TreeView.LeadingVisual>
             <svg aria-hidden={true} />
           </TreeView.LeadingVisual>
@@ -169,15 +171,15 @@ describe('Markup', () => {
     })
     const {getByLabelText, getByText} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item>
+        <TreeView.Item id="item-1">
           Item 1
           <TreeView.SubTree>
-            <TreeView.Item>Item 1.a</TreeView.Item>
-            <TreeView.Item>Item 1.b</TreeView.Item>
-            <TreeView.Item>Item 1.c</TreeView.Item>
+            <TreeView.Item id="item-1-a">Item 1.a</TreeView.Item>
+            <TreeView.Item id="item-1-b">Item 1.b</TreeView.Item>
+            <TreeView.Item id="item-1-c">Item 1.c</TreeView.Item>
           </TreeView.SubTree>
         </TreeView.Item>
-        <TreeView.Item>
+        <TreeView.Item id="item-2">
           Item 2
           <TreeView.SubTree />
         </TreeView.Item>
@@ -203,20 +205,20 @@ describe('Keyboard interactions', () => {
     it('moves focus to the next visible treeitem', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="item-1" defaultExpanded>
             Item 1
             <TreeView.SubTree>
-              <TreeView.Item>Item 1.1</TreeView.Item>
+              <TreeView.Item id="item-1-1">Item 1.1</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="item-2">
             Item 2
             <TreeView.SubTree>
-              <TreeView.Item>Item 2.1</TreeView.Item>
-              <TreeView.Item>Item 2.2</TreeView.Item>
+              <TreeView.Item id="item-2-1">Item 2.1</TreeView.Item>
+              <TreeView.Item id="item-2-2">Item 2.2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>Item 3</TreeView.Item>
+          <TreeView.Item id="item-3">Item 3</TreeView.Item>
         </TreeView>
       )
 
@@ -258,20 +260,20 @@ describe('Keyboard interactions', () => {
     it('moves focus to the previous visible treeitem', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="item-1" defaultExpanded>
             Item 1
             <TreeView.SubTree>
-              <TreeView.Item>Item 1.1</TreeView.Item>
+              <TreeView.Item id="item-1-1">Item 1.1</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="item-2">
             Item 2
             <TreeView.SubTree>
-              <TreeView.Item>Item 2.1</TreeView.Item>
-              <TreeView.Item>Item 2.2</TreeView.Item>
+              <TreeView.Item id="item-2-1">Item 2.1</TreeView.Item>
+              <TreeView.Item id="item-2-2">Item 2.2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>Item 3</TreeView.Item>
+          <TreeView.Item id="item-3">Item 3</TreeView.Item>
         </TreeView>
       )
 
@@ -322,10 +324,10 @@ describe('Keyboard interactions', () => {
     it('collapses an expanded item', () => {
       const {getByRole, queryByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -361,10 +363,10 @@ describe('Keyboard interactions', () => {
     it('does nothing on a root-level collapsed item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="parent">
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -391,7 +393,7 @@ describe('Keyboard interactions', () => {
     it('does nothing on a root-level end item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>Item</TreeView.Item>
+          <TreeView.Item id="item">Item</TreeView.Item>
         </TreeView>
       )
 
@@ -410,11 +412,11 @@ describe('Keyboard interactions', () => {
     it('moves focus to parent of end item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child 1</TreeView.Item>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -443,14 +445,14 @@ describe('Keyboard interactions', () => {
     it('moves focus to parent of collapsed item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
-              <TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
+              <TreeView.Item id="nested-parent">
                 Nested parent
                 <TreeView.SubTree>
-                  <TreeView.Item>Nested child</TreeView.Item>
+                  <TreeView.Item id="nested-child">Nested child</TreeView.Item>
                 </TreeView.SubTree>
               </TreeView.Item>
             </TreeView.SubTree>
@@ -483,10 +485,10 @@ describe('Keyboard interactions', () => {
     it('expands a collapsed item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="parent">
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -518,10 +520,10 @@ describe('Keyboard interactions', () => {
     it('moves focus to first child of an expanded item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -550,11 +552,11 @@ describe('Keyboard interactions', () => {
     it('does nothing on an end item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child 1</TreeView.Item>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -584,10 +586,10 @@ describe('Keyboard interactions', () => {
     it('should move focus to the parent item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -606,10 +608,10 @@ describe('Keyboard interactions', () => {
     it('should not collapse an expanded item', () => {
       const {getByRole, queryByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -642,22 +644,22 @@ describe('Keyboard interactions', () => {
     it('moves focus to first visible item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="parent-1">
             Parent 1
             <TreeView.SubTree>
-              <TreeView.Item>Child 1</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="parent-2">
             Parent 2
             <TreeView.SubTree>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="parent-3">
             Parent 3
             <TreeView.SubTree>
-              <TreeView.Item>Child 3</TreeView.Item>
+              <TreeView.Item id="child-3">Child 3</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -687,22 +689,22 @@ describe('Keyboard interactions', () => {
     it('moves focus to last visible item', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="parent-1">
             Parent 1
             <TreeView.SubTree>
-              <TreeView.Item>Child 1</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="parent-2">
             Parent 2
             <TreeView.SubTree>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>
+          <TreeView.Item id="parent-3">
             Parent 3
             <TreeView.SubTree>
-              <TreeView.Item>Child 3</TreeView.Item>
+              <TreeView.Item id="child-3">Child 3</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -738,7 +740,9 @@ describe('Keyboard interactions', () => {
       const onSelect = jest.fn()
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item onSelect={onSelect}>Item</TreeView.Item>
+          <TreeView.Item id="item" onSelect={onSelect}>
+            Item
+          </TreeView.Item>
         </TreeView>
       )
 
@@ -757,11 +761,11 @@ describe('Keyboard interactions', () => {
     it('toggles expanded state if no onSelect function is provided', () => {
       const {getByRole, queryByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="parent">
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child 1</TreeView.Item>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -805,7 +809,7 @@ describe('Keyboard interactions', () => {
       const onSelect = jest.fn()
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.LinkItem href="#" onSelect={onSelect}>
+          <TreeView.LinkItem id="item" href="#" onSelect={onSelect}>
             Item
           </TreeView.LinkItem>
         </TreeView>
@@ -831,15 +835,15 @@ describe('Keyboard interactions', () => {
     it('moves focus to the next item that matches the typed character', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>
+          <TreeView.Item id="apple">
             Apple
             <TreeView.SubTree>
-              <TreeView.Item>Cantalope</TreeView.Item>
+              <TreeView.Item id="cantalope">Cantalope</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
-          <TreeView.Item>Banana</TreeView.Item>
-          <TreeView.Item>Cherry</TreeView.Item>
-          <TreeView.Item>Cucumber</TreeView.Item>
+          <TreeView.Item id="banana">Banana</TreeView.Item>
+          <TreeView.Item id="cherry">Cherry</TreeView.Item>
+          <TreeView.Item id="cucumber">Cucumber</TreeView.Item>
         </TreeView>
       )
 
@@ -862,10 +866,10 @@ describe('Keyboard interactions', () => {
     it('does nothing if no items match the typed character', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>Apple</TreeView.Item>
-          <TreeView.Item>Banana</TreeView.Item>
-          <TreeView.Item>Cherry</TreeView.Item>
-          <TreeView.Item>Durian</TreeView.Item>
+          <TreeView.Item id="apple">Apple</TreeView.Item>
+          <TreeView.Item id="banana">Banana</TreeView.Item>
+          <TreeView.Item id="cherry">Cherry</TreeView.Item>
+          <TreeView.Item id="durian">Durian</TreeView.Item>
         </TreeView>
       )
 
@@ -887,11 +891,11 @@ describe('Keyboard interactions', () => {
     it('supports multiple typed characters', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>Apple</TreeView.Item>
-          <TreeView.Item>Banana</TreeView.Item>
-          <TreeView.Item>Cherry</TreeView.Item>
-          <TreeView.Item>Cantalope 1</TreeView.Item>
-          <TreeView.Item>Cantalope 2</TreeView.Item>
+          <TreeView.Item id="apple">Apple</TreeView.Item>
+          <TreeView.Item id="banana">Banana</TreeView.Item>
+          <TreeView.Item id="cherry">Cherry</TreeView.Item>
+          <TreeView.Item id="cantalope-1">Cantalope 1</TreeView.Item>
+          <TreeView.Item id="cantalope-2">Cantalope 2</TreeView.Item>
         </TreeView>
       )
 
@@ -913,9 +917,11 @@ describe('Keyboard interactions', () => {
     it('prioritizes items following the current aria-activedescendant', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>Cucumber</TreeView.Item>
-          <TreeView.Item current>Cherry</TreeView.Item>
-          <TreeView.Item>Cantalope</TreeView.Item>
+          <TreeView.Item id="cucumber">Cucumber</TreeView.Item>
+          <TreeView.Item id="cherry" current>
+            Cherry
+          </TreeView.Item>
+          <TreeView.Item id="cantalope">Cantalope</TreeView.Item>
         </TreeView>
       )
 
@@ -942,10 +948,12 @@ describe('Keyboard interactions', () => {
     it('wraps around to the beginning if no items match after the current aria-activedescendant', () => {
       const {getByRole} = renderWithTheme(
         <TreeView aria-label="Test tree">
-          <TreeView.Item>Cucumber</TreeView.Item>
-          <TreeView.Item>Cherry</TreeView.Item>
-          <TreeView.Item current>Cantalope</TreeView.Item>
-          <TreeView.Item>Apple</TreeView.Item>
+          <TreeView.Item id="cucumber">Cucumber</TreeView.Item>
+          <TreeView.Item id="cherry">Cherry</TreeView.Item>
+          <TreeView.Item id="cantalope" current>
+            Cantalope
+          </TreeView.Item>
+          <TreeView.Item id="apple">Apple</TreeView.Item>
         </TreeView>
       )
 
@@ -971,16 +979,67 @@ describe('Keyboard interactions', () => {
   })
 })
 
-describe('Controlled state', () => {
+describe('State', () => {
+  it('persists expanded state of nested items', () => {
+    const {getByRole} = renderWithTheme(
+      <TreeView aria-label="Test tree">
+        <TreeView.Item id="item-1" defaultExpanded>
+          Item 1
+          <TreeView.SubTree>
+            <TreeView.Item id="item-2">
+              Item 2
+              <TreeView.SubTree>
+                <TreeView.Item id="item-3">Item 3</TreeView.Item>
+              </TreeView.SubTree>
+            </TreeView.Item>
+          </TreeView.SubTree>
+        </TreeView.Item>
+      </TreeView>
+    )
+
+    const item1 = getByRole('treeitem', {name: 'Item 1'})
+    const item2 = getByRole('treeitem', {name: 'Item 2'})
+
+    // Item 2 should be collapsed by default
+    expect(item2).toHaveAttribute('aria-expanded', 'false')
+
+    // Focus item 2
+    item2.focus()
+
+    // Press Enter to expand item 2
+    fireEvent.keyDown(document.activeElement || document.body, {key: 'Enter'})
+
+    // Item 2 should be expanded
+    expect(item2).toHaveAttribute('aria-expanded', 'true')
+
+    // Press ↑ to move focus to item 1
+    fireEvent.keyDown(document.activeElement || document.body, {key: 'ArrowUp'})
+
+    // Press Enter to collapse item 1
+    fireEvent.keyDown(document.activeElement || document.body, {key: 'Enter'})
+
+    // Item 1 should be collapsed
+    expect(item1).toHaveAttribute('aria-expanded', 'false')
+
+    // Press Enter again to expand item 1
+    fireEvent.keyDown(document.activeElement || document.body, {key: 'Enter'})
+
+    // Item 1 should be expanded
+    expect(item1).toHaveAttribute('aria-expanded', 'true')
+
+    // Item 2 should still be expanded
+    expect(getByRole('treeitem', {name: 'Item 2'})).toHaveAttribute('aria-expanded', 'true')
+  })
+
   it('can be controlled', () => {
     function TestTree() {
       const [expanded, setExpanded] = React.useState(true)
       return (
         <TreeView aria-label="Test tree">
-          <TreeView.Item expanded={expanded} onExpandedChange={setExpanded}>
+          <TreeView.Item id="parent" expanded={expanded} onExpandedChange={setExpanded}>
             Parent
             <TreeView.SubTree>
-              <TreeView.Item>Child</TreeView.Item>
+              <TreeView.Item id="child">Child</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -1018,10 +1077,10 @@ describe('Asyncronous loading', () => {
           {/* Mimic the completion of async loading by clicking the button */}
           <button onClick={() => setState('done')}>Done</button>
           <TreeView aria-label="Test tree">
-            <TreeView.Item defaultExpanded>
+            <TreeView.Item id="parent" defaultExpanded>
               Parent
               <TreeView.SubTree state={state}>
-                <TreeView.Item>Child</TreeView.Item>
+                <TreeView.Item id="child">Child</TreeView.Item>
               </TreeView.SubTree>
             </TreeView.Item>
           </TreeView>
@@ -1058,11 +1117,11 @@ describe('Asyncronous loading', () => {
 
       return (
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree state={state}>
-              <TreeView.Item>Child 1</TreeView.Item>
-              <TreeView.Item>Child 2</TreeView.Item>
+              <TreeView.Item id="child-1">Child 1</TreeView.Item>
+              <TreeView.Item id="child-2">Child 2</TreeView.Item>
             </TreeView.SubTree>
           </TreeView.Item>
         </TreeView>
@@ -1100,7 +1159,7 @@ describe('Asyncronous loading', () => {
 
       return (
         <TreeView aria-label="Test tree">
-          <TreeView.Item defaultExpanded>
+          <TreeView.Item id="parent" defaultExpanded>
             Parent
             <TreeView.SubTree>
               {error ? (
@@ -1145,11 +1204,11 @@ describe('Asyncronous loading', () => {
   it('ignores arrow keys when error dialog is open', async () => {
     const {getByRole} = renderWithTheme(
       <TreeView aria-label="Test tree">
-        <TreeView.Item defaultExpanded>
+        <TreeView.Item id="parent" defaultExpanded>
           Parent
           <TreeView.SubTree>
             <TreeView.ErrorDialog>Opps</TreeView.ErrorDialog>
-            <TreeView.Item>Child</TreeView.Item>
+            <TreeView.Item id="child">Child</TreeView.Item>
           </TreeView.SubTree>
         </TreeView.Item>
       </TreeView>
@@ -1193,6 +1252,7 @@ describe('Asyncronous loading', () => {
       return (
         <TreeView aria-label="Test tree">
           <TreeView.Item
+            id="item-1"
             onExpandedChange={expanded => {
               if (expanded) {
                 timeoutId.current = setTimeout(() => {
diff --git a/src/TreeView/TreeView.tsx b/src/TreeView/TreeView.tsx
index b5a4993a919..669f1d24364 100644
--- a/src/TreeView/TreeView.tsx
+++ b/src/TreeView/TreeView.tsx
@@ -27,8 +27,13 @@ import {useTypeahead} from './useTypeahead'
 
 const RootContext = React.createContext<{
   announceUpdate: (message: string) => void
+  // We cache the expanded state of tree items so we can preserve the state
+  // across remounts. This is necessary because we unmount tree items
+  // when their parent is collapsed.
+  expandedStateCache: React.RefObject<Map<string, boolean> | null>
 }>({
-  announceUpdate: () => {}
+  announceUpdate: () => {},
+  expandedStateCache: {current: new Map()}
 })
 
 const ItemContext = React.createContext<{
@@ -37,7 +42,7 @@ const ItemContext = React.createContext<{
   isSubTreeEmpty: boolean
   setIsSubTreeEmpty: React.Dispatch<React.SetStateAction<boolean>>
   isExpanded: boolean
-  setIsExpanded: React.Dispatch<React.SetStateAction<boolean>>
+  setIsExpanded: (isExpanded: boolean) => void
   leadingVisualId: string
   trailingVisualId: string
 }>({
@@ -257,8 +262,19 @@ const Root: React.FC<TreeViewProps> = ({'aria-label': ariaLabel, 'aria-labelledb
     }
   })
 
+  const expandedStateCache = React.useRef<Map<string, boolean> | null>(null)
+
+  if (expandedStateCache.current === null) {
+    expandedStateCache.current = new Map()
+  }
+
   return (
-    <RootContext.Provider value={{announceUpdate}}>
+    <RootContext.Provider
+      value={{
+        announceUpdate,
+        expandedStateCache
+      }}
+    >
       <>
         <VisuallyHidden role="status" aria-live="polite" aria-atomic="true">
           {ariaLiveMessage}
@@ -277,6 +293,7 @@ Root.displayName = 'TreeView'
 // TreeView.Item
 
 export type TreeViewItemProps = {
+  id: string
   children: React.ReactNode
   current?: boolean
   defaultExpanded?: boolean
@@ -288,14 +305,21 @@ export type TreeViewItemProps = {
 const {Slots, Slot} = createSlots(['LeadingVisual', 'TrailingVisual'])
 
 const Item = React.forwardRef<HTMLElement, TreeViewItemProps>(
-  ({current: isCurrentItem = false, defaultExpanded = false, expanded, onExpandedChange, onSelect, children}, ref) => {
-    const itemId = useSSRSafeId()
+  (
+    {id: itemId, current: isCurrentItem = false, defaultExpanded, expanded, onExpandedChange, onSelect, children},
+    ref
+  ) => {
+    const {expandedStateCache} = React.useContext(RootContext)
     const labelId = useSSRSafeId()
     const leadingVisualId = useSSRSafeId()
     const trailingVisualId = useSSRSafeId()
     const [isExpanded, setIsExpanded] = useControllableState({
       name: itemId,
-      defaultValue: defaultExpanded,
+      // If the item was previously mounted, it's expanded state might be cached.
+      // We check the cache first, and then fall back to the defaultExpanded prop.
+      // If defaultExpanded is not provided, we default to false unless the item
+      // is the current item, in which case we default to true.
+      defaultValue: () => expandedStateCache.current?.get(itemId) ?? defaultExpanded ?? isCurrentItem,
       value: expanded,
       onChange: onExpandedChange
     })
@@ -303,27 +327,22 @@ const Item = React.forwardRef<HTMLElement, TreeViewItemProps>(
     const {hasSubTree, subTree, childrenWithoutSubTree} = useSubTree(children)
     const [isSubTreeEmpty, setIsSubTreeEmpty] = React.useState(!hasSubTree)
 
+    // Set the expanded state and cache it
+    const setIsExpandedWithCache = React.useCallback(
+      (newIsExpanded: boolean) => {
+        setIsExpanded(newIsExpanded)
+        expandedStateCache.current?.set(itemId, newIsExpanded)
+      },
+      [itemId, setIsExpanded, expandedStateCache]
+    )
+
     // Expand or collapse the subtree
     const toggle = React.useCallback(
       (event?: React.MouseEvent | React.KeyboardEvent) => {
-        setIsExpanded(!isExpanded)
+        setIsExpandedWithCache(!isExpanded)
         event?.stopPropagation()
       },
-      // setIsExpanded is stable
-      // eslint-disable-next-line react-hooks/exhaustive-deps
-      [isExpanded]
-    )
-
-    // If this item is the current item, expand it
-    React.useLayoutEffect(
-      () => {
-        if (isCurrentItem) {
-          setIsExpanded(true)
-        }
-      },
-      // setIsExpanded is stable
-      // eslint-disable-next-line react-hooks/exhaustive-deps
-      [isCurrentItem]
+      [isExpanded, setIsExpandedWithCache]
     )
 
     const handleKeyDown = React.useCallback(
@@ -339,16 +358,16 @@ const Item = React.forwardRef<HTMLElement, TreeViewItemProps>(
           case 'ArrowRight':
             event.preventDefault()
             event.stopPropagation()
-            setIsExpanded(true)
+            setIsExpandedWithCache(true)
             break
           case 'ArrowLeft':
             event.preventDefault()
             event.stopPropagation()
-            setIsExpanded(false)
+            setIsExpandedWithCache(false)
             break
         }
       },
-      [onSelect, setIsExpanded, toggle]
+      [onSelect, setIsExpandedWithCache, toggle]
     )
 
     return (
@@ -359,7 +378,7 @@ const Item = React.forwardRef<HTMLElement, TreeViewItemProps>(
           isSubTreeEmpty,
           setIsSubTreeEmpty,
           isExpanded,
-          setIsExpanded,
+          setIsExpanded: setIsExpandedWithCache,
           leadingVisualId,
           trailingVisualId
         }}
@@ -659,12 +678,12 @@ type LoadingItemProps = {
   count?: number
 }
 
-const LoadingItem = React.forwardRef<HTMLElement, LoadingItemProps>((props, ref) => {
-  const {count} = props
+const LoadingItem = React.forwardRef<HTMLElement, LoadingItemProps>(({count}, ref) => {
+  const itemId = useSSRSafeId()
 
   if (count) {
     return (
-      <Item ref={ref}>
+      <Item id={itemId} ref={ref}>
         {Array.from({length: count}).map((_, i) => {
           return <SkeletonItem aria-hidden={true} key={i} />
         })}
@@ -674,7 +693,7 @@ const LoadingItem = React.forwardRef<HTMLElement, LoadingItemProps>((props, ref)
   }
 
   return (
-    <Item ref={ref}>
+    <Item id={itemId} ref={ref}>
       <LeadingVisual>
         <Spinner size="small" />
       </LeadingVisual>
diff --git a/src/hooks/useControllableState.ts b/src/hooks/useControllableState.ts
index 90efbee4fbc..dd4a00d2aa2 100644
--- a/src/hooks/useControllableState.ts
+++ b/src/hooks/useControllableState.ts
@@ -10,7 +10,7 @@ type ControllableStateOptions<T> = {
    * The default value used for the state. This will be
    * the fallback value used if `value` is not defined.
    * */
-  defaultValue: T
+  defaultValue: T | (() => T)
   /**
    * A controlled value. Omitting this means that the state is uncontrolled.
    */