Infer preloadedstate properly from provided slice

Author: EskiMojo14

packages/toolkit/src/combineSlices.ts

@@ -1,4 +1,9 @@
-import type { Reducer, StateFromReducersMapObject, UnknownAction } from 'redux'
+import type {
+  PreloadedStateShapeFromReducersMapObject,
+  Reducer,
+  StateFromReducersMapObject,
+  UnknownAction,
+} from 'redux'
 import { combineReducers } from 'redux'
 import { nanoid } from './nanoid'
 import type {
@@ -10,9 +15,9 @@ import type {
 } from './tsHelpers'
 import { getOrInsertComputed } from './utils'
 
-type SliceLike<ReducerPath extends string, State> = {
+type SliceLike<ReducerPath extends string, State, PreloadedState = State> = {
   reducerPath: ReducerPath
-  reducer: Reducer<State>
+  reducer: Reducer<State, any, PreloadedState>
 }
 
 type AnySliceLike = SliceLike<string, any>
@@ -21,18 +26,26 @@ type SliceLikeReducerPath<A extends AnySliceLike> =
   A extends SliceLike<infer ReducerPath, any> ? ReducerPath : never
 
 type SliceLikeState<A extends AnySliceLike> =
-  A extends SliceLike<any, infer State> ? State : never
+  A extends SliceLike<any, infer State, any> ? State : never
+
+type SliceLikePreloadedState<A extends AnySliceLike> =
+  A extends SliceLike<any, any, infer PreloadedState> ? PreloadedState : never
 
 export type WithSlice<A extends AnySliceLike> = {
   [Path in SliceLikeReducerPath<A>]: SliceLikeState<A>
 }
 
+export type WithSlicePreloadedState<A extends AnySliceLike> = {
+  [Path in SliceLikeReducerPath<A>]: SliceLikePreloadedState<A>
+}
+
 type ReducerMap = Record<string, Reducer>
 
-type ExistingSliceLike<DeclaredState> = {
+type ExistingSliceLike<DeclaredState, PreloadedState> = {
   [ReducerPath in keyof DeclaredState]: SliceLike<
     ReducerPath & string,
-    NonUndefined<DeclaredState[ReducerPath]>
+    NonUndefined<DeclaredState[ReducerPath]>,
+    NonUndefined<PreloadedState[ReducerPath & keyof PreloadedState]>
   >
 }[keyof DeclaredState]
 
@@ -48,8 +61,11 @@ export type InjectConfig = {
  */
 export interface CombinedSliceReducer<
   InitialState,
-  DeclaredState = InitialState,
-> extends Reducer<DeclaredState, UnknownAction, Partial<DeclaredState>> {
+  DeclaredState extends InitialState = InitialState,
+  PreloadedState extends Partial<
+    Record<keyof PreloadedState, any>
+  > = Partial<DeclaredState>,
+> extends Reducer<DeclaredState, UnknownAction, PreloadedState> {
   /**
    * Provide a type for slices that will be injected lazily.
    *
@@ -79,9 +95,10 @@ export interface CombinedSliceReducer<
    * const withCustom = rootReducer.inject({ reducerPath: "customName", reducer: customSlice.reducer })
    * ```
    */
-  withLazyLoadedSlices<Lazy = {}>(): CombinedSliceReducer<
+  withLazyLoadedSlices<Lazy = {}, LazyPreloaded = Lazy>(): CombinedSliceReducer<
     InitialState,
-    Id<DeclaredState & Partial<Lazy>>
+    Id<DeclaredState & Partial<Lazy>>,
+    Id<PreloadedState & Partial<LazyPreloaded>>
   >
 
   /**
@@ -96,10 +113,14 @@ export interface CombinedSliceReducer<
    * ```
    *
    */
-  inject<Sl extends Id<ExistingSliceLike<DeclaredState>>>(
+  inject<Sl extends Id<ExistingSliceLike<DeclaredState, PreloadedState>>>(
     slice: Sl,
     config?: InjectConfig,
-  ): CombinedSliceReducer<InitialState, Id<DeclaredState & WithSlice<Sl>>>
+  ): CombinedSliceReducer<
+    InitialState,
+    Id<DeclaredState & WithSlice<Sl>>,
+    Id<PreloadedState & Partial<WithSlicePreloadedState<Sl>>>
+  >
 
   /**
    * Inject a slice.
@@ -113,15 +134,21 @@ export interface CombinedSliceReducer<
    * ```
    *
    */
-  inject<ReducerPath extends string, State>(
+  inject<ReducerPath extends string, State, PreloadedState = State>(
     slice: SliceLike<
       ReducerPath,
-      State & (ReducerPath extends keyof DeclaredState ? never : State)
+      State & (ReducerPath extends keyof DeclaredState ? never : State),
+      PreloadedState &
+        (ReducerPath extends keyof PreloadedState ? never : PreloadedState)
     >,
     config?: InjectConfig,
   ): CombinedSliceReducer<
     InitialState,
-    Id<DeclaredState & WithSlice<SliceLike<ReducerPath, State>>>
+    Id<DeclaredState & WithSlice<SliceLike<ReducerPath, State>>>,
+    Id<
+      PreloadedState &
+        WithSlicePreloadedState<SliceLike<ReducerPath, State, PreloadedState>>
+    >
   >
 
   /**
@@ -301,14 +328,23 @@ type InitialState<Slices extends Array<AnySliceLike | ReducerMap>> =
       : never
   >
 
+type InitialPreloadedState<Slices extends Array<AnySliceLike | ReducerMap>> =
+  UnionToIntersection<
+    Slices[number] extends infer Slice
+      ? Slice extends AnySliceLike
+        ? WithSlicePreloadedState<Slice>
+        : PreloadedStateShapeFromReducersMapObject<Slice>
+      : never
+  >
+
 const isSliceLike = (
   maybeSliceLike: AnySliceLike | ReducerMap,
 ): maybeSliceLike is AnySliceLike =>
   'reducerPath' in maybeSliceLike &&
   typeof maybeSliceLike.reducerPath === 'string'
 
 const getReducers = (slices: Array<AnySliceLike | ReducerMap>) =>
-  slices.flatMap((sliceOrMap) =>
+  slices.flatMap<[string, Reducer]>((sliceOrMap) =>
     isSliceLike(sliceOrMap)
       ? [[sliceOrMap.reducerPath, sliceOrMap.reducer] as const]
       : Object.entries(sliceOrMap),
@@ -370,7 +406,11 @@ const noopReducer: Reducer<Record<string, any>> = (state = emptyObject) => state
 
 export function combineSlices<Slices extends Array<AnySliceLike | ReducerMap>>(
   ...slices: Slices
-): CombinedSliceReducer<Id<InitialState<Slices>>> {
+): CombinedSliceReducer<
+  Id<InitialState<Slices>>,
+  Id<InitialState<Slices>>,
+  Partial<Id<InitialPreloadedState<Slices>>>
+> {
   const reducerMap = Object.fromEntries<Reducer>(getReducers(slices))
 
   const getReducer = () =>

packages/toolkit/src/index.ts

@@ -208,7 +208,11 @@ export type { AutoBatchOptions } from './autoBatchEnhancer'
 
 export { combineSlices } from './combineSlices'
 
-export type { CombinedSliceReducer, WithSlice } from './combineSlices'
+export type {
+  CombinedSliceReducer,
+  WithSlice,
+  WithSlicePreloadedState,
+} from './combineSlices'
 
 export type {
   ExtractDispatchExtensions as TSHelpersExtractDispatchExtensions,

packages/toolkit/src/tests/combineSlices.test-d.ts

@@ -1,4 +1,10 @@
-import type { Reducer, Slice, WithSlice } from '@reduxjs/toolkit'
+import type {
+  Action,
+  Reducer,
+  Slice,
+  WithSlice,
+  WithSlicePreloadedState,
+} from '@reduxjs/toolkit'
 import { combineSlices } from '@reduxjs/toolkit'
 import { createApi, fetchBaseQuery } from '@reduxjs/toolkit/query'
 
@@ -8,6 +14,13 @@ declare const numberSlice: Slice<number, {}, 'number'>
 
 declare const booleanReducer: Reducer<boolean>
 
+declare const mixedReducer: Reducer<string, Action, number>
+
+declare const mixedSliceLike: {
+  reducerPath: 'mixedSlice'
+  reducer: typeof mixedReducer
+}
+
 const exampleApi = createApi({
   baseQuery: fetchBaseQuery(),
   endpoints: (build) => ({
@@ -21,16 +34,31 @@ type ExampleApiState = ReturnType<typeof exampleApi.reducer>
 
 describe('type tests', () => {
   test('combineSlices correctly combines static state', () => {
-    const rootReducer = combineSlices(stringSlice, numberSlice, exampleApi, {
-      boolean: booleanReducer,
-    })
+    const rootReducer = combineSlices(
+      stringSlice,
+      numberSlice,
+      exampleApi,
+      {
+        boolean: booleanReducer,
+        mixed: mixedReducer,
+      },
+      mixedSliceLike,
+    )
 
     expectTypeOf(rootReducer(undefined, { type: '' })).toEqualTypeOf<{
       string: string
       number: number
       boolean: boolean
       api: ExampleApiState
+      mixed: string
+      mixedSlice: string
     }>()
+
+    // test for correct preloaded state handling
+    expectTypeOf(rootReducer).toBeCallableWith(
+      { mixed: 9, mixedSlice: 9 },
+      { type: '' },
+    )
   })
 
   test('combineSlices allows passing no initial reducers', () => {
@@ -63,7 +91,21 @@ describe('type tests', () => {
   test('inject marks injected keys as required', () => {
     const rootReducer = combineSlices(stringSlice).withLazyLoadedSlices<
       WithSlice<typeof numberSlice> &
-        WithSlice<typeof exampleApi> & { boolean: boolean }
+        WithSlice<typeof exampleApi> & { boolean: boolean } & WithSlice<
+          typeof mixedSliceLike
+        > &
+        WithSlice<{
+          reducerPath: 'mixedReducer'
+          reducer: typeof mixedReducer
+        }>,
+      WithSlicePreloadedState<typeof numberSlice> &
+        WithSlicePreloadedState<typeof exampleApi> & {
+          boolean: boolean
+        } & WithSlicePreloadedState<typeof mixedSliceLike> &
+        WithSlicePreloadedState<{
+          reducerPath: 'mixedReducer'
+          reducer: typeof mixedReducer
+        }>
     >()
 
     expectTypeOf(rootReducer(undefined, { type: '' }).number).toEqualTypeOf<
@@ -78,6 +120,14 @@ describe('type tests', () => {
       ExampleApiState | undefined
     >()
 
+    expectTypeOf(rootReducer(undefined, { type: '' }).mixedSlice).toEqualTypeOf<
+      string | undefined
+    >()
+
+    expectTypeOf(
+      rootReducer(undefined, { type: '' }).mixedReducer,
+    ).toEqualTypeOf<string | undefined>()
+
     const withNumber = rootReducer.inject(numberSlice)
 
     expectTypeOf(withNumber(undefined, { type: '' }).number).toBeNumber()
@@ -94,6 +144,21 @@ describe('type tests', () => {
     expectTypeOf(
       withApi(undefined, { type: '' }).api,
     ).toEqualTypeOf<ExampleApiState>()
+
+    const withMixedSlice = rootReducer.inject(mixedSliceLike)
+
+    expectTypeOf(
+      withMixedSlice(undefined, { type: '' }).mixedSlice,
+    ).toBeString()
+
+    const withMixedReducer = rootReducer.inject({
+      reducerPath: 'mixedReducer',
+      reducer: mixedReducer,
+    })
+
+    expectTypeOf(
+      withMixedReducer(undefined, { type: '' }).mixedReducer,
+    ).toBeString()
   })
 
   test('selector() allows defining selectors with injected reducers defined', () => {