refactor: update and maintain __tres parent/objects graph (#741)

src/core/nodeOps.test.ts

@@ -255,6 +255,55 @@ describe('nodeOps', () => {
         expect(parent.children.length).toBe(0)
       }
     })
+
+    it('adds a material to parent.__tres.objects', () => {
+      const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+      const material = nodeOps.createElement('MeshNormalMaterial')
+      nodeOps.insert(material, parent)
+      expect(parent.__tres.objects.map(child => child.uuid)).toStrictEqual([material.uuid])
+    })
+
+    it('adds a geometry to parent.__tres.objects', () => {
+      const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+      const geometry = nodeOps.createElement('BoxGeometry')
+      nodeOps.insert(geometry, parent)
+      expect(parent.__tres.objects.map(child => child.uuid)).toStrictEqual([geometry.uuid])
+    })
+
+    it('adds a fog to parent.__tres.objects', () => {
+      const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+      const fog = nodeOps.createElement('Fog')
+      nodeOps.insert(fog, parent)
+      expect(parent.__tres.objects.map(child => child.uuid)).toStrictEqual([fog.uuid])
+    })
+
+    it('adds parent to child.__tres.parent', () => {
+      const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+      const material = nodeOps.createElement('MeshNormalMaterial')
+      const geometry = nodeOps.createElement('BoxGeometry')
+      const fog = nodeOps.createElement('Fog')
+      nodeOps.insert(material, parent)
+      nodeOps.insert(geometry, parent)
+      nodeOps.insert(fog, parent)
+      expect(material.__tres.parent).toBe(parent)
+      expect(geometry.__tres.parent).toBe(parent)
+      expect(fog.__tres.parent).toBe(parent)
+    })
+
+    it('adds non-Object3D children to parent.__tres.objects, but no more than once', () => {
+      const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+      const material = nodeOps.createElement('MeshNormalMaterial')
+      const geometry = nodeOps.createElement('BoxGeometry')
+      const fog = nodeOps.createElement('Fog')
+      nodeOps.insert(material, parent)
+      nodeOps.insert(geometry, parent)
+      nodeOps.insert(fog, parent)
+      expect(parent.__tres.objects.length).toBe(3)
+      const objectSet = new Set(parent.__tres.objects)
+      expect(objectSet.has(material)).toBe(true)
+      expect(objectSet.has(geometry)).toBe(true)
+      expect(objectSet.has(fog)).toBe(true)
+    })
   })
 
   describe('remove', () => {
@@ -515,6 +564,35 @@ describe('nodeOps', () => {
         expect(grandChild1.parent.uuid).toBe(primitiveMesh.uuid)
       })
     })
