From 48a2ca38064d90a69934a35ef2644c4b36fb182e Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva13@gmail.com>
Date: Tue, 16 Apr 2024 17:55:06 +0200
Subject: [PATCH] refactor: Revert "refactor: Revert "fix(types): fix
 storeToRefs state return type (#2574) (#2604)""

This reverts commit f550c6080514a851f77a652becbf5b90818479d1.
---
 packages/pinia/src/storeToRefs.ts             | 34 ++++++++-
 .../pinia/test-dts/customizations.test-d.ts   | 73 +++++++++++++++++++
 2 files changed, 103 insertions(+), 4 deletions(-)

diff --git a/packages/pinia/src/storeToRefs.ts b/packages/pinia/src/storeToRefs.ts
index 76e91c333b..a6b91217e2 100644
--- a/packages/pinia/src/storeToRefs.ts
+++ b/packages/pinia/src/storeToRefs.ts
@@ -11,7 +11,15 @@ import {
   toRefs,
 } from 'vue-demi'
 import { StoreGetters, StoreState } from './store'
-import type { PiniaCustomStateProperties, StoreGeneric } from './types'
+import type {
+  _ActionsTree,
+  _GettersTree,
+  _UnwrapAll,
+  PiniaCustomStateProperties,
+  StateTree,
+  Store,
+  StoreGeneric,
+} from './types'
 
 type ToComputedRefs<T> = {
   [K in keyof T]: ToRef<T[K]> extends Ref<infer U>
@@ -19,13 +27,31 @@ type ToComputedRefs<T> = {
     : ToRef<T[K]>
 }
 
+/**
+ * Extracts the refs of a state object from a store. If the state value is a Ref or type that extends ref, it will be kept as is.
+ * Otherwise, it will be converted into a Ref.
+ * @internal
+ */
+type ToStateRefs<SS> =
+  SS extends Store<
+    string,
+    infer UnwrappedState,
+    _GettersTree<StateTree>,
+    _ActionsTree
+  >
+    ? UnwrappedState extends _UnwrapAll<Pick<infer State, infer Key>>
+      ? {
+          [K in Key]: ToRef<State[K]>
+        }
+      : ToRefs<UnwrappedState>
+    : ToRefs<StoreState<SS>>
+
 /**
  * Extracts the return type for `storeToRefs`.
  * Will convert any `getters` into `ComputedRef`.
  */
-export type StoreToRefs<SS extends StoreGeneric> = ToRefs<
-  StoreState<SS> & PiniaCustomStateProperties<StoreState<SS>>
-> &
+export type StoreToRefs<SS extends StoreGeneric> = ToStateRefs<SS> &
+  ToRefs<PiniaCustomStateProperties<StoreState<SS>>> &
   ToComputedRefs<StoreGetters<SS>>
 
 /**
diff --git a/packages/pinia/test-dts/customizations.test-d.ts b/packages/pinia/test-dts/customizations.test-d.ts
index b937c989a2..54bd5b3033 100644
--- a/packages/pinia/test-dts/customizations.test-d.ts
+++ b/packages/pinia/test-dts/customizations.test-d.ts
@@ -193,3 +193,76 @@ expectType<{
     })()
   )
 )
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', () => {
+      const n = ref(1)
+      const customN = ref(1) as Ref<number> & { plusOne: () => void }
+      const double = computed(() => n.value * 2)
+      return {
+        n,
+        customN,
+        double,
+      }
+    })()
+  )
+)
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', () => {
+      const n = ref(1)
+      const customN = ref(1) as Ref<number> & { plusOne: () => void }
+      const double = computed(() => n.value * 2)
+
+      function plusOne() {
+        customN.value++
+      }
+
+      return {
+        n,
+        customN,
+        double,
+        plusOne,
+      }
+    })()
+  )
+)
+
+expectType<{
+  n: Ref<number>
+  customN: Ref<number> & { plusOne: () => void }
+  double: ComputedRef<number>
+  myState: Ref<number>
+  stateOnly: Ref<number>
+}>(
+  storeToRefs(
+    defineStore('a', {
+      state: () => ({
+        n: 1,
+        customN: ref(1) as Ref<number> & { plusOne: () => void },
+      }),
+      getters: {
+        double: (state) => state.n * 2,
+      },
+      actions: {
+        plusOne() {
+          this.n++
+        },
+      },
+    })()
+  )
+)