From 3ef5d8740871a2d12879c2d0c0b54c2c3fe4e673 Mon Sep 17 00:00:00 2001
From: Andrew Kushnir <akushnir@google.com>
Date: Thu, 16 Mar 2023 16:51:26 -0700
Subject: [PATCH] refactor(core): post-hydration cleanup of unclaimed views
 (#49455)

This commit adds a logic to remove all views that were not cleaimed during the hydration. The process is started once the ApplicationRef becomes stable on the client (which matches the timing of serialization on the server).

PR Close #49455
---
 packages/core/src/hydration/api.ts            |  13 +
 packages/core/src/hydration/cleanup.ts        |  59 +++-
 packages/core/src/util/ng_dev_mode.ts         |   2 +
 .../platform-server/test/hydration_spec.ts    | 330 ++++++++++++++++++
 4 files changed, 402 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/hydration/api.ts b/packages/core/src/hydration/api.ts
index 6e0d051199a..602748927b9 100644
--- a/packages/core/src/hydration/api.ts
+++ b/packages/core/src/hydration/api.ts
@@ -6,6 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
+import {APP_BOOTSTRAP_LISTENER, ApplicationRef} from '../application_ref';
 import {PLATFORM_ID} from '../application_tokens';
 import {ENVIRONMENT_INITIALIZER, EnvironmentProviders, makeEnvironmentProviders} from '../di';
 import {inject} from '../di/injector_compatibility';
@@ -16,6 +17,7 @@ import {enableApplyRootElementTransformImpl} from '../render3/instructions/share
 import {enableLocateOrCreateContainerAnchorImpl} from '../render3/instructions/template';
 import {enableLocateOrCreateTextNodeImpl} from '../render3/instructions/text';
 
+import {cleanupDehydratedViews} from './cleanup';
 import {IS_HYDRATION_FEATURE_ENABLED, PRESERVE_HOST_CONTENT} from './tokens';
 import {enableRetrieveHydrationInfoImpl} from './utils';
 import {enableFindMatchingDehydratedViewImpl} from './views';
@@ -129,6 +131,17 @@ export function provideHydrationSupport(): EnvironmentProviders {
       // environment. On a server, an application is rendered
       // from scratch, so the host content needs to be empty.
       useFactory: () => isBrowser(),
+    },
+    {
+      provide: APP_BOOTSTRAP_LISTENER,
+      useFactory: () => {
+        if (isBrowser()) {
+          const appRef = inject(ApplicationRef);
+          return () => cleanupDehydratedViews(appRef);
+        }
+        return () => {};  // noop
+      },
+      multi: true,
     }
   ]);
 }
diff --git a/packages/core/src/hydration/cleanup.ts b/packages/core/src/hydration/cleanup.ts
index 59739aa5faa..8b8e4ea4587 100644
--- a/packages/core/src/hydration/cleanup.ts
+++ b/packages/core/src/hydration/cleanup.ts
@@ -6,15 +6,20 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {DEHYDRATED_VIEWS, LContainer} from '../render3/interfaces/container';
+import {first} from 'rxjs/operators';
+
+import {ApplicationRef} from '../application_ref';
+import {CONTAINER_HEADER_OFFSET, DEHYDRATED_VIEWS, LContainer} from '../render3/interfaces/container';
 import {Renderer} from '../render3/interfaces/renderer';
 import {RNode} from '../render3/interfaces/renderer_dom';
-import {PARENT, RENDERER} from '../render3/interfaces/view';
+import {isLContainer} from '../render3/interfaces/type_checks';
+import {HEADER_OFFSET, HOST, LView, PARENT, RENDERER, TVIEW} from '../render3/interfaces/view';
 import {nativeRemoveNode} from '../render3/node_manipulation';
 import {EMPTY_ARRAY} from '../util/empty';
 
 import {validateSiblingNodeExists} from './error_handling';
 import {DehydratedContainerView, NUM_ROOT_NODES} from './interfaces';
