wgpu/src/resources/buffer.ts

import { BufferError, WebGPUError } from '../utils/errors.js'
import { TypedArray, TypedArrayConstructor } from '../utils/index.js'
import { ResourceType } from '../utils/internal-enums.js'

type SmallTypedArray = Exclude<TypedArray, BigInt64Array | BigUint64Array>
type SmallTypedArrayConstructor = Exclude<TypedArrayConstructor, BigInt64ArrayConstructor | BigUint64ArrayConstructor>

export class Buffer {
  _device: GPUDevice
  _handle: GPUBuffer

  _mapped = false

  _defaultStagingBuffer: GPUBuffer

  get handle() {
    return this._handle
  }

  get size() {
    return this._handle.size
  }

  get usage() {
    return this._handle.usage
  }

  get resourceType() {
    return ResourceType.Buffer
  }

  constructor(device: GPUDevice, texture: GPUBuffer) {
    this._device = device
    this._handle = texture
  }

  static create(device: GPUDevice, descriptor: GPUBufferDescriptor) {
    try {
      return new Buffer(
        device,
        device.createBuffer(descriptor)
      )
    } catch (err) {
      throw BufferError.from(err)
    }
  }

  _createStagingOptions(size: number = this.size) {
    return {
      size,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
    }
  }

  _getStagingBuffer(size: number) {
    return this._device.createBuffer(this._createStagingOptions(size))
  }

  write(data: ArrayBufferView | ArrayBuffer, offset: number = 0, dataOffset: number = 0) {
    if (!(this.usage & GPUBufferUsage.COPY_DST)) {
      console.warn('Buffer usage does not include COPY_DST. Buffer.write may fail.')
    }

    if (!this._device?.queue) {
      throw WebGPUError.deviceUnavailable()
    }

    this._device.queue.writeBuffer(
      this._handle,
      offset,
      data,
      dataOffset
    )
  }


  async read(out: SmallTypedArray | DataView | undefined, byteOffset: number = 0, byteSize: number = -1) {
    if (!this._device) {
      throw WebGPUError.deviceUnavailable()
    }

    if (!(this.usage & GPUBufferUsage.MAP_READ)) {
      throw BufferError.invalidRead()
    }

    if (byteOffset < 0) { throw new RangeError('Read byteOffset cannot be negative') }
    if (byteSize < 0) { byteSize = this.size - byteOffset }
    if (byteSize < 0) { throw new RangeError(`Invalid calculated byteSize (${byteSize})`) }
    if (byteSize === 0) { return out ?? new ArrayBuffer(0) }
    if (byteOffset + byteSize > this.size) { throw new RangeError(`Read range exceeds buffer size`) }


    if (out != null) {
      if (!ArrayBuffer.isView(out)) { throw new TypeError('"out" parameter must be a TypedArray or DataView.') }

      if (out.byteLength < byteSize) { throw new RangeError(`Provided output buffer too small`) }
    }

    let result: SmallTypedArray | ArrayBuffer | DataView<ArrayBufferLike>
    let range: ArrayBuffer

    try {
      await this.handle.mapAsync(GPUMapMode.READ, byteOffset, byteSize)
      range = this.handle.getMappedRange(byteOffset, byteSize)

      if (out != null) {
        const SourceView = out.constructor as SmallTypedArrayConstructor
        const bytesPerElement = SourceView.BYTES_PER_ELEMENT

        if (!bytesPerElement) {
          if (out instanceof DataView) {
            new Uint8Array(
              out.buffer,
              out.byteOffset,
              byteSize
            ).set(new Uint8Array(range))
          } else {
            throw new TypeError('"out" is not a standard TypedArray or DataView')
          }
        } else {
          if (byteSize % bytesPerElement !== 0) {
            throw new RangeError(`"byteSize" (${byteSize}) incompatible with "out" byte size (${bytesPerElement})`)
          }

          const view = new SourceView(range)
          const target = new SourceView(
            out.buffer,
            out.byteOffset,
            byteSize / bytesPerElement
          )

          target.set(view)
        }

        result = out
      } else {
        result = range.slice(0)
      }
    } catch (err) {
      throw BufferError.from(err)
    } finally {
      this.handle.unmap()
    }

    return result
  }

  toBindingResource({ offset, size }: { offset?: number; size?: number } = {}) {
    return {
      buffer: this._handle,
      offset: offset || 0,
      size: size || this.size
    }
  }

  destroy() {
    this._handle?.destroy()
    this._handle = null
  }
}


export class UniformBuffer extends Buffer {
  constructor(device: GPUDevice, buffer: GPUBuffer) {
    super(device, buffer)
  }

  static create(device: GPUDevice, descriptor: Omit<GPUBufferDescriptor, 'usage'>) {
    const usage = GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
    return super.create(device, {
      usage,
      ...descriptor
    })
  }
}