Try stronger typing

src/arcolObjectStore.ts

@@ -80,7 +80,7 @@ type GConstructor<T = {}> = new (...args: any[]) => T;
 // The fields defined in mixins could be local fields, and we need to aggregate the list of local fields.
 // That's why all mixins are required to declare their local fields in static properties so we can
 // aggregate them.
-type ArcolObjectBase<T extends object> = GConstructor<ArcolObject<any, any>> & {
+type ArcolObjectBase<T extends object> = GConstructor<ArcolObject<any, any, any>> & {
   LocalFieldsWithDefaults: T,
 };
 type MixinClass<T extends object> = GConstructor<any> & {
@@ -113,7 +113,8 @@ export function applyArcolObjectMixins<
  */
 export class ArcolObject<
   I extends string,
-  T extends ArcolObject<I, T>
+  S extends ArcolObjectFields,
+  T extends ArcolObject<I, any, T>
 > {
   /**
    * Stores the values of the fields of the object.
@@ -126,20 +127,20 @@ export class ArcolObject<
    * Setters should call `.set`, NOT mutate `fields` directly.
    */
   public readonly id: I;
-  protected fields: ArcolObjectFields;
+  protected fields: S;
   protected localFields: { [key: string]: true };
 
   constructor(
     protected store: ArcolObjectStore<I, T>,
-    protected node: LiveObject<FileFormat.ObjectShared<I> & ArcolObjectFields>,
+    protected node: LiveObject<FileFormat.ObjectShared<I> & S>,
     /**
      * List of fields that should not be persisted.
      */
     localFieldsWithDefaults: { [key: string]: unknown },
   ) {
     const { id, ...fields } = node.toObject();
     this.id = id;
-    this.fields = fields;
+    this.fields = fields as unknown as S;
     this.localFields = {};
     for (const key in localFieldsWithDefaults) {
       this.localFields[key] = true;
@@ -162,6 +163,14 @@ export class ArcolObject<
     return { ...this.fields };
   }
 
+  public get<K extends keyof S>(key: K): S[K] {
+    return this.fields[key];
+  }
+
+  public set<K extends keyof S>(key: K, value: S[K]) {
+    return this.setAny(key as string, value);
+  }
+
   public getAny(key: string): any {
     return this.fields[key];
   }
@@ -177,15 +186,15 @@ export class ArcolObject<
     if (!(key in this.localFields)) {
       this.node.set(key, value);
     }
-    this.fields[key] = value;
+    this.fields[key as keyof S] = value;
     this.store._internalOnFieldSet(this as any, key as string, oldValue, value);
   }
 
   /**
    * To be called from `ArcolObjectStore` only.
    */
   public _internalUpdateField(key: string, value: any) {
-    this.fields[key] = value;
+    this.fields[key as keyof S] = value;
   }
 
   /**
@@ -208,7 +217,7 @@ export type StoreName = Brand<string, "store-name">;
  * <T> is the type of the objects that we are putting in the store, probably a discriminated union
  * of all the different subtypes of objects.
  */
-export abstract class ArcolObjectStore<I extends string, T extends ArcolObject<I, T>> {
+export abstract class ArcolObjectStore<I extends string, T extends ArcolObject<I, any, T>> {
   protected objects = new Map<I, T>;
   protected listeners = new Set<ObjectListener<T>>();
 

src/elements/element.ts

@@ -7,18 +7,18 @@ import { Sketch } from "./sketch";
 
 export type Element = Sketch | Extrusion | Group | Level;
 
-interface Hideable {
+export interface Hideable {
   hidden: boolean;
 }
 
 export class HideableMixin {
   static LocalFieldsWithDefaults = { hidden: false } satisfies Partial<Hideable>;
 
   get hidden(): Hideable["hidden"] {
-    return (this as unknown as ArcolObject<ElementId, Element>).getFields().hidden;
+    return (this as unknown as ArcolObject<ElementId, any, Element>).getFields().hidden;
   }
 
   set hidden(value: Hideable["hidden"]) {
-    (this as unknown as ArcolObject<ElementId, Element>).setAny("hidden", value);
+    (this as unknown as ArcolObject<ElementId, any, Element>).setAny("hidden", value);
   }
 };

src/elements/extrusion.ts

@@ -5,7 +5,7 @@ import { ProjectStore } from "../project";
 import { Element, HideableMixin } from "./element";
 import { HierarchyMixin } from "../hierarchyMixin";
 
-export class Extrusion extends ArcolObject<ElementId, Element> {
+export class Extrusion extends ArcolObject<ElementId, FileFormat.Extrusion, Element> {
   static LocalFieldsWithDefaults = {
     ...HideableMixin.LocalFieldsWithDefaults,
   };
@@ -28,7 +28,7 @@ export class Extrusion extends ArcolObject<ElementId, Element> {
   }
 
   set height(value: FileFormat.Extrusion["height"]) {
-    this.setAny("height", value);
+    this.set("height", value);
   }
 }
 

src/elements/group.ts

@@ -5,7 +5,7 @@ import { ProjectStore } from "../project";
 import { Element, HideableMixin } from "./element";
 import { HierarchyMixin } from "../hierarchyMixin";
 
-export class Group extends ArcolObject<ElementId, Element> {
+export class Group extends ArcolObject<ElementId, FileFormat.Group, Element> {
   static LocalFieldsWithDefaults = {
     ...HideableMixin.LocalFieldsWithDefaults,
   };

src/elements/level.ts

@@ -5,7 +5,7 @@ import { ProjectStore } from "../project";
 import { Element, HideableMixin } from "./element";
 import { HierarchyMixin } from "../hierarchyMixin";
 
-export class Level extends ArcolObject<ElementId, Element> {
+export class Level extends ArcolObject<ElementId, FileFormat.Level, Element> {
   static LocalFieldsWithDefaults = {
     ...HideableMixin.LocalFieldsWithDefaults,
   };

src/elements/sketch.ts

@@ -5,7 +5,7 @@ import { Element, HideableMixin } from "./element";
 import { HierarchyMixin } from "../hierarchyMixin";
 import { ArcolObject, applyArcolObjectMixins } from "../arcolObjectStore";
 
-export class Sketch extends ArcolObject<ElementId, Element> {
+export class Sketch extends ArcolObject<ElementId, FileFormat.Sketch, Element> {
   static LocalFieldsWithDefaults = {
     ...HideableMixin.LocalFieldsWithDefaults,
   };
@@ -28,7 +28,7 @@ export class Sketch extends ArcolObject<ElementId, Element> {
   }
 
   set translate(value: FileFormat.Sketch["translate"]) {
-    this.setAny("translate", value);
+    this.set("translate", value);
   }
 
   get color(): FileFormat.Sketch["color"] {

src/hierarchyMixin.ts

@@ -10,7 +10,7 @@ import { FileFormat } from "./fileFormat";
  * - Implement HierarchyObserver to updated the cached values. It should probably be the first
  *   observer to run considering that subsequent observers are likely to read the children list.
  */
-export class HierarchyMixin<I extends string, T extends ArcolObject<I, T> & HierarchyMixin<I, T>> {
+export class HierarchyMixin<I extends string, T extends ArcolObject<I, any, T> & HierarchyMixin<I, T>> {
   static MixinLocalFieldsWithDefaults = {};
 
   /**
@@ -113,15 +113,15 @@ export class HierarchyMixin<I extends string, T extends ArcolObject<I, T> & Hier
    * We could make this a symbol private to this field to enforce it more strongly, but I think it's
    * not worth the readability hit.
    */
-  public _internalAddChild(child: ArcolObject<I, T>) {
+  public _internalAddChild(child: ArcolObject<I, any, T>) {
     this.childrenSet.add(child.id);
     this.cachedChildren = null;
   }
 
   /**
    * To be called from `ArcolObjectStore` only.
    */
-  public _internalRemoveChild(child: ArcolObject<I, T>) {
+  public _internalRemoveChild(child: ArcolObject<I, any, T>) {
     this.childrenSet.delete(child.id);
     this.cachedChildren = null;
   }
@@ -139,7 +139,7 @@ export class HierarchyMixin<I extends string, T extends ArcolObject<I, T> & Hier
  */
 export class HierarchyObserver<
     I extends string,
-    T extends ArcolObject<I, T> & HierarchyMixin<I, T>
+    T extends ArcolObject<I, any, T> & HierarchyMixin<I, T>
 > implements ObjectObserver<T> {
 
   constructor(private store: ArcolObjectStore<I, T>) { }

src/project.ts

@@ -1,7 +1,7 @@
 import { LiveObject, Room } from "@liveblocks/client";
 import { ElementId, FileFormat } from "./fileFormat";
 import { Sketch } from "./elements/sketch";
-import { ArcolObjectStore, ObjectChange, ObjectListener, ObjectObserver, StoreName } from "./arcolObjectStore";
+import { ArcolObjectFields, ArcolObjectStore, ChangeOrigin, ObjectChange, ObjectObserver, StoreName } from "./arcolObjectStore";
 import { Extrusion } from "./elements/extrusion";
 import { Group } from "./elements/group";
 import { Level } from "./elements/level";
@@ -10,9 +10,70 @@ import { generateKeyBetween } from "fractional-indexing";
 import { HierarchyObserver } from "./hierarchyMixin";
 import { ChangeManager } from "./changeManager";
 
-export type ElementListener = ObjectListener<Element>;
 export type ElementObserver = ObjectObserver<Element>;
 
+/**
+ * A more strongly typed version of {@link ObjectChange}.
+ *
+ * It uses mapped types and index access types to produce a union of all the possible changes for
+ * a given object type.
+ *
+ * e.g.
+ *   TypedObjectChange<{ k1: V1, k2: V2, ... }> =
+ *     | { type: "create" | "delete" }
+ *     | { type: "update", property: "k1", oldValue: V1 }
+ *     | { type: "update", property: "k2", oldValue: V2 }
+ *     ...
+ */
+type TypedObjectChange<T extends ArcolObjectFields> =
+  | { type: "create" | "delete" }
+  | { [K in keyof T]: { type: "update", property: K, oldValue: T[K] } }[keyof T]
+
+/**
+ * A more strongly typed version of {@link ObjectListener} for `Element`. The argument to the
+ * callback is a single `params` object which repeats the `type` from `obj.type`. This allows
+ * narrowing both `obj` and `change` based on the `type` of the `Element`.
+ *
+ * In more concrete terms, it allows the following to give the desired types:
+ * ```
+ *   if (params.type === "someElementType") {
+ *     if (params.change.type === "update") {
+ *       // params.change.property is a union of the possible fields of `someElementType`
+ *       if (params.change.property) {
+ *         // params.change.value is the type of the field [params.change.property] in `someElementType`
+ *       }
+ *     }
+ *   }
+ * ```
+ *
+ * The weird `T extends any` syntax is called "Distributive Conditional Types" and allows writing
+ * generic that maps a union type to a different union of types.
+ * https://www.typescriptlang.org/docs/handbook/2/conditional-types.html#distributive-conditional-types
+ */
+type ElementListener<T extends Element = Element> = (params: T extends any ? {
+  type: T["type"],
+  obj: T,
+  change: TypedObjectChange<ReturnType<T["getFields"]>>,
+  origin: ChangeOrigin,
+} : never) => void
+
+// This is just here as a test that our TypeScript types are able to perform the desired narrowing.
+function foo(l: ElementListener) {}
+foo((params) => {
+    if (params.type === "sketch") {
+        const obj: Sketch = params.obj;
+        if (params.change.type === "update") {
+          const property: "id" | "type" | "parentId" | "parentIndex" | "translate" | "color" = params.change.property;
+          if (params.change.property === "translate") {
+            const old: FileFormat.Vec3 = params.change.oldValue
+            void old;
+          }
+          void property;
+        }
+    }
+})
+
+
 class DeleteEmptyExtrusionObserver implements ElementObserver {
   private elementsToCheck = new Set<ElementId>();
 
@@ -100,6 +161,12 @@ export class ProjectStore extends ArcolObjectStore<ElementId, Element> {
     });
   }
 
+  public subscribeElementChange(listener: ElementListener): () => void {
+    return this.subscribeObjectChange((obj, change, origin) => {
+      listener({ type: obj.type, obj, change, origin } as any);
+    });
+  }
+
   public removeElement(element: Element) {
     const removeRecursive = (element: Element) => {
       if (!this.objects.has(element.id)) {

src/relationsStore.ts

@@ -13,7 +13,7 @@ export class Relation<
   IA extends string,
   IB extends string,
   O extends Relation<IA, IB, O>
-> extends ArcolObject<`${IA}<>${IB}`, O> {
+> extends ArcolObject<`${IA}<>${IB}`, any, O> {
   public readonly keyA: IA;
   public readonly keyB: IB;
 
@@ -52,9 +52,9 @@ export abstract class RelationsStore<
 
   // The owner of this class should put these observers in the respective stores whose objects
   // are being related.
-  public readonly observerA: ObjectObserver<ArcolObject<IA, any>> =
+  public readonly observerA: ObjectObserver<ArcolObject<IA, any, any>> =
     { onChange: this.onObjectChangeA.bind(this) };
-  public readonly observerB: ObjectObserver<ArcolObject<IB, any>> =
+  public readonly observerB: ObjectObserver<ArcolObject<IB, any, any>> =
     { onChange: this.onObjectChangeB.bind(this) };
 
   protected initialize() {
@@ -120,7 +120,7 @@ export abstract class RelationsStore<
   /**
    * Cleans up relations related to an object instance of type A when it is deleted.
    */
-  private onObjectChangeA(obj: ArcolObject<IA, any>, change: ObjectChange) {
+  private onObjectChangeA(obj: ArcolObject<IA, any, any>, change: ObjectChange) {
     if (change.type === "delete") {
       const relations = this.relationsFromA.get(obj.id);
       if (relations) {
@@ -136,7 +136,7 @@ export abstract class RelationsStore<
   /**
    * Cleans up relations related to an object instance of type B when it is deleted.
    */
-  private onObjectChangeB(obj: ArcolObject<IB, any>, change: ObjectChange) {
+  private onObjectChangeB(obj: ArcolObject<IB, any, any>, change: ObjectChange) {
     if (change.type === "delete") {
       const relations = this.relationsFromB.get(obj.id);
       if (relations) {

src/undoRedo.ts

@@ -2,7 +2,7 @@ import { ArcolObject, ArcolObjectFields, ArcolObjectStore, ChangeOrigin, ObjectC
 import { Editor } from "./editor";
 import { ElementSelection, useAppState } from "./global";
 
-type AnyObject = ArcolObject<any, any>;
+type AnyObject = ArcolObject<any, any, any>;
 type AnyObjectStore = ArcolObjectStore<any, any>;
 
 // When "op" is "create" or "delete", properties includes all of the object.