[lit-html] Fix ref bug when auto-bound class method used as as call…

…back (#2691)

* [lit-html] Key "last element" map on context in addition to callback

* Add changeset

* Flip order of context and callback lookups

* Add lit to changeset

* Update .changeset/short-cooks-deliver.md

Co-authored-by: Andrew Jakubowicz <spyr1014@gmail.com>

* Updates based on feedback

Co-authored-by: Andrew Jakubowicz <spyr1014@gmail.com>

.changeset/short-cooks-deliver.md

@@ -0,0 +1,6 @@
+---
+'lit-html': patch
+'lit': patch
+---
+
+Fixes `ref` bug when auto-bound class method used as a callback could incorrectly receive `undefined`.

packages/lit-html/src/directives/ref.ts

@@ -30,16 +30,20 @@ interface RefInternal {
 
 // When callbacks are used for refs, this map tracks the last value the callback
 // was called with, for ensuring a directive doesn't clear the ref if the ref
-// has already been rendered to a new spot
-const lastElementForCallback: WeakMap<Function, Element | undefined> =
-  new WeakMap();
+// has already been rendered to a new spot. It is double-keyed on both the
+// context (`options.host`) and the callback, since we auto-bind class methods
+// to `options.host`.
+const lastElementForContextAndCallback: WeakMap<
+  object,
+  WeakMap<Function, Element | undefined>
+> = new WeakMap();
 
 export type RefOrCallback = Ref | ((el: Element | undefined) => void);
 
 class RefDirective extends AsyncDirective {
   private _element?: Element;
   private _ref?: RefOrCallback;
-  private _context: unknown;
+  private _context?: object;
 
   render(_ref: RefOrCallback) {
     return nothing;
@@ -69,7 +73,17 @@ class RefDirective extends AsyncDirective {
       // way regardless of whether a ref might be moving up in the tree (in
       // which case it would otherwise be called with the new value before the
       // previous one unsets it) and down in the tree (where it would be unset
-      // before being set)
+      // before being set). Note that element lookup is keyed by
+      // both the context and the callback, since we allow passing unbound
+      // functions that are called on options.host, and we want to treat
+      // these as unique "instances" of a function.
+      const context = this._context ?? globalThis;
+      let lastElementForCallback =
+        lastElementForContextAndCallback.get(context);
+      if (lastElementForCallback === undefined) {
+        lastElementForCallback = new WeakMap();
+        lastElementForContextAndCallback.set(context, lastElementForCallback);
+      }
       if (lastElementForCallback.get(this._ref) !== undefined) {
         this._ref.call(this._context, undefined);
       }
@@ -85,7 +99,9 @@ class RefDirective extends AsyncDirective {
 
   private get _lastElementForRef() {
     return typeof this._ref === 'function'
-      ? lastElementForCallback.get(this._ref)
+      ? lastElementForContextAndCallback
+          .get(this._context ?? globalThis)
+          ?.get(this._ref)
       : this._ref?.value;
   }
 

packages/lit-html/src/test/directives/ref_test.ts

@@ -140,41 +140,76 @@ suite('ref', () => {
   });
 
   test('calls callback bound to options.host', () => {
-    let queriedEl: Element | null;
-    const host = {
-      calls: [] as Array<string | undefined>,
+    class Host {
+      bool = false;
+      calls: Array<string | undefined> = [];
+      root = document.createElement('div');
       elCallback(e: Element | undefined) {
         this.calls.push(e?.tagName);
-      },
+      }
+      render() {
+        return render(
+          this.bool
+            ? html`<div ${ref(this.elCallback)}></div>`
+            : html`<span ${ref(this.elCallback)}></span>`,
+          this.root,
+          {host: this}
+        );
+      }
+    }
+
+    const testRef = (
+      host: Host,
+      initialCalls: Array<string | undefined> = []
+    ) => {
+      let queriedEl: Element | null;
+      host.bool = true;
+      host.render();
+      queriedEl = host.root.firstElementChild;
+      assert.equal(queriedEl?.tagName, 'DIV');
+      assert.deepEqual(host.calls, [...initialCalls, 'DIV']);
+
+      host.bool = true;
+      host.render();
+      queriedEl = host.root.firstElementChild;
+      assert.equal(queriedEl?.tagName, 'DIV');
+      assert.deepEqual(host.calls, [...initialCalls, 'DIV']);
+
+      host.bool = false;
+      host.render();
+      queriedEl = host.root.firstElementChild;
+      assert.equal(queriedEl?.tagName, 'SPAN');
+      assert.deepEqual(host.calls, [...initialCalls, 'DIV', undefined, 'SPAN']);
+
+      host.bool = true;
+      host.render();
+      queriedEl = host.root.firstElementChild;
+      assert.equal(queriedEl?.tagName, 'DIV');
+      assert.deepEqual(host.calls, [
+        ...initialCalls,
+        'DIV',
+        undefined,
+        'SPAN',
+        undefined,
+        'DIV',
+      ]);
     };
-    const go = (x: boolean) =>
-      render(
-        x
-          ? html`<div ${ref(host.elCallback)}></div>`
-          : html`<span ${ref(host.elCallback)}></span>`,
-        container,
-        {host}
-      );
 
-    go(true);
-    queriedEl = container.firstElementChild;
-    assert.equal(queriedEl?.tagName, 'DIV');
-    assert.deepEqual(host.calls, ['DIV']);
-
-    go(true);
-    queriedEl = container.firstElementChild;
-    assert.equal(queriedEl?.tagName, 'DIV');
-    assert.deepEqual(host.calls, ['DIV']);
-
-    go(false);
-    queriedEl = container.firstElementChild;
-    assert.equal(queriedEl?.tagName, 'SPAN');
-    assert.deepEqual(host.calls, ['DIV', undefined, 'SPAN']);
-
-    go(true);
-    queriedEl = container.firstElementChild;
-    assert.equal(queriedEl?.tagName, 'DIV');
-    assert.deepEqual(host.calls, ['DIV', undefined, 'SPAN', undefined, 'DIV']);
+    // Test first instance
+    const host1 = new Host();
+    testRef(host1);
+
+    // Test second instance
+    const host2 = new Host();
+    testRef(host2);
+
+    // Test on first instance again
+    // (reset boolean to render SPAN, so we see an initial change
+    // back to DIV)
+    host1.bool = false;
+    host1.render();
+    // Add in an undefined call for the initial switch from SPAN back to DIV
+    testRef(host1, [...host1.calls, undefined]);
   });
 
   test('two refs', () => {