Messy react + regl framework for simple instanced and batched rendering

example.js

const defaultColor = [1,1,1];
const ExampleCommand = regl => ({
  vert: "…", 
  frag: "…", 
  uniforms: {
    color: ({color}) => color || defaultColor,
  },
  attributes: {
    position: …,
  },
  // Apart from this part everything here is standard regl
  instancedAttributes: {
    translation: new Float32Array(3),
  },
});

const Example = () => {
  const { Element, Batch } = useCommand(ExampleCommand);
  // 1 <Element> = 1 instance added to the command, or the closest <Batch> of the command that contains the <Element>
  
  return (
    <ReglEngine>
      {/* Batch props get added to the command's context */}
      <Batch color={[1,1,1]}>
        {/* Element props are sent to the command's buffers in `instancedAttributes` */}
        <Element translation={new Float32Array([0,2,0])} />
      </Batch>
    
      <Element translation={new Float32Array([0,0,0])} />
      
      <Batch color={[1,0,0]}>
        <Element translation={new Float32Array([0,5,0])} />
        <Element translation={new Float32Array([0,10,0])} />
      </Batch>
    </ReglEngine>
  );
}

ReglEngine.js

import React, {
  useContext,
  useEffect,
  useRef,
  createContext,
  useState,
  useCallback,
} from "react";
import createCamera from "perspective-camera";

let regl;
if (process.env.NODE_ENV === "development") {
  regl = require("regl");
} else {
  regl = require("regl/dist/regl.unchecked.js");
}

// WARNING: won't work with multiple engine instances
const subscribers = new Set();
const useLoopCallback = callback => {
  useEffect(() => {
    if (callback) {
      subscribers.add(callback);

      return () => {
        subscribers.delete(callback);
      };
    }
  }, [callback]);
};

const EngineContext = createContext();
export const useEngine = () => useContext(EngineContext);
export const useLoop = callback => useContext(EngineContext).useLoop(callback);

const CameraContext = createContext();
export const useCamera = () => useEngine().context.camera;

export const EngineWithCamera = ({ camera, context, onResize, ...props }) => {
  const cameraInstance = useRef(createCamera(camera)).current;

  const handleResize = (width, height) => {
    const vmax = Math.max(width, height);
    cameraInstance.viewport[0] = -(vmax / height);
    cameraInstance.viewport[1] = -(vmax / width);
    cameraInstance.viewport[2] = vmax / height;
    cameraInstance.viewport[3] = vmax / width;
    cameraInstance.update();

    if (onResize) {
      onResize(width, height);
    }
  };

  return (
    <CameraContext.Provider value={cameraInstance}>
      <Engine context={{ ...context, camera: cameraInstance }} onResize={handleResize} {...props} />
    </CameraContext.Provider>
  );
};

