From 9e74ef645d1d821b9678c7d14e6749b499fa2cb7 Mon Sep 17 00:00:00 2001
From: Bjorn Stromberg <bjorn@bjornstar.com>
Date: Thu, 7 Apr 2022 12:14:18 +0200
Subject: [PATCH] Add scaleOverride property (Fixes #367)

---
 packages/react-three-cannon/src/hooks.ts      | 15 ++++++----
 .../react-three-cannon/src/physics-context.ts |  4 +++
 .../src/physics-provider.tsx                  | 29 +++++++++++--------
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/packages/react-three-cannon/src/hooks.ts b/packages/react-three-cannon/src/hooks.ts
index a14653cb..5b0ccd74 100644
--- a/packages/react-three-cannon/src/hooks.ts
+++ b/packages/react-three-cannon/src/hooks.ts
@@ -75,6 +75,7 @@ export type WorkerApi = {
   applyTorque: (torque: Triplet) => void
   quaternion: QuaternionApi
   rotation: VectorApi
+  scaleOverride: (scale: Triplet) => void
   sleep: () => void
   wakeUp: () => void
 }
@@ -130,10 +131,10 @@ function subscribe<T extends SubscriptionName>(
   }
 }
 
-function prepare(object: Object3D, props: BodyProps) {
-  object.userData = props.userData || {}
-  object.position.set(...(props.position || [0, 0, 0]))
-  object.rotation.set(...(props.rotation || [0, 0, 0]))
+function prepare(object: Object3D, { position = [0, 0, 0], rotation = [0, 0, 0], userData = {} }: BodyProps) {
+  object.userData = userData
+  object.position.set(...position)
+  object.rotation.set(...rotation)
   object.updateMatrix()
 }
 
@@ -161,7 +162,7 @@ function useBody<B extends BodyProps<unknown[]>>(
 ): Api {
   const ref = useForwardedRef(fwdRef)
 
-  const { events, refs, subscriptions, worker } = usePhysicsContext()
+  const { events, refs, scaleOverrides, subscriptions, worker } = usePhysicsContext()
   const debugApi = useDebugContext()
 
   useLayoutEffect(() => {
@@ -332,6 +333,10 @@ function useBody<B extends BodyProps<unknown[]>>(
         position: makeVec('position', index),
         quaternion: makeQuaternion(index),
         rotation: makeRotation(index),
+        scaleOverride(scale) {
+          const uuid = getUUID(ref, index)
+          if (uuid) scaleOverrides[uuid] = new Vector3(...scale)
+        },
         sleep() {
           const uuid = getUUID(ref, index)
           uuid && worker.sleep({ uuid })
diff --git a/packages/react-three-cannon/src/physics-context.ts b/packages/react-three-cannon/src/physics-context.ts
index a74a1f6b..eaa65b9c 100644
--- a/packages/react-three-cannon/src/physics-context.ts
+++ b/packages/react-three-cannon/src/physics-context.ts
@@ -8,6 +8,7 @@ import type {
   Subscriptions,
 } from '@pmndrs/cannon-worker-api'
 import { createContext, useContext } from 'react'
+import type { Vector3 } from 'three'
 
 type CannonEvent = CollideBeginEvent | CollideEndEvent | CollideEvent | RayhitEvent
 type CallbackByType<T extends { type: string }> = {
@@ -16,10 +17,13 @@ type CallbackByType<T extends { type: string }> = {
 
 export type CannonEvents = { [uuid: string]: Partial<CallbackByType<CannonEvent>> }
 
+export type ScaleOverrides = { [uuid: string]: Vector3 }
+
 export type PhysicsContext = {
   bodies: { [uuid: string]: number }
   events: CannonEvents
   refs: Refs
+  scaleOverrides: ScaleOverrides
   subscriptions: Subscriptions
   worker: CannonWorkerAPI
 }
diff --git a/packages/react-three-cannon/src/physics-provider.tsx b/packages/react-three-cannon/src/physics-provider.tsx
index 8cc7ac77..5fdb3a86 100644
--- a/packages/react-three-cannon/src/physics-provider.tsx
+++ b/packages/react-three-cannon/src/physics-provider.tsx
@@ -31,13 +31,15 @@ const s = new Vector3(1, 1, 1)
 const q = new Quaternion()
 const m = new Matrix4()
 
-function apply(index: number, positions: Float32Array, quaternions: Float32Array, object?: Object3D) {
+function apply(
+  index: number,
+  positions: Float32Array,
+  quaternions: Float32Array,
+  scale = s,
+  object?: Object3D,
+) {
   if (index !== undefined) {
-    m.compose(
-      v.fromArray(positions, index * 3),
-      q.fromArray(quaternions, index * 4),
-      object ? object.scale : s,
-    )
+    m.compose(v.fromArray(positions, index * 3), q.fromArray(quaternions, index * 4), scale)
     if (object) {
       object.matrixAutoUpdate = false
       object.matrix.copy(m)
@@ -67,10 +69,11 @@ export const PhysicsProvider: FC<PhysicsProviderProps> = ({
 }) => {
   const { invalidate } = useThree()
 
-  const [{ bodies, events, refs, subscriptions, worker }] = useState<PhysicsContext>(() => ({
+  const [{ bodies, events, refs, scaleOverrides, subscriptions, worker }] = useState<PhysicsContext>(() => ({
     bodies: {},
     events: {},
     refs: {},
+    scaleOverrides: {},
     subscriptions: {},
     worker: new CannonWorkerAPI({
       allowSleep,
@@ -178,14 +181,16 @@ export const PhysicsProvider: FC<PhysicsProviderProps> = ({
       for (const ref of Object.values(refs)) {
         if (ref instanceof InstancedMesh) {
           for (let i = 0; i < ref.count; i++) {
-            const index = bodies[`${ref.uuid}/${i}`]
+            const uuid = `${ref.uuid}/${i}`
+            const index = bodies[uuid]
             if (index !== undefined) {
-              ref.setMatrixAt(i, apply(index, positions, quaternions))
+              ref.setMatrixAt(i, apply(index, positions, quaternions, scaleOverrides[uuid]))
+              ref.instanceMatrix.needsUpdate = true
             }
-            ref.instanceMatrix.needsUpdate = true
           }
         } else {
-          apply(bodies[ref.uuid], positions, quaternions, ref)
+          const scale = scaleOverrides[ref.uuid] || ref.scale
+          apply(bodies[ref.uuid], positions, quaternions, scale, ref)
         }
       }
       if (shouldInvalidate) {
@@ -241,7 +246,7 @@ export const PhysicsProvider: FC<PhysicsProviderProps> = ({
   }, [tolerance])
 
   const value = useMemo<PhysicsContext>(
-    () => ({ bodies, events, refs, subscriptions, worker }),
+    () => ({ bodies, events, refs, scaleOverrides, subscriptions, worker }),
     [bodies, events, refs, subscriptions, worker],
   )