From 22a9668b5cafe68bff5cd6bea570bce00aaa4b9a Mon Sep 17 00:00:00 2001
From: Steve Ruiz <steveruizok@gmail.com>
Date: Tue, 28 Sep 2021 11:12:36 +0100
Subject: [PATCH] [feature] Add `isStateful` flag for shape definitions (#130)

* Add `isStateful` flag for shape definitions

* Update useShapeTree.tsx
---
 packages/core/src/hooks/useShapeTree.tsx | 37 ++++++++++++++++++------
 packages/core/src/shapes/createShape.tsx |  2 ++
 packages/core/src/types.ts               |  2 ++
 packages/dev/src/core/label.tsx          |  2 ++
 packages/dev/src/core/state.ts           | 10 +++++++
 5 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/hooks/useShapeTree.tsx b/packages/core/src/hooks/useShapeTree.tsx
index 4b995a077..ec581a589 100644
--- a/packages/core/src/hooks/useShapeTree.tsx
+++ b/packages/core/src/hooks/useShapeTree.tsx
@@ -10,6 +10,7 @@ import type {
   TLCallbacks,
   TLBinding,
   TLBounds,
+  TLShapeUtil,
 } from '+types'
 import { Utils } from '+utils'
 import { Vec } from '@tldraw/vec'
@@ -69,6 +70,13 @@ function shapeIsInViewport(bounds: TLBounds, viewport: TLBounds) {
   return Utils.boundsContain(viewport, bounds) || Utils.boundsCollide(viewport, bounds)
 }
 
+function getShapeUtils<T extends TLShape, E extends Element>(
+  shape: T,
+  shapeUtils: TLShapeUtils<T, E>
+) {
+  return shapeUtils[shape.type] as TLShapeUtil<T, E>
+}
+
 export function useShapeTree<
   T extends TLShape,
   E extends Element,
@@ -111,18 +119,29 @@ export function useShapeTree<
   shapesIdsToRender.clear()
 
   Object.values(page.shapes)
+    .filter(
+      (shape) =>
+        // Always render shapes that are flagged as stateful
+        getShapeUtils(shape, shapeUtils).isStateful ||
+        // Always render selected shapes (this preserves certain drag interactions)
+        selectedIds.includes(shape.id) ||
+        // Otherwise, only render shapes that are in view
+        shapeIsInViewport(shapeUtils[shape.type as T['type']].getBounds(shape), viewport)
+    )
     .sort((a, b) => a.childIndex - b.childIndex)
     .forEach((shape) => {
-      // Don't hide selected shapes (this breaks certain drag interactions)
-      if (
-        selectedIds.includes(shape.id) ||
-        shapeIsInViewport(shapeUtils[shape.type as T['type']].getBounds(shape), viewport)
-      ) {
-        if (shape.parentId === page.id) {
-          shapesIdsToRender.add(shape.id)
-          shapesToRender.add(shape)
-        }
+      // If the shape's parent is the page, add it to our sets of shapes to render
+      if (shape.parentId === page.id) {
+        shapesIdsToRender.add(shape.id)
+        shapesToRender.add(shape)
+        return
       }
+
+      // If the shape's parent is a different shape (e.g. a group),
+      // add the parent to the sets of shapes to render. The parent's
+      // children will all be rendered.
+      shapesIdsToRender.add(shape.parentId)
+      shapesToRender.add(page.shapes[shape.parentId])
     })
 
   // Call onChange callback when number of rendering shapes changes
diff --git a/packages/core/src/shapes/createShape.tsx b/packages/core/src/shapes/createShape.tsx
index a516db99e..7565cf320 100644
--- a/packages/core/src/shapes/createShape.tsx
+++ b/packages/core/src/shapes/createShape.tsx
@@ -22,6 +22,8 @@ export const ShapeUtil = function <T extends TLShape, E extends Element, M = any
 
     canBind: false,
 
+    isStateful: false,
+
     isAspectRatioLocked: false,
 
     create: (props) => {
diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
index 7d1053035..25074c3ad 100644
--- a/packages/core/src/types.ts
+++ b/packages/core/src/types.ts
@@ -324,6 +324,8 @@ export type TLShapeUtil<
 
   canBind: boolean
 
+  isStateful: boolean
+
   getRotatedBounds(this: TLShapeUtil<T, E, M>, shape: T): TLBounds
 
   hitTest(this: TLShapeUtil<T, E, M>, shape: T, point: number[]): boolean
diff --git a/packages/dev/src/core/label.tsx b/packages/dev/src/core/label.tsx
index 818bf130b..c2f368cb4 100644
--- a/packages/dev/src/core/label.tsx
+++ b/packages/dev/src/core/label.tsx
@@ -17,6 +17,8 @@ export interface LabelShape extends TLShape {
 export const Label = new ShapeUtil<LabelShape, HTMLDivElement, { isDarkMode: boolean }>(() => ({
   type: 'label',
 
+  isStateful: true,
+
   defaultProps: {
     id: 'example1',
     type: 'label',
diff --git a/packages/dev/src/core/state.ts b/packages/dev/src/core/state.ts
index 025c5e340..fab961b1a 100644
--- a/packages/dev/src/core/state.ts
+++ b/packages/dev/src/core/state.ts
@@ -122,6 +122,16 @@ export const appState = new AppState({
         name: 'Label',
         childIndex: 2,
         type: 'label',
+        point: [-200, -200],
+        rotation: 0,
+        text: 'My shape is stateful, I should still render while off screen!',
+      },
+      label2: {
+        id: 'label2',
+        parentId: 'page1',
+        name: 'Label',
+        childIndex: 2,
+        type: 'label',
         point: [200, 200],
         rotation: 0,
         text: 'Hello world!',