Type-safe useCanvas hook for React canvas animations

App.tsx

import React from "react";
import useCanvas from "./useCanvas";
import { useWindowSize } from "@uidotdev/usehooks";

const App = () => {
  const { width, height } = useWindowSize();
  const canvasRef = useCanvas(
    (ctx, frame) => {
      ctx.fillRect(100, 100 + frame, 100, 100);
    },
    {
      beforeDraw: (ctx) => {
        ctx.clearRect(0, 0, width!, height!);
      },
    },
  );

  return <canvas ref={canvasRef} width={width!} height={height!} />;
};

export default App;

useCanvas.ts

import { useRef, useEffect } from "react";

/**
 * Callback function to draw on the canvas
 *
 * @params
 * ctx     CanvasRenderingContext2D
 * frame?  number
 */
type DrawFunction = (ctx: CanvasRenderingContext2D, frame: number) => void;

interface Args {
  beforeDraw: DrawFunction;
  afterDraw: DrawFunction;
}

/**
 * Hook for accessing a stateful canvas
 *
 * @params
 * fn                 DrawFunction
 * args.beforeDraw?   DrawFunction
 * args.afterDraw?    DrawFunction
 */
const useCanvas = (fn: DrawFunction, args: Partial<Args> = {}) => {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  const contextRef = useRef<CanvasRenderingContext2D | null>(null);
  const drawRef = useRef<(() => void) | null>(null);
  const frameRef = useRef<number>(0);
  const animationFrameIdRef = useRef<number>(0);

  // Wrapper for draw function to manage animation frame and frame count state
  drawRef.current = () => {
    frameRef.current += 1;

    if (args.beforeDraw) args.beforeDraw(contextRef.current!, frameRef.current);
    fn(contextRef.current!, frameRef.current);
    if (args.afterDraw) args.afterDraw(contextRef.current!, frameRef.current);

    animationFrameIdRef.current = window.requestAnimationFrame(
      drawRef.current!,
    );
  };

  useEffect(() => {
    if (!canvasRef.current) return;

    if (!contextRef.current) {
      contextRef.current = canvasRef.current.getContext("2d");
    }

    if (args.beforeDraw) args.beforeDraw(contextRef.current!, frameRef.current);
    drawRef.current!();
    if (args.afterDraw) args.afterDraw(contextRef.current!, frameRef.current);

    resizeCanvas(canvasRef.current);

    return () => window.cancelAnimationFrame(animationFrameIdRef.current);
  }, [fn, args]);

  return canvasRef;
};

export default useCanvas;

/**
 * Manage canvas resize when window size changes
 * Includes logic for high-density pixel devices
 */
const resizeCanvas = (canvas: HTMLCanvasElement) => {
  const { width, height } = canvas.getBoundingClientRect();

  if (canvas.width == width && canvas.height == height) return;

  const { devicePixelRatio: ratio = 1 } = window;
  const context = canvas.getContext("2d");

  // Maintain aspect ratio for high-density pixel devices
  if (context) {
    canvas.width = width * ratio;
    canvas.height = height * ratio;
    context.scale(ratio, ratio);
    // Context not available, do what we can
  } else {
    canvas.width = width;
    canvas.height = height;
  }
};