fix(cdk/a11y): focusVia not accounting for focused child node (#21512)

In #20966 some logic was added so that calling `focusVia` on an element that already
has focus would change the origin to the passed-in one. The problem is that the new
logic doesn't account for when a parent element is monitored and `focusVia` is called
on a child.

These changes add some more logic that will look through all the monitored elements
that have `checkChildren: true` and will switch the origin accordingly.

Fixes #21500.

(cherry picked from commit f8df9f8)

Author: crisbeto

src/cdk/a11y/focus-monitor/focus-monitor.spec.ts

@@ -352,6 +352,40 @@ describe('FocusMonitor', () => {
     expect(buttonElement.focus).toHaveBeenCalledTimes(1);
   }));
 
+  it('focusVia should change the focus origin when called a focused child node', fakeAsync(() => {
+    const parent = fixture.nativeElement.querySelector('.parent');
+    focusMonitor.stopMonitoring(buttonElement); // The button gets monitored by default.
+    focusMonitor.monitor(parent, true).subscribe(changeHandler);
+    spyOn(buttonElement, 'focus').and.callThrough();
+    focusMonitor.focusVia(buttonElement, 'keyboard');
+    flush();
+    fakeActiveElement = buttonElement;
+
+    expect(parent.classList.length)
+        .toBe(3, 'Parent should have exactly 2 focus classes and the `parent` class');
+    expect(parent.classList.contains('cdk-focused'))
+        .toBe(true, 'Parent should have cdk-focused class');
+    expect(parent.classList.contains('cdk-keyboard-focused'))
+        .toBe(true, 'Parent should have cdk-keyboard-focused class');
+    expect(changeHandler).toHaveBeenCalledTimes(1);
+    expect(changeHandler).toHaveBeenCalledWith('keyboard');
+    expect(buttonElement.focus).toHaveBeenCalledTimes(1);
+
+    focusMonitor.focusVia(buttonElement, 'mouse');
+    flush();
+    fakeActiveElement = buttonElement;
+
+    expect(parent.classList.length)
+        .toBe(3, 'Parent should have exactly 2 focus classes and the `parent` class');
+    expect(parent.classList.contains('cdk-focused'))
+        .toBe(true, 'Parent should have cdk-focused class');
+    expect(parent.classList.contains('cdk-mouse-focused'))
+        .toBe(true, 'Parent should have cdk-mouse-focused class');
+    expect(changeHandler).toHaveBeenCalledTimes(2);
+    expect(changeHandler).toHaveBeenCalledWith('mouse');
+    expect(buttonElement.focus).toHaveBeenCalledTimes(1);
+  }));
+
 });
 
 describe('FocusMonitor with "eventual" detection', () => {

src/cdk/a11y/focus-monitor/focus-monitor.ts

@@ -313,8 +313,9 @@ export class FocusMonitor implements OnDestroy {
     // If the element is focused already, calling `focus` again won't trigger the event listener
     // which means that the focus classes won't be updated. If that's the case, update the classes
     // directly without waiting for an event.
-    if (nativeElement === focusedElement && this._elementInfo.has(nativeElement)) {
-      this._originChanged(nativeElement, origin, this._elementInfo.get(nativeElement)!);
+    if (nativeElement === focusedElement) {
+      this._getClosestElementsInfo(nativeElement)
+        .forEach(([currentElement, info]) => this._originChanged(currentElement, origin, info));
     } else {
       this._setOriginForCurrentEventQueue(origin);
 
@@ -553,6 +554,23 @@ export class FocusMonitor implements OnDestroy {
     this._emitOrigin(elementInfo.subject, origin);
     this._lastFocusOrigin = origin;
   }
+
+  /**
+   * Collects the `MonitoredElementInfo` of a particular element and
+   * all of its ancestors that have enabled `checkChildren`.
+   * @param element Element from which to start the search.
+   */
+  private _getClosestElementsInfo(element: HTMLElement): [HTMLElement, MonitoredElementInfo][] {
+    const results: [HTMLElement, MonitoredElementInfo][] = [];
+
+    this._elementInfo.forEach((info, currentElement) => {
+      if (currentElement === element || (info.checkChildren && currentElement.contains(element))) {
+        results.push([currentElement, info]);
+      }
+    });
+
+    return results;
+  }
 }
 
 /** Gets the target of an event, accounting for Shadow DOM. */