Building an LRU VRAM Cache for Tile Buffers

export interface TileBuffers {
  vertex: GPUBuffer;
  index: GPUBuffer;
  uniform: GPUBuffer;
}

interface Node {
  key: string;
  buffers: TileBuffers;
  bytes: number;      // aligned resident footprint of this tile
  prev: Node | null;
  next: Node | null;
}

export interface LruOptions {
  capacityBytes: number;   // hard ceiling for resident tile buffers
  highWatermark?: number;  // fraction of capacity that triggers eviction (default 0.9)
  lowWatermark?: number;   // fraction eviction drains down to (default 0.7)
}

// Canonical z/x/y key so two code paths that build the same tile share a slot.
export const tileKey = (z: number, x: number, y: number): string => `${z}/${x}/${y}`;

export class TileBufferCache {
  private map = new Map<string, Node>();
  private head: Node | null = null;   // most-recently-used
  private tail: Node | null = null;   // least-recently-used
  private residentBytes = 0;
  private readonly high: number;
  private readonly low: number;
  private retiring: TileBuffers[] = []; // evicted, awaiting safe destroy()

  constructor(private device: GPUDevice, private opts: LruOptions) {
    this.high = (opts.highWatermark ?? 0.9) * opts.capacityBytes;
    this.low = (opts.lowWatermark ?? 0.7) * opts.capacityBytes;
  }

  get sizeBytes(): number { return this.residentBytes; }
  get count(): number { return this.map.size; }

  // Look up a tile, marking it most-recently-used if present.
  get(key: string): TileBuffers | undefined {
    const node = this.map.get(key);
    if (!node) return undefined;
    this.moveToHead(node);   // touch: O(1) relink, no sort
    return node.buffers;
  }

  // Insert or replace a tile, then evict down to the low watermark if needed.
  // `protectedKeys` are the current working set that must never be evicted.
  set(key: string, buffers: TileBuffers, bytes: number, protectedKeys: Set<string>): void {
    const existing = this.map.get(key);
    if (existing) {
      this.residentBytes -= existing.bytes;
      this.retire(existing.buffers);        // replacing: retire the old backing
      existing.buffers = buffers;
      existing.bytes = bytes;
      this.residentBytes += bytes;
      this.moveToHead(existing);
    } else {
      const node: Node = { key, buffers, bytes, prev: null, next: null };
      this.map.set(key, node);
      this.residentBytes += bytes;
      this.linkAtHead(node);
    }
    if (this.residentBytes > this.high) this.evictToLow(protectedKeys);
  }

  // Evict least-recently-used tiles until under the low watermark, skipping
  // the protected working set. Walks from the tail (oldest) toward the head.
  private evictToLow(protectedKeys: Set<string>): void {
    let node = this.tail;
    while (node && this.residentBytes > this.low) {
      const prev = node.prev;               // capture before unlink
      if (!protectedKeys.has(node.key)) {
        this.unlink(node);
        this.map.delete(node.key);
        this.residentBytes -= node.bytes;
        this.retire(node.buffers);          // defer destroy to reclaim()
      }
      node = prev;
    }
  }

  // Hold evicted buffers until the GPU is provably done reading them.
  private retire(b: TileBuffers): void {
    this.retiring.push(b);
  }

  // Call once per frame after queue.submit(): the safe point to free VRAM.
  async reclaim(): Promise<void> {
    if (this.retiring.length === 0) return;
    const batch = this.retiring;
    this.retiring = [];
    await this.device.queue.onSubmittedWorkDone(); // no in-flight pass reads these
    for (const b of batch) {
      b.vertex.destroy();
      b.index.destroy();
      b.uniform.destroy();
    }
  }

  // ---- intrusive doubly linked list helpers (all O(1)) ----

  private linkAtHead(node: Node): void {
    node.prev = null;
    node.next = this.head;
    if (this.head) this.head.prev = node;
    this.head = node;
    if (!this.tail) this.tail = node;
  }

  private unlink(node: Node): void {
    if (node.prev) node.prev.next = node.next; else this.head = node.next;
    if (node.next) node.next.prev = node.prev; else this.tail = node.prev;
    node.prev = node.next = null;
  }

  private moveToHead(node: Node): void {
    if (node === this.head) return;
    this.unlink(node);
    this.linkAtHead(node);
  }

  // Tear down the whole cache, e.g. on device loss or unmount.
  destroy(): void {
    for (const node of this.map.values()) {
      node.buffers.vertex.destroy();
      node.buffers.index.destroy();
      node.buffers.uniform.destroy();
    }
    this.map.clear();
    this.head = this.tail = null;
    this.residentBytes = 0;
    this.retiring = [];
  }
}

function renderFrame(
  device: GPUDevice,
  cache: TileBufferCache,
  visible: { z: number; x: number; y: number }[],
  buildTile: (z: number, x: number, y: number) => { buffers: TileBuffers; bytes: number },
): void {
  const protectedKeys = new Set(visible.map((t) => tileKey(t.z, t.x, t.y)));

  for (const t of visible) {
    const key = tileKey(t.z, t.x, t.y);
    let buffers = cache.get(key);        // touch on hit
    if (!buffers) {
      const built = buildTile(t.z, t.x, t.y);      // uploads via writeBuffer
      cache.set(key, built.buffers, built.bytes, protectedKeys);
      buffers = built.buffers;
    }
    // ... encode a draw that binds buffers.vertex / index / uniform ...
  }

  // ... device.queue.submit([encoder.finish()]) ...
  void cache.reclaim();   // free evicted VRAM after work is submitted
}