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

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.

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,12 @@ 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) {
+      const info = this._getClosestElementInfo(nativeElement);
+
+      if (info) {
+        this._originChanged(info[0], origin, info[1]);
+      }
     } else {
       this._setOriginForCurrentEventQueue(origin);
 
@@ -553,6 +557,29 @@ export class FocusMonitor implements OnDestroy {
     this._emitOrigin(elementInfo.subject, origin);
     this._lastFocusOrigin = origin;
   }
+
+  /**
+   * Gets the `MonitoredElementInfo` of an element, or the closest
+   * monitored parent that has `checkChildren` enabled.
+   * @param element Element from which to start the search.
+   */
+  private _getClosestElementInfo(element: HTMLElement): [HTMLElement, MonitoredElementInfo] | null {
+    const ownInfo = this._elementInfo.get(element);
+
+    // Do a quick constant-time lookup for the element itself.
+    if (ownInfo) {
+      return [element, ownInfo];
+    }
+
+    // Otherwise do a linear lookup only for the elements that are monitoring child nodes.
+    for (const current of this._elementInfo) {
+      if (current[1].checkChildren && current[0].contains(element)) {
+        return current;
+      }
+    }
+
+    return null;
+  }
 }
 
 /** Gets the target of an event, accounting for Shadow DOM. */