useVirtualList.tsx

useVirtualList.tsx

import { RefObject, UIEvent, useCallback, useEffect, useRef, useState } from 'react'

type GetItemSize = (index: number) => number

type UseVirtualList = (
  count: number,
  height: number,
  getItemSize: GetItemSize,
) => {
  outerProps: {
    onScroll: (event: UIEvent<HTMLDivElement>) => void
    ref: RefObject<HTMLDivElement>
  }
  listProps: {
    ref: RefObject<HTMLDivElement>
  }
  contentProps: {
    ref: RefObject<HTMLDivElement>
  }
  visibleItems: number[]
}

function calculateHeights(count: number, getItemSize: GetItemSize) {
  const calculatedHeights: number[] = []
  for (let i = 0; i < count; i++) {
    const base = calculatedHeights.length > 0 ? calculatedHeights[i - 1] : 0
    calculatedHeights.push(base + getItemSize(i))
  }

  return calculatedHeights
}

function getScrollIndexes(scrollTop: number, heights: number[], height: number, count: number): [number, number] {
  let startIndex = 0
  let endIndex = 0
  const maxEnd = scrollTop + height * 2 < heights[heights.length - 1] ? height * 2 : height

  // we always have to overscan the entire list to prevent flicker and empty rows
  // this can happen because the heights get updated before the animation of the toggle has finished

  for (let i = 0; i < heights.length; i++) {
    const startOfItem = i > 0 ? heights[i - 1] : 0
    if (scrollTop - height >= startOfItem) {
      startIndex = i
    }

    if (heights[i] >= scrollTop + maxEnd) {
      endIndex = i
      break
    }
  }

  if (startIndex < 0) startIndex = 0
  if (endIndex > count - 1) endIndex = count - 1

  return [startIndex, endIndex]
}

export function getRange(start: number, stop: number): number[] {
  const range = []
  let length = stop - start

  for (let i = 0; i <= length; i++) {
    range[i] = start
    start++
  }

  return range
}

export const useVirtualList: UseVirtualList = (count, height, getItemSize) => {
  const scrollRef = useRef<HTMLDivElement>(null)
  const listRef = useRef<HTMLDivElement>(null)
  const contentRef = useRef<HTMLDivElement>(null)

  const [heights, setHeights] = useState<number[]>(() => calculateHeights(count, getItemSize))
  const [visibleIndexes, setVisibleIndexes] = useState<[number, number]>(() => {
    return getScrollIndexes(0, heights, height, count)
  })

  const onScroll = useCallback(() => {
    const scrollTop = scrollRef.current ? scrollRef.current.scrollTop : 0
    const [startIndex, endIndex] = getScrollIndexes(scrollTop, heights, height, count)

    if (startIndex !== visibleIndexes[0] || endIndex !== visibleIndexes[1]) {
      setVisibleIndexes([startIndex, endIndex])

      const translate = startIndex > 0 ? heights[startIndex - 1] : 0
      if (contentRef.current) contentRef.current.style.transform = `translate3d(0,${translate}px,0)`
    }
  }, [count, height, heights, visibleIndexes])

  useEffect(() => {
    // this effect reruns if the outer `getItemSize` has changed – which is the case if the
    // array of toggled elements has changed
    const calculatedHeights = calculateHeights(count, getItemSize)
    setHeights(calculatedHeights)

    if (listRef.current) listRef.current.style.height = calculatedHeights[calculatedHeights.length - 1] + 'px'
  }, [count, getItemSize])

  return {
    outerProps: { onScroll, ref: scrollRef },
    listProps: { ref: listRef },
    contentProps: { ref: contentRef },
    visibleItems: getRange(...visibleIndexes),
  }
}