@loro-extended/change 2.0.0 → 4.0.0

package/src/typed-refs/tree.test.ts

@@ -0,0 +1,384 @@
+import { describe, expect, it } from "vitest"
+import { change } from "../functional-helpers.js"
+import { Shape } from "../shape.js"
+import { createTypedDoc } from "../typed-doc.js"
+
+// Define the state machine schema from the plan
+const StateNodeDataShape = Shape.struct({
+  name: Shape.text(),
+  facts: Shape.record(Shape.plain.any()),
+  rules: Shape.list(
+    Shape.plain.struct({
+      name: Shape.plain.string(),
+      rego: Shape.plain.string(),
+      description: Shape.plain.string().nullable(),
+    }),
+  ),
+})
+
+const ResmSchema = Shape.doc({
+  states: Shape.tree(StateNodeDataShape),
+  currentPath: Shape.list(Shape.plain.string()),
+  input: Shape.record(Shape.plain.any()),
+})
+
+describe("TreeRef", () => {
+  describe("basic operations", () => {
+    it("should create a root node with typed data", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        expect(root).toBeDefined()
+        expect(root.id).toBeDefined()
+      })
+    })
+
+    it("should create a root node and set data", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        // Set data after creation
+        root.data.name.insert(0, "idle")
+      })
+
+      // Verify the data was set
+      const json = typedDoc.toJSON()
+      expect(json.states).toHaveLength(1)
+      expect(json.states[0].data.name).toBe("idle")
+    })
+
+    it("should create child nodes", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "idle")
+
+        const child = root.createNode()
+        child.data.name.insert(0, "running")
+      })
+
+      const json = typedDoc.toJSON()
+      expect(json.states).toHaveLength(1)
+      expect(json.states[0].data.name).toBe("idle")
+      expect(json.states[0].children).toHaveLength(1)
+      expect(json.states[0].children[0].data.name).toBe("running")
+    })
+
+    it("should access node.data with type safety", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "idle")
+
+        // Access nested list container
+        root.data.rules.push({
+          name: "rule1",
+          rego: "package test",
+          description: null,
+        })
+      })
+
+      const json = typedDoc.toJSON()
+      expect(json.states[0].data.name).toBe("idle")
+      expect(json.states[0].data.rules).toEqual([
+        { name: "rule1", rego: "package test", description: null },
+      ])
+    })
+
+    it("should access record containers in node.data", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "idle")
+
+        // Access nested record container - use set method
+        const facts = root.data.facts
+        facts.set("key1", "value1")
+      })
+
+      const json = typedDoc.toJSON()
+      expect(json.states[0].data.name).toBe("idle")
+      expect(json.states[0].data.facts).toEqual({ key1: "value1" })
+    })
+  })
+
+  describe("tree navigation", () => {
+    it("should get roots", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        draft.states.createNode()
+        draft.states.createNode()
+      })
+
+      change(typedDoc, draft => {
+        const roots = draft.states.roots()
+        expect(roots).toHaveLength(2)
+      })
+    })
+
+    it("should get all nodes", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.createNode()
+        root.createNode()
+      })
+
+      change(typedDoc, draft => {
+        const nodes = draft.states.nodes()
+        expect(nodes).toHaveLength(3) // 1 root + 2 children
+      })
+    })
+
+    it("should navigate parent/children relationships", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let rootId: string | undefined
+      let childId: string | undefined
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "root")
+        rootId = root.id
+
+        const child = root.createNode()
+        child.data.name.insert(0, "child")
+        childId = child.id
+      })
+
+      change(typedDoc, draft => {
+        const root = draft.states.getNodeByID(rootId as string)
+        expect(root).toBeDefined()
+        expect(root?.children()).toHaveLength(1)
+
+        const child = draft.states.getNodeByID(childId as string)
+        expect(child).toBeDefined()
+        expect(child?.parent()?.id).toBe(rootId)
+      })
+    })
+
+    it("should get node by ID", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let nodeId: string | undefined
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "test")
+        nodeId = root.id
+      })
+
+      change(typedDoc, draft => {
+        const node = draft.states.getNodeByID(nodeId as string)
+        expect(node).toBeDefined()
+        expect(node?.data.name.toString()).toBe("test")
+      })
+    })
+
+    it("should check if node exists with has()", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let nodeId: string | undefined
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        nodeId = root.id
+      })
+
+      change(typedDoc, draft => {
+        expect(draft.states.has(nodeId as string)).toBe(true)
+        expect(draft.states.has("0@999" as any)).toBe(false)
+      })
+    })
+  })
+
+  describe("tree mutations", () => {
+    it("should delete a node", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let nodeId: string | undefined
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        nodeId = root.id
+      })
+
+      change(typedDoc, draft => {
+        draft.states.delete(nodeId as string)
+      })
+
+      const json = typedDoc.toJSON()
+      expect(json.states).toHaveLength(0)
+    })
+
+    it("should move nodes between parents", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let root1Id: string | undefined
+      let root2Id: string | undefined
+      let childId: string | undefined
+
+      change(typedDoc, draft => {
+        const root1 = draft.states.createNode()
+        root1.data.name.insert(0, "root1")
+        root1Id = root1.id
+
+        const root2 = draft.states.createNode()
+        root2.data.name.insert(0, "root2")
+        root2Id = root2.id
+
+        const child = root1.createNode()
+        child.data.name.insert(0, "child")
+        childId = child.id
+      })
+
+      // Move child from root1 to root2
+      change(typedDoc, draft => {
+        const child = draft.states.getNodeByID(childId as string)
+        const root2 = draft.states.getNodeByID(root2Id as string)
+        if (child && root2) {
+          child.move(root2)
+        }
+      })
+
+      change(typedDoc, draft => {
+        const root1 = draft.states.getNodeByID(root1Id as string)
+        const root2 = draft.states.getNodeByID(root2Id as string)
+
+        expect(root1?.children()).toHaveLength(0)
+        expect(root2?.children()).toHaveLength(1)
+      })
+    })
+  })
+
+  describe("serialization", () => {
+    it("should serialize tree to nested JSON with toJSON()", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "root")
+
+        const child1 = root.createNode()
+        child1.data.name.insert(0, "child1")
+
+        const child2 = root.createNode()
+        child2.data.name.insert(0, "child2")
+
+        const grandchild = child1.createNode()
+        grandchild.data.name.insert(0, "grandchild")
+      })
+
+      const json = typedDoc.toJSON()
+      expect(json.states).toHaveLength(1)
+      expect(json.states[0].data.name).toBe("root")
+      expect(json.states[0].children).toHaveLength(2)
+      expect(json.states[0].children[0].children).toHaveLength(1)
+    })
+
+    it("should serialize tree to flat array with toArray()", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "root")
+
+        const child = root.createNode()
+        child.data.name.insert(0, "child")
+      })
+
+      change(typedDoc, draft => {
+        const array = draft.states.toArray()
+        expect(array).toHaveLength(2)
+        // Each item should have id, parent, index, fractionalIndex, data
+        expect(array[0]).toHaveProperty("id")
+        expect(array[0]).toHaveProperty("parent")
+        expect(array[0]).toHaveProperty("index")
+        expect(array[0]).toHaveProperty("data")
+      })
+    })
+  })
+
+  describe("fractional index", () => {
+    it("should enable fractional index", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        draft.states.enableFractionalIndex(8)
+        const root = draft.states.createNode()
+        expect(root.fractionalIndex()).toBeDefined()
+      })
+    })
+
+    it("should get node index among siblings", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        const child1 = root.createNode()
+        const child2 = root.createNode()
+
+        expect(child1.index()).toBe(0)
+        expect(child2.index()).toBe(1)
+      })
+    })
+  })
+
+  describe("absorbPlainValues", () => {
+    it("should propagate absorbPlainValues to all nodes", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        root.data.name.insert(0, "test")
+        root.data.rules.push({
+          name: "rule1",
+          rego: "package test",
+          description: null,
+        })
+
+        // absorbPlainValues is called automatically at the end of change()
+      })
+
+      // Verify data was persisted
+      const json = typedDoc.toJSON()
+      expect(json.states[0].data.name).toBe("test")
+      expect(json.states[0].data.rules).toEqual([
+        { name: "rule1", rego: "package test", description: null },
+      ])
+    })
+  })
+
+  describe("node deletion tracking", () => {
+    it("should track deleted nodes with isDeleted()", () => {
+      const typedDoc = createTypedDoc(ResmSchema)
+
+      let nodeId: string | undefined
+
+      change(typedDoc, draft => {
+        const root = draft.states.createNode()
+        nodeId = root.id
+      })
+
+      change(typedDoc, draft => {
+        const node = draft.states.getNodeByID(nodeId as string)
+        if (node) {
+          expect(node.isDeleted()).toBe(false)
+          draft.states.delete(nodeId as string)
+        }
+      })
+
+      // After deletion, Loro's has() still returns true (node exists in history)
+      // but isDeleted() returns true and the node won't appear in toJSON()
+      const json = typedDoc.toJSON()
+      expect(json.states).toHaveLength(0)
+    })
+  })
+})

