Writing custom shaderMaterial animating with react-spring

[phaminhieuu]

Hi, I'm writing a custom shader material using 'shaderMaterial' from 'react-three/drei' and it can animate with react-spring/three, this is my code:

AnimateShaderMaterial

import { a } from "@react-spring/three";
import { shaderMaterial } from "@react-three/drei";
import { ShaderMaterialProps, extend, useFrame } from "@react-three/fiber";
import { Ref, forwardRef, useImperativeHandle, useMemo, useRef } from "react";
import { ShaderMaterial } from "three";

interface Props extends ShaderMaterialProps {
  uni?: { [key: string]: any };
}

export const AnimateShaderMaterial = a(
  forwardRef(
    (
      { uni, vertexShader, fragmentShader, ...props }: Props,
      ref: Ref<ShaderMaterial>,
    ) => {
      const shader = useRef<ShaderMaterial>();
      useImperativeHandle(ref, () => shader.current);

      const CustomMaterial = useMemo(
        () =>
          shaderMaterial(
            {
              time: 0,
              uVisibility: 0,
              ...uni,
            },
            vertexShader,
            fragmentShader,
          ),
        [vertexShader, fragmentShader, uni],
      );

      extend({ CustomMaterial });

      useFrame(({ clock }) => {
        shader.current.uniforms.time.value = clock.getElapsedTime();
      });

      return <customMaterial ref={shader} {...props} />;
    },
  ),
);

Parent component

import { SpringValue, useSpring } from "@react-spring/three";

const ParentComponent = () => {
  const [{ visibility }, api] = useSpring(() => ({
    from: {
      visibility: 0,
    },
   to: {
      visibility: 1,
     }
  }));

  return (
    <mesh>
      <boxGeometry args={[2,2,2]} />
      <AnimateShaderMaterial
        uni={{
          uResolution: new Vector2(),
        }}
        vertexShader={flexibleVert}
        fragmentShader={flexibleFrag}
        transparent
        uniforms-uVisibility-value={visibility}
      />
      {children}
    </mesh>
  );
};

It work pretty well when I use it multiple times in my project, but I wondering is it a good practice when I put 'extend' inside component not outside like document, because I need to pass 'uni' from the outside? Do I need to memorize the "CustomMaterial" to avoid rerendering multiple times ? Thanks