+import {getComponentLViewForHydration} from './utils';
 
 /**
  * Removes all dehydrated views from a given LContainer:
@@ -53,3 +58,53 @@ function removeDehydratedView(dehydratedView: DehydratedContainerView, renderer:
     }
   }
 }
+
+/**
+ * Walks over all views within this LContainer invokes dehydrated views
+ * cleanup function for each one.
+ */
+function cleanupLContainer(lContainer: LContainer) {
+  removeDehydratedViews(lContainer);
+  for (let i = CONTAINER_HEADER_OFFSET; i < lContainer.length; i++) {
+    cleanupLView(lContainer[i] as LView);
+  }
+}
+
+/**
+ * Walks over `LContainer`s and components registered within
+ * this LView and invokes dehydrated views cleanup function for each one.
+ */
+function cleanupLView(lView: LView) {
+  const tView = lView[TVIEW];
+  for (let i = HEADER_OFFSET; i < tView.bindingStartIndex; i++) {
+    if (isLContainer(lView[i])) {
+      const lContainer = lView[i];
+      cleanupLContainer(lContainer);
+    } else if (Array.isArray(lView[i])) {
+      // This is a component, enter the `cleanupLView` recursively.
+      cleanupLView(lView[i]);
+    }
+  }
+}
+
+/**
+ * Walks over all views registered within the ApplicationRef and removes
+ * all dehydrated views from all `LContainer`s along the way.
+ */
+export function cleanupDehydratedViews(appRef: ApplicationRef) {
+  // Wait once an app becomes stable and cleanup all views that
+  // were not claimed during the application bootstrap process.
+  // The timing is similar to when we kick off serialization on the server.
+  return appRef.isStable.pipe(first((isStable: boolean) => isStable)).toPromise().then(() => {
+    const viewRefs = appRef._views;
+    for (const viewRef of viewRefs) {
+      const lView = getComponentLViewForHydration(viewRef);
+      // An `lView` might be `null` if a `ViewRef` represents
+      // an embedded view (not a component view).
+      if (lView !== null && lView[HOST] !== null) {
+        cleanupLView(lView);
+        ngDevMode && ngDevMode.dehydratedViewsCleanupRuns++;
+      }
+    }
+  });
+}
diff --git a/packages/core/src/util/ng_dev_mode.ts b/packages/core/src/util/ng_dev_mode.ts
index faf58e6ff45..fea7327022f 100644
--- a/packages/core/src/util/ng_dev_mode.ts
+++ b/packages/core/src/util/ng_dev_mode.ts
@@ -51,6 +51,7 @@ declare global {
     hydratedNodes: number;
     hydratedComponents: number;
     dehydratedViewsRemoved: number;
+    dehydratedViewsCleanupRuns: number;
   }
 }
 
@@ -83,6 +84,7 @@ export function ngDevModeResetPerfCounters(): NgDevModePerfCounters {
     hydratedNodes: 0,
     hydratedComponents: 0,
     dehydratedViewsRemoved: 0,
+    dehydratedViewsCleanupRuns: 0,
   };
 
   // Make sure to refer to ngDevMode as ['ngDevMode'] for closure.