package/src/typed-refs/tree.ts

@@ -1,40 +1,268 @@
-import type { TreeContainerShape } from "../shape.js"
+import type { LoroDoc, LoroTree, LoroTreeNode, TreeID } from "loro-crdt"
+import type {
+  StructContainerShape,
+  TreeContainerShape,
+  TreeNodeJSON,
+} from "../shape.js"
 import type { Infer } from "../types.js"
-import { TypedRef } from "./base.js"
+import { TreeNodeRef } from "./tree-node.js"
 
-// Tree typed ref
-export class TreeRef<T extends TreeContainerShape> extends TypedRef<T> {
-  absorbPlainValues() {
-    // TODO(duane): implement for trees
+/**
+ * Parameters for creating a TreeRef.
+ */
+export interface TreeRefParams<DataShape extends StructContainerShape> {
+  shape: TreeContainerShape<DataShape>
+  placeholder?: never[]
+  getContainer: () => LoroTree
+  autoCommit?: boolean
+  batchedMutation?: boolean
+  getDoc?: () => LoroDoc
+}
+
+/**
+ * Typed ref for tree (forest) containers.
+ * Wraps LoroTree with type-safe access to node metadata.
+ *
+ * @example
+ * ```typescript
+ * const StateNodeDataShape = Shape.struct({
+ *   name: Shape.text(),
+ *   facts: Shape.record(Shape.plain.any()),
+ * })
+ *
+ * doc.$.change(draft => {
+ *   const root = draft.states.createNode({ name: "idle", facts: {} })
+ *   const child = root.createNode({ name: "running", facts: {} })
+ *   child.data.name = "active"
+ * })
+ * ```
+ */
+export class TreeRef<DataShape extends StructContainerShape> {
+  private nodeCache = new Map<TreeID, TreeNodeRef<DataShape>>()
+  private _cachedContainer?: LoroTree
+  protected _params: TreeRefParams<DataShape>
+
+  constructor(params: TreeRefParams<DataShape>) {
+    this._params = params
+  }
+
+  protected get shape(): TreeContainerShape<DataShape> {
+    return this._params.shape
+  }
+
+  protected get container(): LoroTree {
+    if (!this._cachedContainer) {
+      this._cachedContainer = this._params.getContainer()
+    }
+    return this._cachedContainer
+  }
+
+  protected get dataShape(): DataShape {
+    return this.shape.shape
+  }
+
+  protected get autoCommit(): boolean {
+    return !!this._params.autoCommit
+  }
+
+  protected get batchedMutation(): boolean {
+    return !!this._params.batchedMutation
+  }
+
+  protected get doc(): LoroDoc | undefined {
+    return this._params.getDoc?.()
+  }
+
+  /**
+   * Commits changes if autoCommit is enabled.
+   */
+  protected commitIfAuto(): void {
+    if (this.autoCommit && this.doc) {
+      this.doc.commit()
+    }
+  }
+
+  /**
+   * Absorb plain values from all cached nodes.
+   * Called before committing changes to ensure all pending values are written.
+   */
+  absorbPlainValues(): void {
+    for (const nodeRef of this.nodeCache.values()) {
+      nodeRef.absorbPlainValues()
+    }
+  }
+
+  /**
+   * Get or create a TreeNodeRef for the given LoroTreeNode.
+   * Uses caching to ensure the same TreeNodeRef is returned for the same node.
+   */
+  getOrCreateNodeRef(node: LoroTreeNode): TreeNodeRef<DataShape> {
+    const id = node.id
+    let nodeRef = this.nodeCache.get(id)
+    if (!nodeRef) {
+      nodeRef = new TreeNodeRef({
+        node,
+        dataShape: this.dataShape,
+        treeRef: this,
+        autoCommit: this.autoCommit,
+        batchedMutation: this.batchedMutation,
+        getDoc: this._params.getDoc,
+      })
+      this.nodeCache.set(id, nodeRef)
+    }
+    return nodeRef
+  }
+
+  /**
+   * Create a new root node with optional initial data.
+   *
+   * @param initialData - Optional partial data to initialize the node with
+   * @returns The created TreeNodeRef
+   */
+  createNode(initialData?: Partial<Infer<DataShape>>): TreeNodeRef<DataShape> {
+    const loroNode = this.container.createNode()
+    const nodeRef = this.getOrCreateNodeRef(loroNode)
+
+    // Initialize data if provided
+    if (initialData) {
+      for (const [key, value] of Object.entries(initialData)) {
+        ;(nodeRef.data as any)[key] = value
+      }
+    }
+
+    this.commitIfAuto()
+    return nodeRef
+  }
+
+  /**
+   * Get all root nodes (nodes without parents).
+   * Returns nodes in their fractional index order.
+   */
+  roots(): TreeNodeRef<DataShape>[] {
+    return this.container.roots().map(node => this.getOrCreateNodeRef(node))
   }
 
-  toJSON(): Infer<T> {
-    // TODO(duane): implement proper tree serialization
-    return this.container.toJSON() as Infer<T>
+  /**
+   * Get all nodes in the tree (unordered).
+   * Includes all nodes, not just roots.
+   */
+  nodes(): TreeNodeRef<DataShape>[] {
+    return this.container.nodes().map(node => this.getOrCreateNodeRef(node))
   }
 
-  createNode(parent?: any, index?: number): any {
-    this.assertMutable()
-    return this.container.createNode(parent, index)
+  /**
+   * Get a node by its TreeID.
+   *
+   * @param id - The TreeID of the node to find
+   * @returns The TreeNodeRef if found, undefined otherwise
+   */
+  getNodeByID(id: TreeID): TreeNodeRef<DataShape> | undefined {
+    // Check cache first
+    const cached = this.nodeCache.get(id)
+    if (cached) return cached
+
+    // Check if node exists in tree
+    if (!this.container.has(id)) return undefined
+
+    // Find the node in the tree's nodes
+    const nodes = this.container.nodes()
+    const node = nodes.find(n => n.id === id)
+    if (!node) return undefined
+
+    return this.getOrCreateNodeRef(node)
+  }
+
+  /**
+   * Check if a node with the given ID exists in the tree.
+   */
+  has(id: TreeID): boolean {
+    return this.container.has(id)
   }
 
-  move(target: any, parent?: any, index?: number): void {
-    this.assertMutable()
-    this.container.move(target, parent, index)
+  /**
+   * Delete a node and all its descendants.
+   * Also removes the node from the cache.
+   *
+   * @param target - The TreeID or TreeNodeRef to delete
+   */
+  delete(target: TreeID | TreeNodeRef<DataShape>): void {
+    const id = typeof target === "string" ? target : target.id
+    this.container.delete(id)
+    // Remove from cache
+    this.nodeCache.delete(id)
+    this.commitIfAuto()
   }
 
-  delete(target: any): void {
-    this.assertMutable()
-    this.container.delete(target)
+  /**
+   * Enable fractional index generation for ordering.
+   *
+   * @param jitter - Optional jitter value to avoid conflicts (0 = no jitter)
+   */
+  enableFractionalIndex(jitter = 0): void {
+    this.container.enableFractionalIndex(jitter)
   }
 
-  has(target: any): boolean {
-    return this.container.has(target)
+  /**
+   * Serialize the tree to a nested JSON structure.
+   * Each node includes its data and children recursively.
+   */
+  toJSON(): TreeNodeJSON<DataShape>[] {
+    // Use Loro's native toJSON which returns nested structure
+    const nativeJson = this.container.toJSON() as any[]
+    return this.transformNativeJson(nativeJson)
   }
 
-  getNodeByID(id: any): any {
-    return this.container.getNodeByID
-      ? this.container.getNodeByID(id)
-      : undefined
+  /**
+   * Transform Loro's native JSON format to our typed format.
+   */
+  private transformNativeJson(nodes: any[]): TreeNodeJSON<DataShape>[] {
+    return nodes.map(node => ({
+      id: node.id as TreeID,
+      parent: node.parent as TreeID | null,
+      index: node.index as number,
+      fractionalIndex: node.fractional_index as string,
+      data: node.meta as Infer<DataShape>,
+      children: this.transformNativeJson(node.children || []),
+    }))
+  }
+
+  /**
+   * Get a flat array representation of all nodes.
+   * Flattens the nested tree structure into a single array.
+   */
+  toArray(): Array<{
+    id: TreeID
+    parent: TreeID | null
+    index: number
+    fractionalIndex: string
+    data: Infer<DataShape>
+  }> {
+    const result: Array<{
+      id: TreeID
+      parent: TreeID | null
+      index: number
+      fractionalIndex: string
+      data: Infer<DataShape>
+    }> = []
+
+    // Flatten the nested structure
+    const flattenNodes = (nodes: any[]) => {
+      for (const node of nodes) {
+        result.push({
+          id: node.id as TreeID,
+          parent: node.parent as TreeID | null,
+          index: node.index as number,
+          fractionalIndex: node.fractional_index as string,
+          data: node.meta as Infer<DataShape>,
+        })
+        if (node.children && node.children.length > 0) {
+          flattenNodes(node.children)
+        }
+      }
+    }
+
+    const nativeJson = this.container.toJSON() as any[]
+    flattenNodes(nativeJson)
+    return result
   }
 }