From ff78c66ae917c1cf02a83b30aae375860119710a Mon Sep 17 00:00:00 2001
From: mgechev <mgechev@gmail.com>
Date: Tue, 9 Nov 2021 09:22:49 -0800
Subject: [PATCH] fix(devtools): handle ngFor and ngIf contexts in the profiler

Currently, we skip the time spent in evaluation of embedded view
templates because the passed context is either an NgIf or an NgForOf.

This change generalizes the implementation to work with views embedded
into a component.
---
 .../src/lib/hooks/capture.ts                  |  7 +--
 .../src/lib/hooks/profiler/native.ts          | 44 +++++++++++++++----
 2 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/projects/ng-devtools-backend/src/lib/hooks/capture.ts b/projects/ng-devtools-backend/src/lib/hooks/capture.ts
index 28c855a7316..e698b1f8e6e 100644
--- a/projects/ng-devtools-backend/src/lib/hooks/capture.ts
+++ b/projects/ng-devtools-backend/src/lib/hooks/capture.ts
@@ -31,20 +31,20 @@ export const stop = (): ProfilerFrame => {
   return result;
 };
 
-const startEvent = (map: { [key: string]: number }, directive: any, label: string) => {
+const startEvent = (map: Record<string, number>, directive: any, label: string) => {
   const name = getDirectiveName(directive);
   const key = `${name}#${label}`;
   map[key] = performance.now();
 };
 
-const getEventStart = (map: { [key: string]: number }, directive: any, label: string) => {
+const getEventStart = (map: Record<string, number>, directive: any, label: string) => {
   const name = getDirectiveName(directive);
   const key = `${name}#${label}`;
   return map[key];
 };
 
 const getHooks = (onFrame: (frame: ProfilerFrame) => void): Partial<Hooks> => {
-  const timeStartMap: { [key: string]: number } = {};
+  const timeStartMap: Record<string, number> = {};
   return {
     // We flush here because it's possible the current node to overwrite
     // an existing removed node.
@@ -82,6 +82,7 @@ const getHooks = (onFrame: (frame: ProfilerFrame) => void): Partial<Hooks> => {
     },
     onChangeDetectionEnd(component: any): void {
       const profile = eventMap.get(component);
+
       if (profile) {
         let current = profile.changeDetection;
         if (current === undefined) {
diff --git a/projects/ng-devtools-backend/src/lib/hooks/profiler/native.ts b/projects/ng-devtools-backend/src/lib/hooks/profiler/native.ts
index a4c84c126c8..d86a0fcc8f1 100644
--- a/projects/ng-devtools-backend/src/lib/hooks/profiler/native.ts
+++ b/projects/ng-devtools-backend/src/lib/hooks/profiler/native.ts
@@ -10,6 +10,7 @@ type ProfilerCallback = (event: ɵProfilerEvent, instanceOrLView: {}, hookOrList
 export class NgProfiler extends Profiler {
   private _tracker = IdentityTracker.getInstance();
   private _callbacks: ProfilerCallback[] = [];
+  private _lastDirectiveInstance: {} | null = null;
 
   constructor(config: Partial<Hooks> = {}) {
     super(config);
@@ -66,7 +67,7 @@ export class NgProfiler extends Profiler {
     return;
   }
 
-  [ɵProfilerEvent.TemplateUpdateStart](directive: any, _hookOrListener: any): void {
+  [ɵProfilerEvent.TemplateUpdateStart](context: any, _hookOrListener: any): void {
     if (!this._inChangeDetection) {
       this._inChangeDetection = true;
       runOutsideAngular(() => {
@@ -77,19 +78,44 @@ export class NgProfiler extends Profiler {
       });
     }
 
-    const position = this._tracker.getDirectivePosition(directive);
-    const id = this._tracker.getDirectiveId(directive);
+    const position = this._tracker.getDirectivePosition(context);
+    const id = this._tracker.getDirectiveId(context);
 
-    this._onChangeDetectionStart(directive, getDirectiveHostElement(directive), id, position);
+    // If we can find the position and the ID we assume that this is a component instance.
+    // Alternatively, if we can't find the ID or the position, we assume that this is a
+    // context of an embedded view (for example, NgForOfContext, NgIfContext, or a custom one).
+    if (position !== undefined && id !== undefined) {
+      this._lastDirectiveInstance = context;
+    }
+
+    if (id !== undefined && position !== undefined) {
+      this._onChangeDetectionStart(context, getDirectiveHostElement(context), id, position);
+      return;
+    }
+
+    this._onChangeDetectionStart(
+      this._lastDirectiveInstance,
+      getDirectiveHostElement(this._lastDirectiveInstance),
+      this._tracker.getDirectiveId(this._lastDirectiveInstance),
+      this._tracker.getDirectivePosition(this._lastDirectiveInstance)
+    );
   }
 
-  [ɵProfilerEvent.TemplateUpdateEnd](directive: any, _hookOrListener: any): void {
-    const position = this._tracker.getDirectivePosition(directive);
-    const id = this._tracker.getDirectiveId(directive);
+  [ɵProfilerEvent.TemplateUpdateEnd](context: any, _hookOrListener: any): void {
+    const position = this._tracker.getDirectivePosition(context);
+    const id = this._tracker.getDirectiveId(context);
 
-    if (this._tracker.hasDirective(directive) && id !== undefined && position !== undefined) {
-      this._onChangeDetectionEnd(directive, getDirectiveHostElement(directive), id, position);
+    if (this._tracker.hasDirective(context) && id !== undefined && position !== undefined) {
+      this._onChangeDetectionEnd(context, getDirectiveHostElement(context), id, position);
+      return;
     }
+
+    this._onChangeDetectionEnd(
+      this._lastDirectiveInstance,
+      getDirectiveHostElement(this._lastDirectiveInstance),
+      this._tracker.getDirectiveId(this._lastDirectiveInstance),
+      this._tracker.getDirectivePosition(this._lastDirectiveInstance)
+    );
   }
 
   [ɵProfilerEvent.LifecycleHookStart](directive: any, hook: any): void {