Implemented DataTree and associated types for exposure in the viewer API

packages/viewer-sandbox/src/main.ts

@@ -1,4 +1,4 @@
-import { Viewer, DefaultViewerParams } from '@speckle/viewer'
+import { Viewer, DefaultViewerParams, DataTree, SpeckleObject } from '@speckle/viewer'
 
 import './style.css'
 import Sandbox from './Sandbox'
@@ -46,7 +46,7 @@ sandbox.makeSceneUI()
 sandbox.makeFilteringUI()
 // Load demo object
 
-sandbox.loadUrl(
+await sandbox.loadUrl(
   // 'https://speckle.xyz/streams/da9e320dad/commits/5388ef24b8?c=%5B-7.66134,10.82932,6.41935,-0.07739,-13.88552,1.8697,0,1%5D'
   // Revit sample house (good for bim-like stuff with many display meshes)
   'https://speckle.xyz/streams/da9e320dad/commits/5388ef24b8'
@@ -59,3 +59,11 @@ sandbox.loadUrl(
   // IFC story, a subtree of the above
   // 'https://latest.speckle.dev/streams/92b620fb17/objects/8247bbc53865b0e0cb5ee4e252e66216'
 )
+
+const dataTree: DataTree = viewer.getDataTree()
+console.log(`Built data tree `, dataTree)
+console.log(
+  dataTree.findAll((obj: SpeckleObject) => {
+    return obj.speckle_type === 'Objects.Geometry.Mesh'
+  })
+)

packages/viewer/src/IViewer.ts

@@ -1,5 +1,6 @@
 import { Vector3 } from 'three'
 import sampleHdri from './assets/sample-hdri.png'
+import { DataTree } from './modules/tree/DataTree'
 
 export interface ViewerParams {
   postprocessing: boolean
@@ -96,5 +97,7 @@ export interface IViewer {
   applyFilter(filter: unknown): Promise<void>
   getObjectsProperties(includeAll?: boolean): unknown
 
+  getDataTree(): DataTree
+
   dispose(): void
 }

packages/viewer/src/index.ts

@@ -7,6 +7,7 @@ import { WorldTree } from './modules/tree/WorldTree'
 import { SpeckleType } from './modules/converter/GeometryConverter'
 import { GeometryConverter } from './modules/converter/GeometryConverter'
 import { SunLightConfiguration } from './IViewer'
+import { DataTree, ObjectPredicate, SpeckleObject } from './modules/tree/DataTree'
 
 export {
   Viewer,
@@ -18,5 +19,11 @@ export {
   SpeckleType,
   GeometryConverter
 }
-
-export type { IViewer, SelectionEvent, SunLightConfiguration }
+export type {
+  IViewer,
+  SelectionEvent,
+  SunLightConfiguration,
+  DataTree,
+  ObjectPredicate,
+  SpeckleObject
+}

packages/viewer/src/modules/Viewer.ts

@@ -22,6 +22,7 @@ import { TreeNode, WorldTree } from './tree/WorldTree'
 import SpeckleRenderer from './SpeckleRenderer'
 import { FilteringManager, FilterMaterialType } from './FilteringManager'
 import { SpeckleType } from './converter/GeometryConverter'
+import { DataTree } from './tree/DataTree'
 
 export class Viewer extends EventEmitter implements IViewer {
   public speckleRenderer: SpeckleRenderer
@@ -120,6 +121,10 @@ export class Viewer extends EventEmitter implements IViewer {
     ;(window as any).V = this
   }
 
+  getDataTree(): DataTree {
+    return WorldTree.getDataTree()
+  }
+
   public async init(): Promise<void> {
     if (this.startupParams.environmentSrc) {
       Assets.getEnvironment(this.startupParams.environmentSrc)

packages/viewer/src/modules/tree/DataTree.ts

@@ -0,0 +1,71 @@
+import TreeModel from 'tree-model'
+import { TreeNode, WorldTree } from './WorldTree'
+
+export type SpeckleObject = Record<string, unknown>
+export type ObjectPredicate = (obj: SpeckleObject) => boolean
+
+export interface DataTree {
+  findFirst(predicate: ObjectPredicate): SpeckleObject
+  findAll(predicate: ObjectPredicate): SpeckleObject[]
+  walk(predicate: ObjectPredicate): void
+}
+
+class DataTreeInternal implements DataTree {
+  tree: TreeModel
+  root: TreeNode
+
+  public constructor() {
+    this.tree = new TreeModel()
+    this.root = this.tree.parse({ id: 'MOTHERSHIP' })
+  }
+  public findAll(predicate: ObjectPredicate): SpeckleObject[] {
+    return this.root
+      .all((node: TreeNode) => {
+        if (!node.model.data) return false
+        return predicate(node.model.data)
+      })
+      .map((value: TreeNode) => value.model.data)
+  }
+
+  public findFirst(predicate: ObjectPredicate): SpeckleObject {
+    return this.root.first((node: TreeNode) => {
+      if (!node.model.data) return false
+      return predicate(node.model.data)
+    }).model.data
+  }
+
+  public walk(predicate: ObjectPredicate) {
+    this.root.walk((node: TreeNode) => {
+      if (!node.model.data) return true
+      return predicate(node.model.data)
+    })
+  }
+}
+
+export class DataTreeBuilder {
+  public static build(root: TreeNode): DataTree {
+    const dataTree = new DataTreeInternal()
+    let parent = null
+    WorldTree.getInstance().walk((node: TreeNode) => {
+      if (!node.parent) {
+        parent = dataTree.root
+        return true
+      }
+
+      parent = dataTree.root.first((localNode) => {
+        return localNode.model.id === node.parent.model.id
+      })
+
+      const _node: TreeNode = WorldTree.getInstance().parse({
+        id: node.model.id,
+        data: node.model.raw,
+        atomic: node.model.atomic,
+        children: []
+      })
+      parent.addChild(_node)
+
+      return true
+    }, root)
+    return dataTree as DataTree
+  }
+}

packages/viewer/src/modules/tree/WorldTree.ts

@@ -1,4 +1,5 @@
 import TreeModel from 'tree-model'
+import { DataTree, DataTreeBuilder } from './DataTree'
 import { NodeRenderView } from './NodeRenderView'
 import { RenderTree } from './RenderTree'
 
@@ -57,6 +58,10 @@ export class WorldTree {
     return WorldTree.renderTreeInstances[id]
   }
 
+  public static getDataTree(): DataTree {
+    return DataTreeBuilder.build(WorldTree.instance._root)
+  }
+
   private tree: TreeModel
   private _root: TreeNode