const Engine = ({
  children,
  context,
  canvasProps = {},
  pixelRatio = process.browser ? window.devicePixelRatio : 1,
  attributes,
  onResize,
  debug = process.env.NODE_ENV === "development",
  defaultShaders,
  drawOnEveryFrame = true,
  onLoop,
  ...props
}) => {
  // FIXME: can't change init variables after first render
  const [engine, setEngine] = useState();
  const [readyForRendering, setReadyForRendering] = useState(false);
  const engineRef = useRef(null);
  const commands = useRef(new Map()).current;
  const internalCanvasRef = useRef();
  const canvasRef = canvasProps.ref || internalCanvasRef;

  useEffect(() => {
    const canvas = canvasRef.current;
    const engineInstance = regl({
      extensions: ["ANGLE_instanced_arrays", "OES_standard_derivatives"],
      optionalExtensions: debug ? ["EXT_disjoint_timer_query"] : [],
      attributes: {
        antialias: false,
        cull: { enable: true },
        alpha: false,
        premultipliedAlpha: false,
        ...attributes,
      },
      profile: debug,
      pixelRatio,
      canvas,
      ...props,
    });
    setEngine(() => engineInstance);
    engineRef.current = engineInstance;

    return () => {
      console.log("destroying regl instance");
      setEngine();
      engineRef.current = null;
      engineInstance.destroy();
    };
  }, []);

  const draw = useRef();
  useEffect(() => {
    if (engine) {
      const contextCommand = engine({
        context: {
          ...context,
          clear: { color: [0.618, 0.618, 0.618, 1] },
        },
      });

      const batchHolder = [];

      draw.current = () => {
        try {
          contextCommand(context => {
            if (onLoop) {
              onLoop(context);
            }

            engine.clear(context.clear);

            for (const callback of subscribers) {
              callback.call(null, context);
            }

            for (const command of commands.values()) {
              const {
                command: callCommand,
                batches,
                indexes,
                instancedBuffers,
                needsUpdate,
                isInstanced,
                name,
              } = command;

              if (needsUpdate) {
                if (process.env.NODE_ENV === "development") {
                  console.log("updating buffers and indexes for command", name);
                }

                // Refresh indexes cache
                command.indexes.clear();
                let index = 0;
                for (const batch of batches) {
                  batch.offset = index;

                  for (const instance of batch.instances) {
                    command.indexes.set(instance, index);
                    index++;
                  }
                }

                command.totalInstances = index;

                // Refill buffers
                for (const [key, buffer] of instancedBuffers) {
                  const { dimensions, ArrayConstructor } = buffer;
                  const data = new ArrayConstructor(command.totalInstances * dimensions);
                  for (const { instances } of batches) {
                    for (const instance of instances) {
                      data.set(instance[key], indexes.get(instance) * dimensions);
                    }
                  }
                  buffer.buffer(data);
                }

                command.needsUpdate = false;
              }

              if (isInstanced) {
                if (command.totalInstances === 0) {
                  continue;
                }

                batchHolder.length = 0;
                for (const batch of batches) {
                  batchHolder.push(batch);
                }

                callCommand(batchHolder);
              } else {
                callCommand();
              }
            }
          });
        } catch (err) {
          loop.cancel();
          throw err;
        }
      };

      const loop = drawOnEveryFrame && engine.frame(draw.current);
      setReadyForRendering(true);

      return () => {
        if (engineRef.current && loop) {
          loop.cancel();
          console.log("destroying regl loop");
        }
      };
    }
  }, [engine]);

  useEffect(() => {
    const canvas = canvasRef.current;

    const handleResize = () => {
      const { width: canvasWidth, height: canvasHeight } = canvas.getBoundingClientRect();
      const width = canvasWidth;
      const height = canvasHeight;

      canvas.width = width * pixelRatio;
      canvas.height = height * pixelRatio;

      if (!drawOnEveryFrame && draw.current) {
        engine.poll();
        draw.current();
      }

      if (onResize) {
        onResize(width, height);
      }
    };

    window.addEventListener("resize", handleResize);
    handleResize();

    return () => {
      window.removeEventListener("resize", handleResize);
    };
  }, [canvasRef, onResize, pixelRatio, drawOnEveryFrame, engine]);

  useEffect(() => {
    if (engine && debug) {
      const debugInterval = setInterval(() => {
        console.table(
          Object.fromEntries(
            Object.entries(engine.stats).map(([key, value]) => [
              key,
              typeof value === "function" ? value() : value,
            ])
          )
        );
        let commandStats = [];
        for (const [
          ,
          { command, totalInstances, batches, name, drawOrder },
        ] of commands.entries()) {
          commandStats.push({
            "drawOrder": drawOrder,
            "name": name,
            "batches": batches.size,
            "instances": totalInstances,
            "invocations": command.stats.count,
            "CPU %": (command.stats.cpuTime / performance.now()) * 100,
            "CPU/frame %": command.stats.cpuTime / command.stats.count / 16,
            "GPU %": (command.stats.gpuTime / performance.now()) * 100,
            "GPU/frame %": command.stats.gpuTime / command.stats.count / 16,
          });
        }
        console.table(commandStats);
      }, 10000);

      return () => {
        clearInterval(debugInterval);
      };
    }
  }, [engine, debug, commands]);

  return (
    <>
      <canvas ref={canvasRef} {...canvasProps}></canvas>

      {readyForRendering ? (
        <EngineContext.Provider
          value={{
            engine,
            commands,
            useLoop: useLoopCallback,
            defaultShaders,
            context,
            draw: () => {
              engine.poll();
              draw.current();
            },
          }}
        >
          {children}
        </EngineContext.Provider>
      ) : null}
    </>
  );
};
export default Engine;

