From 1147af1526f3969c894128cfc83694a690582a5b Mon Sep 17 00:00:00 2001
From: Mikhail Golbakh <golbah@yandex.ru>
Date: Mon, 19 Aug 2024 16:56:51 +0200
Subject: [PATCH] feat(react-query): add memoization for data in infinite
 source (#17)

---
 src/react-query/impl/infinite/hooks.ts | 18 ++++++++++++++++--
 src/react-query/impl/infinite/utils.ts | 19 -------------------
 src/react-query/impl/plain/hooks.ts    |  9 +++++++--
 src/react-query/impl/plain/utils.ts    | 13 -------------
 4 files changed, 23 insertions(+), 36 deletions(-)

diff --git a/src/react-query/impl/infinite/hooks.ts b/src/react-query/impl/infinite/hooks.ts
index eb1ab1d..9b06fb0 100644
--- a/src/react-query/impl/infinite/hooks.ts
+++ b/src/react-query/impl/infinite/hooks.ts
@@ -1,3 +1,5 @@
+import {useMemo} from 'react';
+
 import {useInfiniteQuery} from '@tanstack/react-query';
 
 import type {
@@ -6,9 +8,10 @@ import type {
     DataSourceParams,
     DataSourceState,
 } from '../../../core';
+import {normalizeStatus} from '../../utils/normalizeStatus';
 
 import type {AnyInfiniteQueryDataSource} from './types';
-import {composeOptions, transformResult} from './utils';
+import {composeOptions} from './utils';
 
 export const useInfiniteQueryData = <TDataSource extends AnyInfiniteQueryDataSource>(
     context: DataSourceContext<TDataSource>,
@@ -19,5 +22,16 @@ export const useInfiniteQueryData = <TDataSource extends AnyInfiniteQueryDataSou
     const composedOptions = composeOptions(context, dataSource, params, options);
     const result = useInfiniteQuery(composedOptions);
 
-    return transformResult(result);
+    const transformedData = useMemo<DataSourceState<TDataSource>['data']>(
+        () => result.data?.pages.flat(1) ?? [],
+        [result.data],
+    );
+
+    return {
+        ...result,
+        status: normalizeStatus(result.status, result.fetchStatus),
+        data: transformedData,
+        originalStatus: result.status,
+        originalData: result.data,
+    } as DataSourceState<TDataSource>;
 };
diff --git a/src/react-query/impl/infinite/utils.ts b/src/react-query/impl/infinite/utils.ts
index 76ca8cb..3f3a1a4 100644
--- a/src/react-query/impl/infinite/utils.ts
+++ b/src/react-query/impl/infinite/utils.ts
@@ -2,7 +2,6 @@ import {skipToken} from '@tanstack/react-query';
 import type {
     InfiniteData,
     InfiniteQueryObserverOptions,
-    InfiniteQueryObserverResult,
     QueryFunctionContext,
 } from '@tanstack/react-query';
 
@@ -15,13 +14,10 @@ import type {
     DataSourceOptions,
     DataSourceParams,
     DataSourceResponse,
-    DataSourceState,
 } from '../../../core';
-import {normalizeStatus} from '../../utils/normalizeStatus';
 
 import type {AnyInfiniteQueryDataSource, AnyPageParam} from './types';
 
-const EMPTY_ARRAY: unknown[] = [];
 const EMPTY_OBJECT = {};
 
 export const composeOptions = <TDataSource extends AnyInfiniteQueryDataSource>(
@@ -61,18 +57,3 @@ export const composeOptions = <TDataSource extends AnyInfiniteQueryDataSource>(
         ...options,
     };
 };
-
-export const transformResult = <TDataSource extends AnyInfiniteQueryDataSource>(
-    result: InfiniteQueryObserverResult<
-        InfiniteData<DataSourceData<TDataSource>, AnyPageParam>,
-        DataSourceError<TDataSource>
-    >,
-): DataSourceState<TDataSource> => {
-    return {
-        ...result,
-        status: normalizeStatus(result.status, result.fetchStatus),
-        data: result.data?.pages.flat(1) ?? EMPTY_ARRAY,
-        originalStatus: result.status,
-        originalData: result.data,
-    } as DataSourceState<TDataSource>;
-};
diff --git a/src/react-query/impl/plain/hooks.ts b/src/react-query/impl/plain/hooks.ts
index 80d0928..20a8ea8 100644
--- a/src/react-query/impl/plain/hooks.ts
+++ b/src/react-query/impl/plain/hooks.ts
@@ -6,9 +6,10 @@ import type {
     DataSourceParams,
     DataSourceState,
 } from '../../../core';
+import {normalizeStatus} from '../../utils/normalizeStatus';
 
 import type {AnyPlainQueryDataSource} from './types';
-import {composeOptions, transformResult} from './utils';
+import {composeOptions} from './utils';
 
 export const usePlainQueryData = <TDataSource extends AnyPlainQueryDataSource>(
     context: DataSourceContext<TDataSource>,
@@ -19,5 +20,9 @@ export const usePlainQueryData = <TDataSource extends AnyPlainQueryDataSource>(
     const composedOptions = composeOptions(context, dataSource, params, options);
     const result = useQuery(composedOptions);
 
-    return transformResult(result);
+    return {
+        ...result,
+        status: normalizeStatus(result.status, result.fetchStatus),
+        originalStatus: result.status,
+    } as DataSourceState<TDataSource>;
 };
diff --git a/src/react-query/impl/plain/utils.ts b/src/react-query/impl/plain/utils.ts
index b5b0103..04ff7c6 100644
--- a/src/react-query/impl/plain/utils.ts
+++ b/src/react-query/impl/plain/utils.ts
@@ -1,7 +1,6 @@
 import {
     type QueryFunctionContext,
     type QueryObserverOptions,
-    type QueryObserverResult,
     skipToken,
 } from '@tanstack/react-query';
 
@@ -14,9 +13,7 @@ import type {
     DataSourceOptions,
     DataSourceParams,
     DataSourceResponse,
-    DataSourceState,
 } from '../../../core';
-import {normalizeStatus} from '../../utils/normalizeStatus';
 
 import type {AnyPlainQueryDataSource} from './types';
 
@@ -52,13 +49,3 @@ export const composeOptions = <TDataSource extends AnyPlainQueryDataSource>(
         ...options,
     };
 };
-
-export const transformResult = <TDataSource extends AnyPlainQueryDataSource>(
-    result: QueryObserverResult<DataSourceData<TDataSource>, DataSourceError<TDataSource>>,
-): DataSourceState<TDataSource> => {
-    return {
-        ...result,
-        status: normalizeStatus(result.status, result.fetchStatus),
-        originalStatus: result.status,
-    } as DataSourceState<TDataSource>;
-};