@needle-tools/engine 4.14.0-next.31f837e → 4.14.0-next.52fdb13

package/src/engine/engine_materialpropertyblock.ts

@@ -1,30 +1,56 @@
-import { Material, Object3D, Color, Texture, Vector2, Vector3, Vector4, WebGLRenderer, Scene, Camera, BufferGeometry, Group } from "three";
+import { BufferGeometry, Camera, Color, Euler,Group, Material, Object3D, Scene, Texture, Vector2, Vector3, Vector4, WebGLRenderer } from "three";
 
-type MaterialPropertyType = number | number[] | Color | Texture | Vector2 | Vector3 | Vector4 | null;
 
+// @TODO: we need to detect objects with materials both transparent and NOT transparent. These need to be updated in scene.onBeforeRender to have correct renderlists
+
+/**
+ * Valid types that can be used as material property overrides
+ */
+type MaterialPropertyType = number | number[] | Color | Texture | Vector2 | Vector3 | Vector4 | null | Euler;
+
+/**
+ * Defines offset and repeat transformations for texture coordinates
+ */
 export interface TextureTransform {
+    /** UV offset applied to the texture */
     offset?: Vector2;
+    /** UV repeat/scale applied to the texture */
     repeat?: Vector2;
 }
 
-export interface PropertyBlockOverride {
+/**
+ * Represents a single material property override with optional texture transformation
+ * @template T The type of the property value
+ */
+export interface PropertyBlockOverride<T extends MaterialPropertyType = MaterialPropertyType> {
+    /** The name of the material property to override (e.g., "color", "map", "roughness") */
     name: string;
-    value: MaterialPropertyType;
+    /** The value to set for this property */
+    value: T;
+    /** Optional texture coordinate transformation (only used when value is a Texture) */
     textureTransform?: TextureTransform;
 }
 
+/**
+ * Utility type that extracts only non-function property names from a type
+ * @template T The type to extract property names from
+ */
 type NonFunctionPropertyNames<T> = {
     [K in keyof T]: T[K] extends Function ? never : K
 }[keyof T];
 
-// Centralized registry for all property block related data
+/**
+ * Centralized registry for all property block related data.
+ * Uses WeakMaps to allow automatic garbage collection when objects are destroyed.
+ * @internal
+ */
 class PropertyBlockRegistry {
     // Map from object to its property block
     private objectToBlock = new WeakMap<Object3D, MaterialPropertyBlock>();
 
     // Track which materials belong to which property block (to prevent applying to wrong materials)
     // Use WeakSet for automatic cleanup when materials are garbage collected
-    private objectToMaterials = new WeakMap<Object3D, WeakSet<Material>>();
+    // private objectToMaterials = new WeakMap<Object3D, WeakSet<Material>>();
 
     // Track which meshes have callbacks for which property block owners
     private meshToOwners = new WeakMap<Object3D, Set<Object3D>>();
@@ -45,21 +71,21 @@ class PropertyBlockRegistry {
 
     deleteBlock(object: Object3D): void {
         this.objectToBlock.delete(object);
-        this.objectToMaterials.delete(object);
+        // this.objectToMaterials.delete(object);
     }
 
-    addMaterial(object: Object3D, material: Material): void {
-        let materials = this.objectToMaterials.get(object);
-        if (!materials) {
-            materials = new WeakSet();
-            this.objectToMaterials.set(object, materials);
-        }
-        materials.add(material);
-    }
+    // addMaterial(object: Object3D, material: Material): void {
+    //     let materials = this.objectToMaterials.get(object);
+    //     if (!materials) {
+    //         materials = new WeakSet();
+    //         this.objectToMaterials.set(object, materials);
+    //     }
+    //     materials.add(material);
+    // }
 
-    hasMaterial(object: Object3D, material: Material): boolean {
-        return this.objectToMaterials.get(object)?.has(material) ?? false;
-    }
+    // hasMaterial(object: Object3D, material: Material): boolean {
+    //     return this.objectToMaterials.get(object)?.has(material) ?? false;
+    // }
 
     isHooked(mesh: Object3D, owner: Object3D): boolean {
         return this.meshToOwners.get(mesh)?.has(owner) ?? false;
@@ -95,17 +121,130 @@ class PropertyBlockRegistry {
 
 const registry = new PropertyBlockRegistry();
 
+/**
+ * MaterialPropertyBlock allows per-object material property overrides without creating new material instances.
+ * This is useful for rendering multiple objects with the same base material but different properties
+ * (e.g., different colors, textures, or shader parameters).
+ *
+ * ## How Property Blocks Work
+ *
+ * **Important**: Overrides are registered on the **Object3D**, not on the material.
+ * This means:
+ * - If you change the object's material, the overrides will still be applied to the new material
+ * - Multiple objects can share the same material but have different property overrides
+ * - If you don't want overrides applied after changing a material, you must remove them using {@link removeOveride}, {@link clearAllOverrides}, or {@link dispose}
+ *
+ * The property block system works by:
+ * - Temporarily applying overrides in onBeforeRender
+ * - Restoring original values in onAfterRender
+ * - Managing shader defines and program cache keys for correct shader compilation
+ * - Supporting texture coordinate transforms per object
+ *
+ * ## Common Use Cases
+ *
+ * - **Lightmaps**: Apply unique lightmap textures to individual objects sharing the same material
+ * - **Reflection Probes**: Apply different environment maps per object for localized reflections
+ * - **See-through effects**: Temporarily override transparency/transmission properties for X-ray effects
+ *
+ * ## Getting a MaterialPropertyBlock
+ *
+ * **Important**: Do not use the constructor directly. Instead, use the static {@link MaterialPropertyBlock.get} method:
+ *
+ * ```typescript
+ * const block = MaterialPropertyBlock.get(myMesh);
+ * ```
+ *
+ * This method will either return an existing property block or create a new one if it doesn't exist.
+ * It automatically:
+ * - Creates the property block instance
+ * - Registers it in the internal registry
+ * - Attaches the necessary render callbacks to the object
+ * - Handles Groups by applying overrides to all child meshes
+ *
+ * @example Basic usage
+ * ```typescript
+ * // Get or create a property block for an object
+ * const block = MaterialPropertyBlock.get(myMesh);
+ *
+ * // Override the color property
+ * block.setOverride("color", new Color(1, 0, 0));
+ *
+ * // Override a texture with custom UV transform (useful for lightmaps)
+ * block.setOverride("lightMap", myLightmapTexture, {
+ *   offset: new Vector2(0.5, 0.5),
+ *   repeat: new Vector2(2, 2)
+ * });
+ *
+ * // Set a shader define
+ * block.setDefine("USE_CUSTOM_FEATURE", 1);
+ * ```
+ *
+ * @example Material swapping behavior
+ * ```typescript
+ * const mesh = new Mesh(geometry, materialA);
+ * const block = MaterialPropertyBlock.get(mesh);
+ * block.setOverride("color", new Color(1, 0, 0));
+ *
+ * // The color override is red for materialA
+ *
+ * // Swap the material - overrides persist and apply to the new material!
+ * mesh.material = materialB;
+ * // The color override is now red for materialB too
+ *
+ * // If you don't want overrides on the new material, remove them:
+ * block.clearAllOverrides(); // Remove all overrides
+ * // or
+ * block.removeOveride("color"); // Remove specific override
+ * // or
+ * block.dispose(); // Remove the entire property block
+ * ```
+ *
+ * @example Lightmap usage
+ * ```typescript
+ * const block = MaterialPropertyBlock.get(mesh);
+ * block.setOverride("lightMap", lightmapTexture);
+ * block.setOverride("lightMapIntensity", 1.5);
+ * ```
+ *
+ * @example See-through effect
+ * ```typescript
+ * const block = MaterialPropertyBlock.get(mesh);
+ * block.setOverride("transparent", true);
+ * block.setOverride("opacity", 0.3);
+ * ```
+ *
+ * @template T The material type this property block is associated with
+ */
 export class MaterialPropertyBlock<T extends Material = Material> {
     private _overrides: PropertyBlockOverride[] = [];
     private _defines: Record<string, string | number | boolean> = {};
     private _object: Object3D | null = null;
 
+    /** The object this property block is attached to */
     get object(): Object3D | null { return this._object; }
 
-    constructor(object: Object3D | null = null) {
+    /**
+     * Creates a new MaterialPropertyBlock
+     * @param object The object this property block is for (optional)
+     */
+    protected constructor(object: Object3D | null = null) {
         this._object = object;
     }
 
+    /**
+     * Gets or creates a MaterialPropertyBlock for the given object.
+     * This is the recommended way to obtain a property block instance.
+     *
+     * @template T The material type
+     * @param object The object to get/create a property block for
+     * @returns The MaterialPropertyBlock associated with this object
+     *
+     * @example
+     * ```typescript
+     * const block = MaterialPropertyBlock.get(myMesh);
+     * block.setOverride("roughness", 0.5);
+     * ```
+     */
     static get<T extends Material = Material>(object: Object3D): MaterialPropertyBlock<T> {
         let block = registry.getBlock(object);
         if (!block) {
@@ -116,6 +255,20 @@ export class MaterialPropertyBlock<T extends Material = Material> {
         return block as MaterialPropertyBlock<T>;
     }
 
+    /**
+     * Checks if an object has any property overrides
+     * @param object The object to check
+     * @returns True if the object has a property block with overrides
+     */
+    static hasOverrides(object: Object3D): boolean {
+        const block = registry.getBlock(object);
+        return block ? block.hasOverrides() : false;
+    }
+
+    /**
+     * Disposes this property block and cleans up associated resources.
+     * After calling dispose, this property block should not be used.
+     */
     dispose() {
         if (this._object) {
             registry.deleteBlock(this._object);
@@ -125,6 +278,29 @@ export class MaterialPropertyBlock<T extends Material = Material> {
         this._object = null;
     }
 
+    /**
+     * Sets or updates a material property override.
+     * The override will be applied to the material during rendering.
+     *
+     * @param name The name of the material property to override (e.g., "color", "map", "roughness")
+     * @param value The value to set
+     * @param textureTransform Optional UV transform (only used when value is a Texture)
+     *
+     * @example
+     * ```typescript
+     * // Override a simple property
+     * block.setOverride("roughness", 0.8);
+     *
+     * // Override a color
+     * block.setOverride("color", new Color(0xff0000));
+     *
+     * // Override a texture with UV transform
+     * block.setOverride("map", texture, {
+     *   offset: new Vector2(0, 0),
+     *   repeat: new Vector2(2, 2)
+     * });
+     * ```
+     */
     setOverride<K extends NonFunctionPropertyNames<T>>(name: K, value: T[K], textureTransform?: TextureTransform): void;
     setOverride(name: string, value: MaterialPropertyType, textureTransform?: TextureTransform): void;
     setOverride(name: string, value: MaterialPropertyType, textureTransform?: TextureTransform): void {
@@ -137,36 +313,144 @@ export class MaterialPropertyBlock<T extends Material = Material> {
         }
     }
 
+    /**
+     * Gets the override for a specific property with type-safe value inference
+     * @param name The property name to get
+     * @returns The PropertyBlockOverride with correctly typed value if it exists, undefined otherwise
+     *
+     * @example
+     * ```typescript
+     * const block = MaterialPropertyBlock.get<MeshStandardMaterial>(mesh);
+     *
+     * // Value is inferred as number | undefined
+     * const roughness = block.getOverride("roughness")?.value;
+     *
+     * // Value is inferred as Color | undefined
+     * const color = block.getOverride("color")?.value;
+     *
+     * // Value is inferred as Texture | null | undefined
+     * const map = block.getOverride("map")?.value;
+     *
+     * // Explicitly specify the type for properties not on the base material type
+     * const transmission = block.getOverride<number>("transmission")?.value;
+     *
+     * // Or use a more specific material type
+     * const physicalBlock = block as MaterialPropertyBlock<MeshPhysicalMaterial>;
+     * const transmissionTyped = physicalBlock.getOverride("transmission")?.value; // number
+     * ```
+     */
+    getOverride<K extends NonFunctionPropertyNames<T>>(name: K): PropertyBlockOverride<T[K] & MaterialPropertyType> | undefined;
+    getOverride<V extends MaterialPropertyType = MaterialPropertyType>(name: string): PropertyBlockOverride<V> | undefined;
     getOverride(name: string): PropertyBlockOverride | undefined {
         return this._overrides.find(o => o.name === name);
     }
 
-    clearOverride(name: string): void {
+    /**
+     * Removes a specific property override.
+     * After removal, the material will use its original property value for this property.
+     *
+     * @param name The property name to remove the override for
+     *
+     * @example
+     * ```typescript
+     * const block = MaterialPropertyBlock.get(mesh);
+     *
+     * // Set some overrides
+     * block.setOverride("color", new Color(1, 0, 0));
+     * block.setOverride("roughness", 0.5);
+     * block.setOverride("lightMap", lightmapTexture);
+     *
+     * // Remove a specific override - the material will now use its original color
+     * block.removeOveride("color");
+     *
+     * // Other overrides (roughness, lightMap) remain active
+     * ```
+     */
+    removeOveride<K extends NonFunctionPropertyNames<T>>(name: K | ({} & string)): void {
         const index = this._overrides.findIndex(o => o.name === name);
         if (index >= 0) {
             this._overrides.splice(index, 1);
         }
     }
 
+    /**
+     * Removes all property overrides from this block.
+     * After calling this, the material will use its original values for all properties.
+     *
+     * **Note**: This does NOT remove shader defines. Use {@link clearDefine} or {@link dispose} for that.
+     *
+     * @example Remove all overrides but keep the property block
+     * ```typescript
+     * const block = MaterialPropertyBlock.get(mesh);
+     *
+     * // Set multiple overrides
+     * block.setOverride("color", new Color(1, 0, 0));
+     * block.setOverride("roughness", 0.5);
+     * block.setOverride("lightMap", lightmapTexture);
+     *
+     * // Later, remove all overrides at once
+     * block.clearAllOverrides();
+     *
+     * // The material now uses its original values
+     * // The property block still exists and can be reused with new overrides
+     * ```
+     *
+     * @example Temporarily disable all overrides
+     * ```typescript
+     * const block = MaterialPropertyBlock.get(mesh);
+     *
+     * // Save current overrides if you want to restore them later
+     * const savedOverrides = [...block.overrides];
+     *
+     * // Clear all overrides temporarily
+     * block.clearAllOverrides();
+     *
+     * // Do some rendering without overrides...
+     *
+     * // Restore overrides
+     * savedOverrides.forEach(override => {
+     *   block.setOverride(override.name, override.value, override.textureTransform);
+     * });
+     * ```
+     *
+     * @see {@link removeOveride} - To remove a single override
+     * @see {@link dispose} - To completely remove the property block and clean up resources
+     */
     clearAllOverrides(): void {
         this._overrides = [];
     }
 
-    removeOverride(name: string): void {
-        this.clearOverride(name);
-    }
-
-    get overrides(): PropertyBlockOverride[] {
+    /**
+     * Gets all property overrides as a readonly array
+     * @returns Array of all property overrides
+     */
+    get overrides(): readonly PropertyBlockOverride[] {
         return this._overrides;
     }
 
+    /**
+     * Checks if this property block has any overrides
+     * @returns True if there are any overrides set
+     */
     hasOverrides(): boolean {
         return this._overrides.length > 0;
     }
 
     /**
-     * Set a shader define that will be included in the program cache key
-     * This allows different objects sharing the same material to have different shader programs
+     * Set a shader define that will be included in the program cache key.
+     * This allows different objects sharing the same material to have different shader programs.
+     *
+     * Defines affect shader compilation and are useful for enabling/disabling features per-object.
+     *
+     * @param name The define name (e.g., "USE_LIGHTMAP", "ENABLE_REFLECTIONS")
+     * @param value The define value (typically a boolean, number, or string)
+     *
+     * @example
+     * ```typescript
+     * // Enable a feature for this specific object
+     * block.setDefine("USE_CUSTOM_SHADER", true);
+     * block.setDefine("QUALITY_LEVEL", 2);
+     * ```
      */
     setDefine(name: string, value: string | number | boolean): void {
         this._defines[name] = value;
@@ -174,6 +458,7 @@ export class MaterialPropertyBlock<T extends Material = Material> {
 
     /**
      * Remove a shader define
+     * @param name The define name to remove
      */
     clearDefine(name: string): void {
         this._defines[name] = undefined as any;
@@ -181,11 +466,20 @@ export class MaterialPropertyBlock<T extends Material = Material> {
 
     /**
      * Get all defines set on this property block
+     * @returns A readonly record of all defines
      */
     getDefines(): Readonly<Record<string, string | number | boolean>> {
         return this._defines;
     }
 
+    /**
+     * Generates a cache key based on the current overrides and defines.
+     * This key is used internally to ensure correct shader program selection
+     * when objects share materials but have different property blocks.
+     *
+     * @returns A string representing the current state of this property block
+     * @internal
+     */
     getCacheKey(): string {
         const parts: string[] = [];
 
@@ -227,13 +521,25 @@ export class MaterialPropertyBlock<T extends Material = Material> {
     }
 }
 
+/**
+ * Symbol used to store original material values on the material object
+ * @internal
+ */
 const $originalValues = Symbol("originalValues");
 
+/**
+ * Stores an original material property value before override
+ * @internal
+ */
 interface OriginalValue {
     name: string;
     value: unknown;
 }
 
+/**
+ * Stores saved texture transform state for restoration
+ * @internal
+ */
 interface SavedTextureTransform {
     name: string;
     offsetX: number;
@@ -242,11 +548,22 @@ interface SavedTextureTransform {
     repeatY: number;
 }
 
+/**
+ * Symbol used to store saved texture transforms on the material object
+ * @internal
+ */
 const $savedTextureTransforms = Symbol("savedTextureTransforms");
 
+/**
+ * Type for Three.js object render callbacks
+ * @internal
+ */
 type ObjectRenderCallback = (this: Object3D, renderer: WebGLRenderer, scene: Scene, camera: Camera, geometry: BufferGeometry, material: Material, group: Group) => void;
 
-// Collect all materials from an object and its children
+/**
+ * Collect all materials from an object and its children
+ * @internal
+ */
 function collectMaterials(object: Object3D, materials: Set<Material>): void {
     const obj = object as Object3D & { material?: Material | Material[] };
     if (obj.material) {
@@ -263,8 +580,14 @@ function collectMaterials(object: Object3D, materials: Set<Material>): void {
     }
 }
 
-// Find property block by checking this object and parent (if parent is a Group)
-// Returns both the block and the owner object
+/**
+ * Find property block by checking this object and parent (if parent is a Group).
+ * Returns both the block and the owner object.
+ *
+ * @param obj The object to search from
+ * @returns The property block and its owner object, or undefined if not found
+ * @internal
+ */
 function findPropertyBlockAndOwner(obj: Object3D): { block: MaterialPropertyBlock; owner: Object3D } | undefined {
     // First check if this object itself has a property block
     let block = registry.getBlock(obj);
@@ -279,8 +602,93 @@ function findPropertyBlockAndOwner(obj: Object3D): { block: MaterialPropertyBloc
     return undefined;
 }
 
+/**
+ * Symbol to track which materials are currently being rendered for an object
+ * @internal
+ */
+const currentlyRenderingFlag = Symbol("beforeRenderingFlag");
+
+/**
+ * Tracks original transparent values that were changed during render list building
+ * @internal
+ */
+const beforeRenderListTransparentChanged = new WeakMap<Object3D, boolean>();
+
+/**
+ * Tracks original transmission values that were changed during render list building
+ * @internal
+ */
+const beforeRenderListTransparentChangedTransmission = new WeakMap<Object3D, number>();
+
+/**
+ * Callback invoked before an object is added to the render list.
+ * Used to temporarily override transparency/transmission for correct render list assignment.
+ * @internal
+ */
+const onBeforeRenderListPush = function (this: Object3D, _object: Object3D, _geometry: BufferGeometry, material: Material, _group: Group) {
+    const block = registry.getBlock(_object);
+    if (!block) {
+        return;
+    }
+    if (block.hasOverrides()) {
+        const transmission = block.getOverride<number>("transmission")?.value;
+        const transparent = block.getOverride("transparent")?.value;
+
+        if (transmission !== undefined && typeof transmission === "number" && "transmission" in material && transmission !== material.transmission) {
+            beforeRenderListTransparentChangedTransmission.set(this, material.transmission as number);
+            material.transmission = transmission;
+        }
+        if (transparent !== undefined && typeof transparent === "boolean" && transparent !== material.transparent) {
+            beforeRenderListTransparentChanged.set(this, material.transparent);
+            material.transparent = transparent;
+        }
+    }
+}
+/**
+ * Callback invoked after an object is added to the render list.
+ * Restores the original transparency/transmission values that were overridden in onBeforeRenderListPush.
+ * @internal
+ */
+const onAfterRenderListPush = function (this: Object3D, _object: Object3D, _geometry: BufferGeometry, material: Material, _group: Group) {
+    const prevTransparent = beforeRenderListTransparentChanged.get(_object);
+    if (prevTransparent !== undefined) {
+        beforeRenderListTransparentChanged.delete(_object);
+        material.transparent = prevTransparent;
+    }
+
+    const prevTransmission = beforeRenderListTransparentChangedTransmission.get(_object);
+    if (prevTransmission !== undefined) {
+        beforeRenderListTransparentChangedTransmission.delete(_object);
+        (material as any).transmission = prevTransmission;
+    }
+}
+
 // #region OnBeforeRender
+/**
+ * Main callback invoked before rendering an object.
+ * Applies property block overrides and defines to the material.
+ * @internal
+ */
 const onBeforeRender_MaterialBlock: ObjectRenderCallback = function (this: Object3D, _renderer: WebGLRenderer, _scene: Scene, _camera: Camera, _geometry: BufferGeometry, material: Material, _group: Group) {
+
+
+    // Only run if the material belongs to this object and is a "regular" material (not depth or other override material)
+    const materials = (this as any).material as Array<Material> | Material | undefined;
+    if (!materials) return;
+    if (Array.isArray(materials)) {
+        if (!materials.includes(material)) return;
+    }
+    else if (materials !== material) {
+        return;
+    }
+
+    // Keep track of which materials rendering started for so we can check in onAfterRender if it was processed 
+    // (in case of override materials like depth material where onBeforeRender runs but we don't want to apply overrides)
+    if (this[currentlyRenderingFlag] === undefined) this[currentlyRenderingFlag] = new WeakSet<Material>();
+    this[currentlyRenderingFlag].add(material);
+
+
+    // Before rendering, check if this object (or its parent Group) has a property block with overrides for this material.
     const result = findPropertyBlockAndOwner(this);
     if (!result) {
         return;
@@ -289,9 +697,9 @@ const onBeforeRender_MaterialBlock: ObjectRenderCallback = function (this: Objec
     const { block: propertyBlock, owner } = result;
 
     // Only apply if this material was registered with this property block
-    if (!registry.hasMaterial(owner, material)) {
-        return;
-    }
+    // if (!registry.hasMaterial(owner, material)) {
+    //     return;
+    // }
 
     const overrides = propertyBlock.overrides;
     const mat = material as any;
@@ -373,7 +781,19 @@ const onBeforeRender_MaterialBlock: ObjectRenderCallback = function (this: Objec
 };
 
 // #region OnAfterRender
+/**
+ * Main callback invoked after rendering an object.
+ * Restores the original material property values and defines.
+ * @internal
+ */
 const onAfterRender_MaterialBlock: ObjectRenderCallback = function (this: Object3D, _renderer: WebGLRenderer, _scene: Scene, _camera: Camera, _geometry: BufferGeometry, material: Material, _group: Group) {
+
+    // We don't want to run this logic if onBeforeRender didn't run for this material (e.g. due to DepthMaterial or other override material), so we check the flag set in onBeforeRender
+    if (this[currentlyRenderingFlag] === undefined) return;
+    if (!this[currentlyRenderingFlag].has(material)) return;
+    this[currentlyRenderingFlag].delete(material);
+
+
     const result = findPropertyBlockAndOwner(this);
     if (!result) {
         return;
@@ -382,9 +802,9 @@ const onAfterRender_MaterialBlock: ObjectRenderCallback = function (this: Object
     const { block: propertyBlock, owner } = result;
 
     // Only restore if this material was registered with this property block
-    if (!registry.hasMaterial(owner, material)) {
-        return;
-    }
+    // if (!registry.hasMaterial(owner, material)) {
+    //     return;
+    // }
 
     const overrides = propertyBlock.overrides;
 
@@ -446,25 +866,38 @@ const onAfterRender_MaterialBlock: ObjectRenderCallback = function (this: Object
 
 
 // #region Attach Callbacks
+/**
+ * Attaches the property block render callbacks to an object and its child meshes.
+ * @param object The object to attach callbacks to
+ * @param _propertyBlock The property block being attached (unused but kept for clarity)
+ * @internal
+ */
 function attachPropertyBlockToObject(object: Object3D, _propertyBlock: MaterialPropertyBlock): void {
     // Collect and register all materials that belong to this property block
-    const materials = new Set<Material>();
-    collectMaterials(object, materials);
-    materials.forEach(mat => registry.addMaterial(object, mat));
+    // const materials = new Set<Material>();
+    // collectMaterials(object, materials);
+    // materials.forEach(mat => registry.addMaterial(object, mat));
 
     // Attach callbacks to renderable objects (Mesh, SkinnedMesh)
     // Groups don't render themselves but we still need to handle child meshes
     if (object.type === "Group") {
         object.children.forEach(child => {
             if (child.type === "Mesh" || child.type === "SkinnedMesh") {
-                attachCallbacksToMesh(child, object);
+                attachCallbacksToMesh(child, object, _propertyBlock);
             }
         });
     } else if (object.type === "Mesh" || object.type === "SkinnedMesh") {
-        attachCallbacksToMesh(object, object);
+        attachCallbacksToMesh(object, object, _propertyBlock);
     }
 }
-function attachCallbacksToMesh(mesh: Object3D, propertyBlockOwner: Object3D): void {
+/**
+ * Attaches render callbacks to a specific mesh object.
+ * Chains with existing callbacks if they exist.
+ * @param mesh The mesh to attach callbacks to
+ * @param propertyBlockOwner The object that owns the property block (may be the mesh itself or its parent Group)
+ * @internal
+ */
+function attachCallbacksToMesh(mesh: Object3D, propertyBlockOwner: Object3D, _propertyBlock: MaterialPropertyBlock): void {
     // Check if this specific mesh already has our callbacks attached for this property block owner
     if (registry.isHooked(mesh, propertyBlockOwner)) {
         // Already hooked for this property block owner
@@ -473,6 +906,11 @@ function attachCallbacksToMesh(mesh: Object3D, propertyBlockOwner: Object3D): vo
 
     registry.addHook(mesh, propertyBlockOwner);
 
+    /**
+     * Expose the property block for e.g. Needle Inspector
+     */
+    mesh["needle:materialPropertyBlock"] = _propertyBlock;
+
     if (!mesh.onBeforeRender) {
         mesh.onBeforeRender = onBeforeRender_MaterialBlock;
     } else {
@@ -492,9 +930,28 @@ function attachCallbacksToMesh(mesh: Object3D, propertyBlockOwner: Object3D): vo
             original.call(this, renderer, scene, camera, geometry, material, group);
         };
     }
+
+    /** @ts-ignore patched in three.js */
+    mesh.onBeforeRenderListPush = onBeforeRenderListPush;
+    /** @ts-ignore patched in three.js */
+    mesh.onAfterRenderListPush = onAfterRenderListPush;
+
 }
 //#endregion
 
+/**
+ * Checks if an object has a MaterialPropertyBlock attached to it.
+ *
+ * @param object The object to check
+ * @returns True if the object has a property block registered
+ *
+ * @example
+ * ```typescript
+ * if (objectHasPropertyBlock(myMesh)) {
+ *   console.log("This mesh has property overrides");
+ * }
+ * ```
+ */
 export function objectHasPropertyBlock(object: Object3D): boolean {
     return registry.getBlock(object) !== undefined;
 }