From 55acb3e9237f38b4904c77e3e89939df07587262 Mon Sep 17 00:00:00 2001
From: Damian Pieczynski <piecyk@gmail.com>
Date: Tue, 13 Jul 2021 08:13:26 +0200
Subject: [PATCH] feat: add rangeExtractor

---
 README.md               |  6 +++++-
 src/index.js            | 46 +++++++++++++++++++++++++++--------------
 src/tests/index.test.js |  9 ++++++++
 types/index.d.ts        | 12 ++++++++++-
 4 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 0dd2da579..2c5618323 100644
--- a/README.md
+++ b/README.md
@@ -289,7 +289,11 @@ const {
   - Optional
   - This function receives the index of each item and should return the item's unique ID.
   - This function should be passed whenever dynamic measurement rendering is enabled and the size or order of items in the list changes.
-
+- `rangeExtractor: Function(range: {start: number; end: number; overscan: number; size: number}) => number[]`
+  - Optional
+  - Defaults to `defaultRangeExtractor`, is exported
+  - **Must be memoized using `React.useCallback()`**
+  - This function receives visible range parameters and should return array of indexes to render
 ### Returns
 
 - `virtualItems: Array<item>`
diff --git a/src/index.js b/src/index.js
index a2fd8d635..9c02e0680 100644
--- a/src/index.js
+++ b/src/index.js
@@ -13,6 +13,19 @@ const defaultMeasureSize = (el, horizontal) => {
   return el[key]
 }
 
+export const defaultRangeExtractor = range => {
+  const start = Math.max(range.start - range.overscan, 0)
+  const end = Math.min(range.end + range.overscan, range.size - 1)
+
+  const arr = []
+
+  for (let i = start; i <= end; i++) {
+    arr.push(i)
+  }
+
+  return arr
+}
+
 export function useVirtual({
   size = 0,
   estimateSize = defaultEstimateSize,
@@ -27,6 +40,7 @@ export function useVirtual({
   scrollOffsetFn,
   keyExtractor = defaultKeyExtractor,
   measureSize = defaultMeasureSize,
+  rangeExtractor = defaultRangeExtractor,
 }) {
   const sizeKey = horizontal ? 'width' : 'height'
   const scrollKey = horizontal ? 'scrollLeft' : 'scrollTop'
@@ -74,7 +88,6 @@ export function useVirtual({
 
   const totalSize = (measurements[size - 1]?.end || 0) + paddingEnd
 
-  latestRef.current.overscan = overscan
   latestRef.current.measurements = measurements
   latestRef.current.outerSize = outerSize
   latestRef.current.totalSize = totalSize
@@ -120,10 +133,17 @@ export function useVirtual({
   measureSizeRef.current = measureSize
 
   const virtualItems = React.useMemo(() => {
+    const indexes = rangeExtractor({
+      start: range.start,
+      end: range.end,
+      overscan,
+      size: measurements.length,
+    })
+
     const virtualItems = []
-    const end = Math.min(range.end, measurements.length - 1)
 
-    for (let i = range.start; i <= end; i++) {
+    for (let k = 0, len = indexes.length; k < len; k++) {
+      const i = indexes[k]
       const measurement = measurements[i]
 
       const item = {
@@ -153,12 +173,14 @@ export function useVirtual({
 
     return virtualItems
   }, [
-    range.start,
-    range.end,
-    measurements,
-    horizontal,
     defaultScrollToFn,
+    horizontal,
     keyExtractor,
+    measurements,
+    overscan,
+    range.end,
+    range.start,
+    rangeExtractor,
   ])
 
   const mountedRef = React.useRef()
@@ -272,10 +294,7 @@ const findNearestBinarySearch = (low, high, getCurrentValue, value) => {
   }
 }
 
-function calculateRange(
-  { overscan, measurements, outerSize, scrollOffset },
-  prevRange
-) {
+function calculateRange({ measurements, outerSize, scrollOffset }, prevRange) {
   const size = measurements.length - 1
   const getOffset = index => measurements[index].start
 
@@ -286,10 +305,7 @@ function calculateRange(
     end++
   }
 
-  start = Math.max(start - overscan, 0)
-  end = Math.min(end + overscan, size)
-
-  if (!prevRange || prevRange.start !== start || prevRange.end !== end) {
+  if (prevRange.start !== start || prevRange.end !== end) {
     return { start, end }
   }
 
diff --git a/src/tests/index.test.js b/src/tests/index.test.js
index 2c3958bca..45abd8847 100644
--- a/src/tests/index.test.js
+++ b/src/tests/index.test.js
@@ -12,12 +12,14 @@ function List({
   onRef,
   parentRef,
   useVirtual,
+  rangeExtractor,
 }) {
   const rowVirtualizer = useVirtual({
     size,
     parentRef,
     overscan,
     useObserver: React.useCallback(() => ({ height, width }), [height, width]),
+    rangeExtractor,
   })
 
   return (
@@ -123,4 +125,11 @@ describe('useVirtual list', () => {
     expect(screen.queryByText('Row 15')).toBeInTheDocument()
     expect(screen.queryByText('Row 16')).not.toBeInTheDocument()
   })
+  it('should use rangeExtractor', () => {
+    render(<List {...props} rangeExtractor={() => [0, 1]} />)
+
+    expect(screen.queryByText('Row 0')).toBeInTheDocument()
+    expect(screen.queryByText('Row 1')).toBeInTheDocument()
+    expect(screen.queryByText('Row 2')).not.toBeInTheDocument()
+  })
 })
diff --git a/types/index.d.ts b/types/index.d.ts
index ac44d4234..5cfcc3e1b 100644
--- a/types/index.d.ts
+++ b/types/index.d.ts
@@ -15,6 +15,15 @@ export type VirtualItem = {
   measureRef: (el: HTMLElement | null) => void
 }
 
+export interface Range {
+  start: number
+  end: number
+  overscan: number
+  size: number
+}
+
+declare function defaultRangeExtractor(range: Range): number[]
+
 declare function useVirtual<T>(options: {
   size: number
   parentRef: React.RefObject<T>
@@ -37,6 +46,7 @@ declare function useVirtual<T>(options: {
   keyExtractor?: (index: number) => number | string
   onScrollElement?: React.RefObject<HTMLElement>
   scrollOffsetFn?: (event?: Event) => number
+  rangeExtractor?: (range: Range) => number[]
 }): {
   virtualItems: VirtualItem[]
   totalSize: number
@@ -45,4 +55,4 @@ declare function useVirtual<T>(options: {
   measure: () => void
 }
 
-export { useVirtual }
+export { defaultRangeExtractor, useVirtual }