Update instance removal logic on component unmount

src/index.tsx

@@ -1,53 +1,63 @@
 import diff from "microdiff";
 import p5 from "p5";
-import React, { createRef, FC, memo, useState } from "react";
+import React, { createRef, FC, memo, MutableRefObject, useRef } from "react";
 import { useIsomorphicEffect } from "rooks";
 
 type Wrapper = HTMLDivElement;
-
-export interface SketchProps {
+export type Sketch = (instance: P5CanvasInstance) => void;
+export type SketchProps = {
   [key: string]: any;
-}
-
-export interface Sketch {
-  (instance: P5Instance): void;
-}
-
-export interface P5WrapperProps extends SketchProps {
+};
+export type P5WrapperProps = SketchProps & {
   sketch: Sketch;
-}
-
-export interface P5Instance extends p5 {
+};
+export type P5CanvasInstance = p5 & {
   updateWithProps?: (props: SketchProps) => void;
-}
+};
+
+// @TODO: remove in next major version, keep for compatibility reasons for now.
+export type P5Instance = P5CanvasInstance;
 
-function createCanvas(sketch: Sketch, wrapper: Wrapper): P5Instance {
+function createCanvasInstance(
+  sketch: Sketch,
+  wrapper: Wrapper
+): P5CanvasInstance {
   return new p5(sketch, wrapper);
 }
 
+function removeCanvasInstance(
+  canvasInstanceRef: MutableRefObject<P5CanvasInstance | undefined>
+) {
+  canvasInstanceRef.current?.remove();
+  canvasInstanceRef.current = undefined;
+}
+
 const ReactP5WrapperComponent: FC<P5WrapperProps> = ({
   sketch,
   children,
   ...props
 }) => {
   const wrapperRef = createRef<Wrapper>();
-  const [instance, setInstance] = useState<P5Instance>();
+  const canvasInstanceRef = useRef<P5CanvasInstance>();
 
   useIsomorphicEffect(() => {
     if (wrapperRef.current === null) {
       return;
     }
 
-    instance?.remove();
-    const canvas = createCanvas(sketch, wrapperRef.current);
-    setInstance(canvas);
+    removeCanvasInstance(canvasInstanceRef);
+    canvasInstanceRef.current = createCanvasInstance(
+      sketch,
+      wrapperRef.current
+    );
   }, [sketch]);
 
-  useIsomorphicEffect(() => {
-    instance?.updateWithProps?.(props);
-  }, [props, instance]);
+  useIsomorphicEffect(
+    () => canvasInstanceRef.current?.updateWithProps?.(props),
+    [props]
+  );
 
-  useIsomorphicEffect(() => () => instance?.remove(), []);
+  useIsomorphicEffect(() => () => removeCanvasInstance(canvasInstanceRef), []);
 
   return <div ref={wrapperRef}>{children}</div>;
 };