const BatchContext = createContext();
export const useBatch = () => useContext(BatchContext);

export const useCommand = (draw, name) => {
  const engineObject = useEngine();
  const { engine, commands, defaultShaders } = engineObject;

  // If command hasn't been created, create it
  if (!commands.has(draw)) {
    const { vert, frag, uniforms, attributes, instancedAttributes, drawOrder = 0, ...rest } = draw(
      engine
    );

    const instancedBuffers = new Map();
    const instancedAttributesWithBuffers = {};

    for (const attribute in instancedAttributes) {
      const value = instancedAttributes[attribute];
      const buffer = {
        buffer: engine.buffer({ usage: "dynamic", type: "float32" }),
        dimensions: value.length,
        BYTES_PER_ELEMENT: value.BYTES_PER_ELEMENT,
        ArrayConstructor: value.constructor,
      };

      instancedBuffers.set(attribute, buffer);

      instancedAttributesWithBuffers[attribute] = {
        buffer: ({ instancedBuffers }) => instancedBuffers.get(attribute).buffer,
        divisor: 1,
        offset: ({ instancedBuffers }, { offset }) => {
          const buffer = instancedBuffers.get(attribute);
          return offset * buffer.BYTES_PER_ELEMENT * buffer.dimensions;
        },
      };
    }

    const context = {
      batches: new Set(),
      rootBatch: { instances: new Set() },
      indexes: new Map(),
      instancedBuffers,
      name: name || draw.name || "unnamed command",
      isInstanced: !!instancedAttributes,
      needsUpdate: !!instancedAttributes,
      totalInstances: 0,
      drawOrder,
    };

    const Batch = ({ children, ...props }) => {
      const state = useRef({ instances: new Set() });
      for (const key in props) {
        state.current[key] = props[key];
      }
      return <BatchContext.Provider value={state.current}>{children}</BatchContext.Provider>;
    };

    const Element = ({ children = null, onLoop, ...props }) => {
      // Create a stable identity for this component instance
      const instance = useRef(props).current;

      // Use a batch (unbatched instances get batched together)
      const batch = useBatch() || context.rootBatch;
      const command = context;

      useEffect(() => {
        // Add this instance to its batch within the command
        // and create the batch if it doesn't exist yet
        if (!command.batches.has(batch)) {
          command.batches.add(batch);
        }

        batch.instances.add(instance);
        command.needsUpdate = true;

        return () => {
          // Delete this instance
          batch.instances.delete(instance);
          command.needsUpdate = true;
        };
      }, [command, instance, batch]);

      // A function for updating data in buffers for this instance
      const update = useCallback(
        (key, data) => {
          const { instancedBuffers, indexes } = command;
          const buffer = instancedBuffers.get(key).buffer;

          if (buffer._buffer.byteLength) {
            instancedBuffers
              .get(key)
              .buffer.subdata(data, data.BYTES_PER_ELEMENT * data.length * indexes.get(instance));
          }
        },
        [command, instance]
      );

      // Update buffers from props whenever this component re-renders
      useEffect(() => {
        for (const [key] of instancedBuffers) {
          if (props[key]) {
            instance[key] = props[key];
            update(key, instance[key]);
          }
        }
      });

      const { useLoop } = useEngine();
      useLoop(onLoop && (context => onLoop(update, instance, context)));

      return children;
    };

    const command = engine({
      instances: context.isInstanced ? (c, { instances }) => instances.size : undefined,
      context,
      attributes: {
        ...attributes,
        ...instancedAttributesWithBuffers,
      },
      uniforms,
      vert: vert || defaultShaders.vert,
      frag: frag || defaultShaders.frag,
      ...rest,
    });

    context.command = command;
    context.Batch = Batch;
    context.Element = Element;

    commands.set(draw, context);

    // Sort commands
    const commandArray = Array.from(commands).sort((a, b) => a[1].drawOrder - b[1].drawOrder);
    commands.clear();
    for (const [command, context] of commandArray) {
      commands.set(command, context);
    }
  }

  return commands.get(draw);
};