From 6d98c110fda053d7d32098da5247fff50a83f7c9 Mon Sep 17 00:00:00 2001
From: Andrew Kushnir <akushnir@google.com>
Date: Wed, 7 Aug 2024 21:55:36 -0700
Subject: [PATCH] fix(core): complete post-hydration cleanup in components that
 use ViewContainerRef

Previously, if a component injects a `ViewContainerRef`, the post-hydration cleanup process doesn't visit inner views to cleanup dehydrated views in nested LContainers. This commit updates the logic to recognize this situation and enter host LView to complete cleanup.

Resolves #56989.
---
 packages/core/src/hydration/cleanup.ts        | 13 ++--
 .../platform-server/test/hydration_spec.ts    | 64 +++++++++++++++++++
 2 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/packages/core/src/hydration/cleanup.ts b/packages/core/src/hydration/cleanup.ts
index d0b0664ef4eef..c5b2ea8718ae7 100644
--- a/packages/core/src/hydration/cleanup.ts
+++ b/packages/core/src/hydration/cleanup.ts
@@ -76,6 +76,15 @@ function removeDehydratedView(dehydratedView: DehydratedContainerView, renderer:
  */
 function cleanupLContainer(lContainer: LContainer) {
   removeDehydratedViews(lContainer);
+
+  // The host could be an LView if this container is on a component node.
+  // In this case, descend into host LView for further cleanup. See also
+  // LContainer[HOST] docs for additional information.
+  const hostLView = lContainer[HOST];
+  if (isLView(hostLView)) {
+    cleanupLView(hostLView);
+  }
+
   for (let i = CONTAINER_HEADER_OFFSET; i < lContainer.length; i++) {
     cleanupLView(lContainer[i] as LView);
   }
@@ -114,10 +123,6 @@ export function cleanupDehydratedViews(appRef: ApplicationRef) {
       if (isLView(lNode)) {
         cleanupLView(lNode);
       } else {
-        // Cleanup in the root component view
-        const componentLView = lNode[HOST] as LView<unknown>;
-        cleanupLView(componentLView);
-
         // Cleanup in all views within this view container
         cleanupLContainer(lNode);
       }
diff --git a/packages/platform-server/test/hydration_spec.ts b/packages/platform-server/test/hydration_spec.ts
index 7a538868cd4b0..669f61646f8e5 100644
--- a/packages/platform-server/test/hydration_spec.ts
+++ b/packages/platform-server/test/hydration_spec.ts
@@ -7771,6 +7771,70 @@ describe('platform-server hydration integration', () => {
         verifyAllNodesClaimedForHydration(clientRootNode);
         verifyClientAndSSRContentsMatch(ssrContents, clientRootNode);
       });
+
+      it('should cleanup dehydrated views in routed components that use ViewContainerRef', async () => {
+        @Component({
+          standalone: true,
+          selector: 'cmp-a',
+          template: `
+            @if (isServer) {
+              <p>Server view</p>
+            } @else {
+              <p>Client view</p>
+            }
+          `,
+        })
+        class CmpA {
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+          viewContainerRef = inject(ViewContainerRef);
+        }
+
+        const routes: Routes = [
+          {
+            path: '',
+            component: CmpA,
+          },
+        ];
+
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [RouterOutlet],
+          template: `
+            <router-outlet />
+          `,
+        })
+        class SimpleComponent {}
+
+        const envProviders = [
+          {provide: PlatformLocation, useClass: MockPlatformLocation},
+          provideRouter(routes),
+        ] as unknown as Provider[];
+        const html = await ssr(SimpleComponent, {envProviders});
+        const ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain(`<app ${NGH_ATTR_NAME}`);
+        expect(ssrContents).toContain('Server view');
+        expect(ssrContents).not.toContain('Client view');
+
+        resetTViewsFor(SimpleComponent, CmpA);
+
+        const appRef = await renderAndHydrate(doc, html, SimpleComponent, {envProviders});
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        await whenStable(appRef);
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        // <p> tag is used in a view that is different on a server and
+        // on a client, so it gets re-created (not hydrated) on a client
+        const p = clientRootNode.querySelector('p');
+        verifyAllNodesClaimedForHydration(clientRootNode, [p]);
+
+        expect(clientRootNode.innerHTML).not.toContain('Server view');
+        expect(clientRootNode.innerHTML).toContain('Client view');
+      });
     });
   });
 });