Keep the value and the tag of a tracked field in one cell

packages/@glimmer/validator/lib/tracked-data.ts

@@ -1,36 +1,89 @@
-import { dirtyTagFor, tagFor } from './meta';
+import { DEBUG } from '@glimmer/env';
+import type { UpdatableTag } from '@glimmer/interfaces';
+
+import { debug } from './debug';
+import { tagFor } from './meta';
 import { consumeTag } from './tracking';
+import { unwrap } from './utils';
+import { DIRTY_TAG } from './validators';
 
 export type Getter<T, K extends keyof T> = (self: T) => T[K] | undefined;
 export type Setter<T, K extends keyof T> = (self: T, value: T[K]) => void;
 
+/**
+ * The value and the tag of one field of one instance. A read or a write
+ * finds both with one WeakMap lookup.
+ *
+ * The tag is the one that the tag registry has for the field, so `tagFor`
+ * and `dirtyTagFor` work on the same tag as the field.
+ */
+interface TrackedCell<V> {
+  value: V;
+  tag: UpdatableTag;
+  initialized: boolean;
+}
+
+/**
+ * All cells have one hidden class, and V8 records which kind of value the
+ * `value` field held so far. The first value of another kind makes V8 throw
+ * away the optimized code that reads the field.
+ *
+ * A number and then `undefined` make the field general from the first cell
+ * on, before that code is optimized.
+ */
+function createCell<V>(tag: UpdatableTag, initialized: boolean): TrackedCell<V | undefined> {
+  let cell: TrackedCell<unknown> = { value: 0, tag, initialized };
+
+  cell.value = undefined;
+
+  return cell as TrackedCell<V | undefined>;
+}
+
 export function trackedData<T extends object, K extends keyof T>(
   key: K,
   initializer?: (this: T) => T[K]
 ): { getter: Getter<T, K>; setter: Setter<T, K> } {
-  let values = new WeakMap<T, T[K]>();
+  let cells = new WeakMap<T, TrackedCell<T[K] | undefined>>();
   let hasInitializer = typeof initializer === 'function';
 
+  function cellFor(self: T): TrackedCell<T[K] | undefined> {
+    let cell = cells.get(self);
+
+    if (cell === undefined) {
+      // Other code can ask the registry for this tag before the first read
+      // or write of the field, so the cell takes the tag from there.
+      cell = createCell(tagFor(self, key) as UpdatableTag, !hasInitializer);
+      cells.set(self, cell);
+    }
+
+    return cell;
+  }
+
   function getter(self: T) {
-    consumeTag(tagFor(self, key));
+    const cell = cellFor(self);
 
-    let value;
+    consumeTag(cell.tag);
 
     // If the field has never been initialized, we should initialize it
-    if (hasInitializer && !values.has(self)) {
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion -- @fixme
-      value = initializer!.call(self);
-      values.set(self, value);
-    } else {
-      value = values.get(self);
+    if (!cell.initialized) {
+      cell.initialized = true;
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion -- guarded by initialized
+      cell.value = initializer!.call(self);
     }
 
-    return value;
+    return cell.value;
   }
 
   function setter(self: T, value: T[K]): void {
-    dirtyTagFor(self, key);
-    values.set(self, value);
+    const cell = cellFor(self);
+
+    if (DEBUG) {
+      unwrap(debug.assertTagNotConsumed)(cell.tag, self, key);
+    }
+
+    DIRTY_TAG(cell.tag, true);
+    cell.initialized = true;
+    cell.value = value;
   }
 
   return { getter, setter };

packages/@glimmer/validator/test/tracking-test.ts

@@ -7,11 +7,13 @@ import {
   createTag,
   debug,
   dirtyTag,
+  dirtyTagFor,
   endTrackFrame,
   getValue,
   isConst,
   isTracking,
   resetTracking,
+  tagFor,
   track,
   trackedData,
   untrack,
@@ -571,6 +573,50 @@ module('@glimmer/validator: tracking', () => {
       assert.notOk(validateTag(tag, snapshot));
     });
 
+    test('a write dirties the tag that tagFor gave before the first read', (assert) => {
+      class Foo {
+        foo = 123;
+      }
+
+      let { setter } = trackedData<Foo, keyof Foo>('foo');
+
+      let foo = new Foo();
+      let tag = tagFor(foo, 'foo');
+      let snapshot = valueForTag(tag);
+
+      setter(foo, 456);
+      assert.notOk(validateTag(tag, snapshot));
+    });
+
+    test('a read consumes the tag that tagFor gave before it', (assert) => {
+      class Foo {
+        foo = 123;
+      }
+
+      let { getter } = trackedData<Foo, keyof Foo>('foo');
+
+      let foo = new Foo();
+      let registered = tagFor(foo, 'foo');
+      let tag = track(() => getter(foo));
+
+      assert.strictEqual(tag, registered);
+    });
+
+    test('dirtyTagFor dirties the tag that a read consumed', (assert) => {
+      class Foo {
+        foo = 123;
+      }
+
+      let { getter } = trackedData<Foo, keyof Foo>('foo');
+
+      let foo = new Foo();
+      let tag = track(() => getter(foo));
+      let snapshot = valueForTag(tag);
+
+      dirtyTagFor(foo, 'foo');
+      assert.notOk(validateTag(tag, snapshot));
+    });
+
     if (DEBUG) {
       test('it errors when attempting to update a value already consumed in the same transaction', (assert) => {
         class Foo {