+    describe('in the __tres parent-object graph', () => {
+      it('removes parent-object relationship when object is removed', () => {
+        const parent = nodeOps.createElement('Mesh', undefined, undefined, {})
+        const material = nodeOps.createElement('MeshNormalMaterial')
+        const geometry = nodeOps.createElement('BoxGeometry')
+        const fog = nodeOps.createElement('Fog')
+        nodeOps.insert(material, parent)
+        nodeOps.insert(geometry, parent)
+        nodeOps.insert(fog, parent)
+        expect(material.__tres.parent).toBe(parent)
+        expect(geometry.__tres.parent).toBe(parent)
+        expect(fog.__tres.parent).toBe(parent)
+
+        nodeOps.remove(fog)
+        expect(fog.__tres.parent).toBe(null)
+        expect(parent.__tres.objects.length).toBe(2)
+        expect(parent.__tres.objects.includes(fog)).toBe(false)
+
+        nodeOps.remove(material)
+        expect(material.__tres.parent).toBe(null)
+        expect(parent.__tres.objects.length).toBe(1)
+        expect(parent.__tres.objects.includes(material)).toBe(false)
+
+        nodeOps.remove(geometry)
+        expect(geometry.__tres.parent).toBe(null)
+        expect(parent.__tres.objects.length).toBe(0)
+        expect(parent.__tres.objects.includes(geometry)).toBe(false)
+      })
+    })
   })
 
   describe('patchProp', () => {

src/core/nodeOps.ts

@@ -2,16 +2,12 @@ import type { RendererOptions } from 'vue'
 import { BufferAttribute, Object3D } from 'three'
 import type { TresContext } from '../composables'
 import { useLogger } from '../composables'
-import { deepArrayEqual, disposeObject3D, isHTMLTag, kebabToCamel } from '../utils'
-import type { InstanceProps, TresObject, TresObject3D } from '../types'
+import { deepArrayEqual, disposeObject3D, filterInPlace, isHTMLTag, kebabToCamel } from '../utils'
+import type { InstanceProps, TresInstance, TresObject, TresObject3D } from '../types'
 import * as is from '../utils/is'
+import { invalidateInstance, noop, prepareTresInstance } from '../utils/nodeOpsUtils'
 import { catalogue } from './catalogue'
 
-function noop(fn: string): any {
-  // eslint-disable-next-line no-unused-expressions
-  fn
-}
-
 const { logError } = useLogger()
 
 const supportedPointerEvents = [
@@ -31,16 +27,6 @@ const supportedPointerEvents = [
   'onWheel',
 ]
 
-export function invalidateInstance(instance: TresObject) {
-  const ctx = instance?.__tres?.root
-
-  if (!ctx) { return }
-
-  if (ctx.render && ctx.render.canBeInvalidated.value) {
-    ctx.invalidate()
-  }
-}
-
 export const nodeOps: (context: TresContext) => RendererOptions<TresObject, TresObject | null> = (context) => {
   const scene = context.scene.value
 
@@ -88,14 +74,14 @@ export const nodeOps: (context: TresContext) => RendererOptions<TresObject, Tres
       else if (instance.isBufferGeometry) { instance.attach = 'geometry' }
     }
 
-    instance.__tres = {
+    instance = prepareTresInstance(instance, {
       ...instance.__tres,
       type: name,
       memoizedProps: props,
       eventCount: 0,
       disposable: true,
       primitive: tag === 'primitive',
-    }
+    }, context)
 
     // determine whether the material was passed via prop to
     // prevent it's disposal when node is removed later in it's lifecycle
@@ -108,22 +94,25 @@ export const nodeOps: (context: TresContext) => RendererOptions<TresObject, Tres
 
   function insert(child: TresObject, parent: TresObject) {
     if (!child) { return }
-
-    if (child.__tres) {
-      child.__tres.root = context
-    }
+    const childInstance: TresInstance = (child.__tres ? child as TresInstance : prepareTresInstance(child, {}))
 
     const parentObject = parent || scene
 
     context.registerCamera(child)
     // NOTE: Track onPointerMissed objects separate from the scene
     context.eventManager?.registerPointerMissedObject(child)
 
+    let insertedWithAdd = false
     if (is.object3D(child) && is.object3D(parentObject)) {
       parentObject.add(child)
+      insertedWithAdd = true
       child.dispatchEvent({ type: 'added' })
     }
     else if (is.fog(child)) {
+      // TODO
+      // Currently `material` and `geometry` are attached by
+      // setting `attach` in `createElement`.
+      // Do the same here to eliminate this branch.
       parentObject.fog = child
     }
     else if (typeof child.attach === 'string') {
@@ -132,13 +121,32 @@ export const nodeOps: (context: TresContext) => RendererOptions<TresObject, Tres
         parentObject[child.attach] = child
       }
     }
+
+    // NOTE: Update __tres parent/objects graph
+    childInstance.__tres.parent = parentObject
+    if (parentObject.__tres?.objects && !insertedWithAdd) {
+      if (!parentObject.__tres.objects.includes(child)) {
+        parentObject.__tres.objects.push(child)
+      }
+    }
   }
 
   function remove(node: TresObject | null) {
     if (!node) { return }
     // remove is only called on the node being removed and not on child nodes.
+
+    // TODO:
+    // Figure out why `parent` is being set on `node` here
+    // and remove/refactor.
     node.parent = node.parent || scene
 
+    // NOTE: Update __tres parent/objects graph
+    const parent = node.__tres?.parent || scene
+    if (node.__tres) { node.__tres.parent = null }
+    if (parent.__tres && 'objects' in parent.__tres) {
+      filterInPlace(parent.__tres.objects, obj => obj !== node)
+    }
+
     if (is.object3D(node)) {
       node.removeFromParent?.()
 

src/types/index.ts

@@ -42,15 +42,22 @@ interface TresBaseObject {
 
 export interface LocalState {
   type: string
-  // objects and parent are used when children are added with `attach` instead of being added to the Object3D scene graph
-  objects?: TresObject3D[]
-  parent?: TresObject3D | null
+  eventCount: number
+  root: TresContext
+  handlers: Partial<EventHandlers>
+  memoizedProps: { [key: string]: any }
+  // NOTE:
+  // LocalState holds information about the parent/child relationship
+  // in the Vue graph. If a child is `insert`ed into a parent using
+  // anything but THREE's `add`, it's put into the parent's `objects`.
+  // objects and parent are used when children are added with `attach`
+  // instead of being added to the Object3D scene graph
+  objects: TresObject[]
+  parent: TresObject | null
+  // NOTE: End graph info
+
   primitive?: boolean
-  eventCount?: number
-  handlers?: Partial<EventHandlers>
-  memoizedProps?: { [key: string]: any }
   disposable?: boolean
-  root?: TresContext
 }
 
 // Custom type for geometry and material properties in Object3D
@@ -62,6 +69,8 @@ export interface TresObject3D extends THREE.Object3D<THREE.Object3DEventMap> {
 export type TresObject =
   TresBaseObject & (TresObject3D | THREE.BufferGeometry | THREE.Material | THREE.Fog) & { __tres?: LocalState }
 
+export type TresInstance = TresObject & { __tres: LocalState }
+
 export interface TresScene extends THREE.Scene {
   __tres: {
     root: TresContext

src/utils/index.test.ts

@@ -0,0 +1,62 @@
+import * as utils from './index'
+
+describe('filterInPlace', () => {
+  it('returns the passed array', () => {
+    const arr = [1, 2, 3]
+    const result = utils.filterInPlace(arr, v => v !== 0)
+    expect(result).toBe(arr)
+  })
+  it('removes a single occurence', () => {
+    const arr = [1, 2, 3]
+    utils.filterInPlace(arr, v => v !== 1)
+    expect(arr).toStrictEqual([2, 3])
+  })
+  it('removes every occurence 0', () => {
+    const arr = [1, 1, 2, 1, 3, 1]
+    utils.filterInPlace(arr, v => v !== 1)
+    expect(arr).toStrictEqual([2, 3])
+  })
+
+  it('removes every occurence 1', () => {
+    const [a, b, c] = [{}, {}, {}]
+    const COUNT = 400
+    const arr = []
+    for (const val of [a, b, c]) {
+      for (let i = 0; i < COUNT; i++) {
+        arr.push(val)
+      }
+    }
+    shuffle(arr)
+
+    let filtered = [...arr]
+    utils.filterInPlace(arr, v => v !== b)
+    filtered = filtered.filter(v => v !== b)
+    expect(arr).toStrictEqual(filtered)
+
+    utils.filterInPlace(arr, v => v !== c)
+    filtered = filtered.filter(v => v !== c)
+    expect(arr).toStrictEqual(filtered)
+
+    utils.filterInPlace(arr, v => v !== a)
+    expect(arr).toStrictEqual([])
+  })
+
+  it('sends an index to the callbackFn', () => {
+    const arr = 'abcdefghi'.split('')
+    utils.filterInPlace(arr, (_, i) => i % 2 === 0)
+    expect(arr).toStrictEqual('acegi'.split(''))
+  })
+})
+
+function shuffle(array: any[]) {
+  let currentIndex = array.length
+  while (currentIndex !== 0) {
+    const randomIndex = Math.floor(Math.random() * currentIndex)
+    currentIndex--;
+    [array[currentIndex], array[randomIndex]] = [
+      array[randomIndex],
+      array[currentIndex],
+    ]
+  }
+  return array
+};

src/utils/index.ts

@@ -313,3 +313,20 @@ export function disposeObject3D(object: TresObject): void {
     }
   }
 }
+
+/**
+ * Like Array.filter, but modifies the array in place.
+ * @param array - Array to modify
+ * @param callbackFn - A function called for each element of the array. It should return a truthy value to keep the element in the array.
+ */
+export function filterInPlace<T>(array: T[], callbackFn: (element: T, index: number) => unknown) {
+  let i = 0
+  for (let ii = 0; ii < array.length; ii++) {
+    if (callbackFn(array[ii], ii)) {
+      array[i] = array[ii]
+      i++
+    }
+  }
+  array.length = i
+  return array
+}

src/utils/nodeOpsUtils.ts

@@ -0,0 +1,32 @@
+import type { TresContext } from '../composables/useTresContextProvider'
+import type { LocalState, TresInstance, TresObject } from '../types'
+
+export function prepareTresInstance<T extends TresObject>(obj: T, state: Partial<LocalState>, context: TresContext): TresInstance {
+  const instance = obj as unknown as TresInstance
+  instance.__tres = {
+    type: 'unknown',
+    eventCount: 0,
+    root: context,
+    handlers: {},
+    memoizedProps: {},
+    objects: [],
+    parent: null,
+    ...state,
+  }
+  return instance
+}
+
+export function invalidateInstance(instance: TresObject) {
+  const ctx = instance?.__tres?.root
+
+  if (!ctx) { return }
+
+  if (ctx.render && ctx.render.canBeInvalidated.value) {
+    ctx.invalidate()
+  }
+}
+
+export function noop(fn: string): any {
+  // eslint-disable-next-line no-unused-expressions
+  fn
+}