ContextNode: groups API (#30347)

* ContextNode: groups API

* remove .only

* review fixes

src/context/index.js

@@ -97,3 +97,35 @@ export function setProp(node, prop, setter, value) {
 export function removeProp(node, prop, setter) {
   ContextNode.get(node).values.remove(prop, setter);
 }
+
+/**
+ * @param {!Node} node
+ * @param {string} name
+ * @param {function(!Node):boolean} match
+ * @param {number=} weight
+ */
+export function addGroup(node, name, match, weight = 0) {
+  ContextNode.get(node).addGroup(name, match, weight);
+}
+
+/**
+ * @param {!Node} node
+ * @param {string} groupName
+ * @param {!ContextProp<T>} prop
+ * @param {*} setter
+ * @param {T} value
+ * @template T
+ */
+export function setGroupProp(node, groupName, prop, setter, value) {
+  ContextNode.get(node).group(groupName).values.set(prop, setter, value);
+}
+
+/**
+ * @param {!Node} node
+ * @param {string} groupName
+ * @param {!ContextProp} prop
+ * @param {*} setter
+ */
+export function removeGroupProp(node, groupName, prop, setter) {
+  ContextNode.get(node).group(groupName).values.remove(prop, setter);
+}

src/context/node.js

@@ -33,6 +33,17 @@ const DOCUMENT_NODE = 9;
 // Includes shadow root, template, etc.
 const FRAGMENT_NODE = 11;
 
+/**
+ * The structure for a group of nodes.
+ *
+ * @typedef {{
+ *   cn: !ContextNode,
+ *   match: function(!Node, !Node):boolean,
+ *   weight: number,
+ * }}
+ */
+let GroupDef;
+
 /**
  * The context node is a sparse tree over the DOM tree. Any node that needs
  * to manage and compute state can be attached to the context node tree. The
@@ -51,7 +62,7 @@ export class ContextNode {
   static get(node) {
     let contextNode = /** @type {!ContextNode|undefined} */ (node[NODE_PROP]);
     if (!contextNode) {
-      contextNode = new ContextNode(node);
+      contextNode = new ContextNode(node, null);
       if (getMode().localDev || getMode().test) {
         // The `Object.defineProperty({enumerable: false})` helps tests, but
         // hurts performance. So this is only done in a dev/test modes.
@@ -160,11 +171,15 @@ export class ContextNode {
    * Creates the context node and automatically starts the discovery process.
    *
    * @param {!Node} node
+   * @param {?string} name
    */
-  constructor(node) {
+  constructor(node, name) {
     /** @const {!Node} */
     this.node = node;
 
+    /** @const @package {?string} */
+    this.name = name;
+
     /**
      * Whether this node is a root. The Document DOM nodes are automatically
      * considered as roots. But other nodes can become roots as well
@@ -200,6 +215,9 @@ export class ContextNode {
      */
     this.children = null;
 
+    /** @package {?Map<string, !GroupDef>} */
+    this.groups = null;
+
     /** @package {!Values} */
     this.values = new Values(this);
 
@@ -245,6 +263,9 @@ export class ContextNode {
   discover() {
     if (this.isDiscoverable()) {
       this.scheduleDiscover_();
+    } else if (this.name && this.children) {
+      // Recursively discover the group's children.
+      this.children.forEach(discoverContextNode);
     }
   }
 
@@ -318,6 +339,55 @@ export class ContextNode {
     }
   }
 
+  /**
+   * @param {string} name
+   * @param {function(!Node):boolean} match
+   * @param {number} weight
+   * @return {!ContextNode}
+   */
+  addGroup(name, match, weight) {
+    const groups = this.groups || (this.groups = new Map());
+    const {node, children} = this;
+    const cn = new ContextNode(node, name);
+    groups.set(name, {cn, match, weight});
+    cn.setParent(this);
+    if (children) {
+      children.forEach(discoverContextNode);
+    }
+    return cn;
+  }
+
+  /**
+   * @param {string} name
+   * @return {?ContextNode}
+   */
+  group(name) {
+    const {groups} = this;
+    const group = groups && groups.get(name);
+    return (group && group.cn) || null;
+  }
+
+  /**
+   * @param {!Node} node
+   * @return {?ContextNode}
+   * @protected
+   */
+  findGroup(node) {
+    const {groups} = this;
+    if (!groups) {
+      return null;
+    }
+    let found = null;
+    let maxWeight = Number.NEGATIVE_INFINITY;
+    groups.forEach(({cn, match, weight}) => {
+      if (match(node, this.node) && weight > maxWeight) {
+        found = cn;
+        maxWeight = weight;
+      }
+    });
+    return found;
+  }
+
   /**
    * Add or update a component with a specified ID. If component doesn't
    * yet exist, it will be created using the specified factory. The use
@@ -386,7 +456,9 @@ export class ContextNode {
       // queue.
       return;
     }
-    const parent = ContextNode.closest(this.node, /* includeSelf */ false);
+    const closestNode = ContextNode.closest(this.node, /* includeSelf */ false);
+    const parent =
+      (closestNode && closestNode.findGroup(this.node)) || closestNode;
     this.updateTree_(parent, /* parentOverridden */ false);
   }
 

test/unit/context/test-node-discover.js

@@ -134,10 +134,13 @@ describes.realWin('ContextNode', {}, (env) => {
     }
     if (spec.parent !== undefined) {
       const {parent} = contextNode;
-      expect((parent && parent.node) ?? null, 'parent').to.equal(spec.parent);
+      const parentNode = (parent && parent.node) ?? null;
+      const specNode =
+        ((spec.parent && spec.parent.node) || spec.parent) ?? null;
+      expect(parentNode, 'parent').to.equal(specNode);
     }
     if (spec.children !== undefined) {
-      const children = (contextNode.children || []).map((cn) => cn.node);
+      const children = (contextNode.children || []).map((cn) => cn.node || cn);
       children.sort(domOrderComparator);
       const specChildren = spec.children.slice(0);
       specChildren.sort(domOrderComparator);
@@ -1002,6 +1005,84 @@ describes.realWin('ContextNode', {}, (env) => {
     });
   });
 
+  describe('discover groups', () => {
+    let sibling1;
+    let sibling2;
+    let cousin1;
+    let parent;
+    let grandparent;
+
+    beforeEach(async () => {
+      sibling1 = el('T-1-1-1');
+      sibling2 = el('T-1-1-2');
+      cousin1 = el('T-1-2-1');
+      parent = el('T-1-1');
+      grandparent = el('T-1');
+      await waitForDiscover(grandparent, parent, sibling1, sibling2, cousin1);
+    });
+
+    it('should rediscover children when a new group is added', async () => {
+      const group1 = ContextNode.get(parent).addGroup(
+        'group1',
+        (node) => node == sibling1,
+        0
+      );
+      await waitForDiscover(group1, sibling1);
+      clock.runAll();
+
+      expect(group1.node).to.equal(parent);
+      expectContext(group1, {parent, children: [sibling1]});
+
+      // sibling1 is reassigned to the group.
+      expectContext(sibling1, {parent: group1});
+
+      // sibling2 stays stays unchanged.
+      expectContext(sibling2, {parent});
+    });
+
+    it('should discover a new child', async () => {
+      const sibling3 = el('T-1-1-3');
+      const group1 = ContextNode.get(parent).addGroup(
+        'group1',
+        (node) => node == sibling3,
+        0
+      );
+      await waitForDiscover(group1);
+
+      // Discover the new node.
+      await rediscover(sibling3);
+      expectContext(sibling3, {parent: group1});
+    });
+
+    it('should handle weight', async () => {
+      const group1 = ContextNode.get(parent).addGroup(
+        'group1',
+        (node) => node == sibling1,
+        0
+      );
+      await waitForDiscover(group1, sibling1);
+      expectContext(sibling1, {parent: group1});
+
+      // A lower weight.
+      const group2 = ContextNode.get(parent).addGroup(
+        'group1',
+        (node) => node == sibling1,
+        -1
+      );
+      await waitForDiscover(group2, sibling1);
+      expectContext(sibling1, {parent: group1});
+
+      // A higher weight.
+      const group3 = ContextNode.get(parent).addGroup(
+        'group1',
+        (node) => node == sibling1,
+        1
+      );
+      await waitForDiscover(group3, sibling1);
+      expectContext(sibling1, {parent: group3});
+    });
+  });
+
   describe('scanners', () => {
     const EXCLUDE_SELF = false;