diff --git a/packages/platform-server/test/hydration_spec.ts b/packages/platform-server/test/hydration_spec.ts
index b9624eb34a8..a30d82401ac 100644
--- a/packages/platform-server/test/hydration_spec.ts
+++ b/packages/platform-server/test/hydration_spec.ts
@@ -10,6 +10,7 @@ import {CommonModule, DOCUMENT, isPlatformServer, NgComponentOutlet, NgFor, NgIf
 import {APP_ID, ApplicationRef, Component, ComponentRef, createComponent, destroyPlatform, ElementRef, EnvironmentInjector, getPlatform, inject, Input, PLATFORM_ID, Provider, TemplateRef, Type, ViewChild, ViewContainerRef, ɵgetComponentDef as getComponentDef, ɵprovideHydrationSupport as provideHydrationSupport, ɵsetDocument} from '@angular/core';
 import {TestBed} from '@angular/core/testing';
 import {bootstrapApplication} from '@angular/platform-browser';
+import {first} from 'rxjs/operators';
 
 import {renderApplication} from '../src/utils';
 
@@ -79,6 +80,15 @@ function stripTransferDataScript(input: string): string {
   return input.replace(/<script (.*?)<\/script>/s, '');
 }
 
+function stripExcessiveSpaces(html: string): string {
+  return html.replace(/\s+/g, ' ');
+}
+
+/** Returns a Promise that resolves when the ApplicationRef becomes stable. */
+function whenStable(appRef: ApplicationRef): Promise<unknown> {
+  return appRef.isStable.pipe(first((isStable: boolean) => isStable)).toPromise();
+}
+
 function verifyClientAndSSRContentsMatch(ssrContents: string, clientAppRootElement: HTMLElement) {
   const clientContents =
       stripTransferDataScript(stripUtilAttributes(clientAppRootElement.outerHTML, false));
@@ -618,6 +628,36 @@ describe('platform-server integration', () => {
             verifyClientAndSSRContentsMatch(ssrContents, clientRootNode);
           });
 
+          it('should work with empty containers on ng-container nodes', async () => {
+            @Component({
+              standalone: true,
+              selector: 'app',
+              imports: [NgIf],
+              template: `
+                This is an empty container:
+                <ng-container *ngIf="false" />
+                <div>Post-container element</div>
+              `,
+            })
+            class SimpleComponent {
+            }
+
+            const html = await ssr(SimpleComponent);
+            const ssrContents = getAppContents(html);
+
+            expect(ssrContents).toContain(`<app ${NGH_ATTR_NAME}`);
+
+            resetTViewsFor(SimpleComponent);
+
+            const appRef = await hydrate(html, SimpleComponent);
+            const compRef = getComponentRef<SimpleComponent>(appRef);
+            appRef.tick();
+
+            const clientRootNode = compRef.location.nativeElement;
+            verifyAllNodesClaimedForHydration(clientRootNode);
+            verifyClientAndSSRContentsMatch(ssrContents, clientRootNode);
+          });
+
           it('should work with *ngIf on element nodes', async () => {
             @Component({
               standalone: true,
@@ -646,6 +686,29 @@ describe('platform-server integration', () => {
             verifyClientAndSSRContentsMatch(ssrContents, clientRootNode);
           });
 
+          it('should work with empty containers on element nodes', async () => {
+            @Component({
+              standalone: true,
+              selector: 'app',
+              imports: [NgIf],
+              template: `
+                <h1 *ngIf="false">Hello world!</h1>
+              `,
+            })
+            class SimpleComponent {
+            }
+            const html = await ssr(SimpleComponent);
+            const ssrContents = getAppContents(html);
+            expect(ssrContents).toContain(`<app ${NGH_ATTR_NAME}`);
+            resetTViewsFor(SimpleComponent);
+            const appRef = await hydrate(html, SimpleComponent);
+            const compRef = getComponentRef<SimpleComponent>(appRef);
+            appRef.tick();
+            const clientRootNode = compRef.location.nativeElement;
+            verifyAllNodesClaimedForHydration(clientRootNode);
+            verifyClientAndSSRContentsMatch(ssrContents, clientRootNode);
+          });
+
           it('should work with *ngIf on component host nodes', async () => {
             @Component({
               standalone: true,
@@ -1121,6 +1184,9 @@ describe('platform-server integration', () => {
                  'if there is a mismatch in template ids between the current view ' +
                  '(that is being created) and the first dehydrated view in the list',
              async () => {
+               // Reset the relevant counter.
+               (ngDevMode as any).dehydratedViewsRemoved = 0;
+
                @Component({
                  standalone: true,
                  selector: 'app',
@@ -1766,6 +1832,270 @@ describe('platform-server integration', () => {
       });
     });
 
+    describe('post-hydration cleanup', () => {
+      it('should cleanup unclaimed views in a component (when using elements)', async () => {
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [NgIf],
+          template: `
+            <b *ngIf="isServer">This is a SERVER-ONLY content</b>
+            <i *ngIf="!isServer">This is a CLIENT-ONLY content</i>
+          `,
+        })
+        class SimpleComponent {
+          // This flag is intentionally different between the client
+          // and the server: we use it to test the logic to cleanup
+          // dehydrated views.
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+        }
+
+        const html = await ssr(SimpleComponent);
+        let ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain('<app ngh');
+
+        resetTViewsFor(SimpleComponent);
+
+        const appRef = await hydrate(html, SimpleComponent);
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        ssrContents = stripExcessiveSpaces(stripUtilAttributes(ssrContents, false));
+
+        // In the SSR output we expect to see SERVER content, but not CLIENT.
+        expect(ssrContents).not.toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(ssrContents).toContain('<b>This is a SERVER-ONLY content</b>');
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        await whenStable(appRef);
+
+        const clientContents =
+            stripExcessiveSpaces(stripUtilAttributes(clientRootNode.outerHTML, false));
+
+        // After the cleanup, we expect to see CLIENT content, but not SERVER.
+        expect(clientContents).toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(clientContents).not.toContain('<b>This is a SERVER-ONLY content</b>');
+      });
+
+      it('should cleanup unclaimed views in a component (when using <ng-container>s)', async () => {
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [NgIf],
+          template: `
+            <ng-container *ngIf="isServer">This is a SERVER-ONLY content</ng-container>
+            <ng-container *ngIf="!isServer">This is a CLIENT-ONLY content</ng-container>
+          `,
+        })
+        class SimpleComponent {
+          // This flag is intentionally different between the client
+          // and the server: we use it to test the logic to cleanup
+          // dehydrated views.
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+        }
+
+        const html = await ssr(SimpleComponent);
+        let ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain('<app ngh');
+
+        resetTViewsFor(SimpleComponent);
+
+        const appRef = await hydrate(html, SimpleComponent);
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        ssrContents = stripExcessiveSpaces(stripUtilAttributes(ssrContents, false));
+
+        // In the SSR output we expect to see SERVER content, but not CLIENT.
+        expect(ssrContents).not.toContain('This is a CLIENT-ONLY content<!--ng-container-->');
+        expect(ssrContents).toContain('This is a SERVER-ONLY content<!--ng-container-->');
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        await whenStable(appRef);
+
+        const clientContents =
+            stripExcessiveSpaces(stripUtilAttributes(clientRootNode.outerHTML, false));
+
+        // After the cleanup, we expect to see CLIENT content, but not SERVER.
+        expect(clientContents).toContain('This is a CLIENT-ONLY content<!--ng-container-->');
+        expect(clientContents).not.toContain('This is a SERVER-ONLY content<!--ng-container-->');
+      });
+
+      it('should cleanup within inner containers', async () => {
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [NgIf],
+          template: `
+            <ng-container *ngIf="true">
+              <b *ngIf="isServer">This is a SERVER-ONLY content</b>
+              Outside of the container (must be retained).
+            </ng-container>
+            <i *ngIf="!isServer">This is a CLIENT-ONLY content</i>
+          `,
+        })
+        class SimpleComponent {
+          // This flag is intentionally different between the client
+          // and the server: we use it to test the logic to cleanup
+          // dehydrated views.
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+        }
+
+        const html = await ssr(SimpleComponent);
+        let ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain('<app ngh');
+
+        resetTViewsFor(SimpleComponent);
+
+        const appRef = await hydrate(html, SimpleComponent);
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        ssrContents = stripExcessiveSpaces(stripUtilAttributes(ssrContents, false));
+
+        // In the SSR output we expect to see SERVER content, but not CLIENT.
+        expect(ssrContents).not.toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(ssrContents).toContain('<b>This is a SERVER-ONLY content</b>');
+        expect(ssrContents).toContain('Outside of the container (must be retained).');
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        await whenStable(appRef);
+
+        const clientContents =
+            stripExcessiveSpaces(stripUtilAttributes(clientRootNode.outerHTML, false));
+
+        // After the cleanup, we expect to see CLIENT content, but not SERVER.
+        expect(clientContents).toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(clientContents).not.toContain('<b>This is a SERVER-ONLY content</b>');
+
+        // This line must be preserved (it's outside of the dehydrated container).
+        expect(clientContents).toContain('Outside of the container (must be retained).');
+      });
+
+      it('should reconcile *ngFor-generated views', async () => {
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [NgIf, NgFor],
+          template: `
+            <div>
+              <span *ngFor="let item of items">
+                {{ item }}
+                <b *ngIf="item > 15">is bigger than 15!</b>
+              </span>
+              <main>Hi! This is the main content.</main>
+            </div>
+          `,
+        })
+        class SimpleComponent {
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+          // Note: this is needed to test cleanup/reconciliation logic.
+          items = this.isServer ? [10, 20, 100, 200] : [30, 5, 50];
+        }
+
+        const html = await ssr(SimpleComponent);
+        const ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain('<app ngh');
+
+        resetTViewsFor(SimpleComponent);
+
+        const appRef = await hydrate(html, SimpleComponent);
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        await whenStable(appRef);
+
+        // Post-cleanup should *not* contain dehydrated views.
+        const postCleanupContents = stripExcessiveSpaces(clientRootNode.outerHTML);
+        expect(postCleanupContents)
+            .not.toContain(
+                '<span> 5 <b>is bigger than 15!</b><!--bindings={ "ng-reflect-ng-if": "false" }--></span>');
+        expect(postCleanupContents)
+            .toContain(
+                '<span> 30 <b>is bigger than 15!</b><!--bindings={ "ng-reflect-ng-if": "true" }--></span>');
+        expect(postCleanupContents)
+            .toContain('<span> 5 <!--bindings={ "ng-reflect-ng-if": "false" }--></span>');
+        expect(postCleanupContents)
+            .toContain(
+                '<span> 50 <b>is bigger than 15!</b><!--bindings={ "ng-reflect-ng-if": "true" }--></span>');
+      });
+
+      it('should cleanup dehydrated views within dynamically created components', async () => {
+        @Component({
+          standalone: true,
+          imports: [CommonModule],
+          selector: 'dynamic',
+          template: `
+            <span>This is a content of a dynamic component.</span>
+            <b *ngIf="isServer">This is a SERVER-ONLY content</b>
+            <i *ngIf="!isServer">This is a CLIENT-ONLY content</i>
+            <ng-container *ngIf="isServer">
+              This is also a SERVER-ONLY content, but inside ng-container.
+              <b>With some extra tags</b> and some text inside.
+            </ng-container>
+          `,
+        })
+        class DynamicComponent {
+          isServer = isPlatformServer(inject(PLATFORM_ID));
+        }
+
+        @Component({
+          standalone: true,
+          selector: 'app',
+          imports: [NgIf, NgFor],
+          template: `
+            <div #target></div>
+            <main>Hi! This is the main content.</main>
+          `,
+        })
+        class SimpleComponent {
+          @ViewChild('target', {read: ViewContainerRef}) vcr!: ViewContainerRef;
+
+          ngAfterViewInit() {
+            const compRef = this.vcr.createComponent(DynamicComponent);
+            compRef.changeDetectorRef.detectChanges();
+          }
+        }
+
+        const html = await ssr(SimpleComponent);
+        let ssrContents = getAppContents(html);
+
+        expect(ssrContents).toContain('<app ngh');
+
+        resetTViewsFor(SimpleComponent, DynamicComponent);
+
+        const appRef = await hydrate(html, SimpleComponent);
+        const compRef = getComponentRef<SimpleComponent>(appRef);
+        appRef.tick();
+
+        ssrContents = stripExcessiveSpaces(stripUtilAttributes(ssrContents, false));
+
+        // We expect to see SERVER content, but not CLIENT.
+        expect(ssrContents).not.toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(ssrContents).toContain('<b>This is a SERVER-ONLY content</b>');
+
+        const clientRootNode = compRef.location.nativeElement;
+
+        await whenStable(appRef);
+
+        const clientContents =
+            stripExcessiveSpaces(stripUtilAttributes(clientRootNode.outerHTML, false));
+
+        // After the cleanup, we expect to see CLIENT content, but not SERVER.
+        expect(clientContents).toContain('<i>This is a CLIENT-ONLY content</i>');
+        expect(clientContents).not.toContain('<b>This is a SERVER-ONLY content</b>');
+      });
+    });
+
     describe('content projection', () => {
       it('should project plain text', async () => {
         @Component({