fix(useLoader): make better use of generics for extensions (#2598)

packages/fiber/src/core/hooks.tsx

@@ -5,7 +5,6 @@ import { GLTF } from 'three/examples/jsm/loaders/GLTFLoader'
 import { suspend, preload, clear } from 'suspend-react'
 import { context, RootState, RenderCallback } from './store'
 import { buildGraph, ObjectMap, is, useMutableCallback, useIsomorphicLayoutEffect } from './utils'
-import { LoadingManager } from 'three'
 import { LocalState, Instance } from './renderer'
 
 export interface Loader<T> extends THREE.Loader {
@@ -17,8 +16,13 @@ export interface Loader<T> extends THREE.Loader {
   ): unknown
 }
 
-export type Extensions = (loader: THREE.Loader) => void
+export type LoaderProto<T> = new (...args: any) => Loader<T extends unknown ? any : T>
+export type LoaderReturnType<T, L extends LoaderProto<T>> = T extends unknown
+  ? Awaited<ReturnType<InstanceType<L>['loadAsync']>>
+  : T
+// TODO: this isn't used anywhere, remove in v9
 export type LoaderResult<T> = T extends any[] ? Loader<T[number]> : Loader<T>
+export type Extensions<T extends { prototype: LoaderProto<any> }> = (loader: T['prototype']) => void
 export type ConditionalType<Child, Parent, Truthy, Falsy> = Child extends Parent ? Truthy : Falsy
 export type BranchingReturn<T, Parent, Coerced> = ConditionalType<T, Parent, Coerced, T>
 
@@ -74,8 +78,11 @@ export function useGraph(object: THREE.Object3D) {
   return React.useMemo(() => buildGraph(object), [object])
 }
 
-function loadingFn<T>(extensions?: Extensions, onProgress?: (event: ProgressEvent<EventTarget>) => void) {
-  return function (Proto: new () => LoaderResult<T>, ...input: string[]) {
+function loadingFn<L extends LoaderProto<any>>(
+  extensions?: Extensions<L>,
+  onProgress?: (event: ProgressEvent<EventTarget>) => void,
+) {
+  return function (Proto: L, ...input: string[]) {
     // Construct new loader and run extensions
     const loader = new Proto()
     if (extensions) extensions(loader)
@@ -105,37 +112,37 @@ function loadingFn<T>(extensions?: Extensions, onProgress?: (event: ProgressEven
  * Note: this hook's caller must be wrapped with `React.Suspense`
  * @see https://docs.pmnd.rs/react-three-fiber/api/hooks#useloader
  */
-export function useLoader<T, U extends string | string[]>(
-  Proto: new (manager?: LoadingManager) => LoaderResult<T>,
+export function useLoader<T, U extends string | string[], L extends LoaderProto<T>, R = LoaderReturnType<T, L>>(
+  Proto: L,
   input: U,
-  extensions?: Extensions,
+  extensions?: Extensions<L>,
   onProgress?: (event: ProgressEvent<EventTarget>) => void,
-): U extends any[] ? BranchingReturn<T, GLTF, GLTF & ObjectMap>[] : BranchingReturn<T, GLTF, GLTF & ObjectMap> {
+): U extends any[] ? BranchingReturn<R, GLTF, GLTF & ObjectMap>[] : BranchingReturn<R, GLTF, GLTF & ObjectMap> {
   // Use suspense to load async assets
   const keys = (Array.isArray(input) ? input : [input]) as string[]
-  const results = suspend(loadingFn<T>(extensions, onProgress), [Proto, ...keys], { equal: is.equ })
+  const results = suspend(loadingFn<L>(extensions, onProgress), [Proto, ...keys], { equal: is.equ })
   // Return the object/s
   return (Array.isArray(input) ? results : results[0]) as U extends any[]
-    ? BranchingReturn<T, GLTF, GLTF & ObjectMap>[]
-    : BranchingReturn<T, GLTF, GLTF & ObjectMap>
+    ? BranchingReturn<R, GLTF, GLTF & ObjectMap>[]
+    : BranchingReturn<R, GLTF, GLTF & ObjectMap>
 }
 
 /**
  * Preloads an asset into cache as a side-effect.
  */
-useLoader.preload = function <T, U extends string | string[]>(
-  Proto: new () => LoaderResult<T>,
+useLoader.preload = function <T, U extends string | string[], L extends LoaderProto<T>>(
+  Proto: L,
   input: U,
-  extensions?: Extensions,
+  extensions?: Extensions<L>,
 ) {
   const keys = (Array.isArray(input) ? input : [input]) as string[]
-  return preload(loadingFn<T>(extensions), [Proto, ...keys])
+  return preload(loadingFn<L>(extensions), [Proto, ...keys])
 }
 
 /**
  * Removes a loaded asset from cache.
  */
-useLoader.clear = function <T, U extends string | string[]>(Proto: new () => LoaderResult<T>, input: U) {
+useLoader.clear = function <T, U extends string | string[], L extends LoaderProto<T>>(Proto: L, input: U) {
   const keys = (Array.isArray(input) ? input : [input]) as string[]
   return clear([Proto, ...keys])
 }

packages/fiber/tests/core/hooks.test.tsx

@@ -211,6 +211,21 @@ describe('hooks', () => {
     expect(scene.children[0]).toBe(MockMesh)
   })
 
+  it('can handle useLoader with a loader extension', async () => {
+    class Loader extends THREE.Loader {
+      load = (_url: string) => null
+    }
+
+    let proto!: Loader
+
+    function Test() {
+      return useLoader(Loader, '', (loader) => (proto = loader))
+    }
+    await act(async () => createRoot(canvas).render(<Test />))
+
+    expect(proto).toBeInstanceOf(Loader)
+  })
+
   it('can handle useGraph hook', async () => {
     const group = new THREE.Group()
     const mat1 = new THREE.MeshBasicMaterial()