From 3ba70e49b5856c53611c314d4855d679a546a7df Mon Sep 17 00:00:00 2001
From: Evan You <evan@vuejs.org>
Date: Wed, 17 Jul 2024 18:20:23 +0800
Subject: [PATCH] feat: useTemplateRef()

---
 packages/dts-test/ref.test-d.ts               |  8 +++
 .../__tests__/apiTemplateRef.spec.ts          | 71 +++++++++++++++++++
 packages/runtime-core/src/apiTemplateRef.ts   | 25 +++++++
 packages/runtime-core/src/index.ts            |  1 +
 4 files changed, 105 insertions(+)
 create mode 100644 packages/runtime-core/__tests__/apiTemplateRef.spec.ts
 create mode 100644 packages/runtime-core/src/apiTemplateRef.ts

diff --git a/packages/dts-test/ref.test-d.ts b/packages/dts-test/ref.test-d.ts
index 1456c5232..9de9f3f0a 100644
--- a/packages/dts-test/ref.test-d.ts
+++ b/packages/dts-test/ref.test-d.ts
@@ -17,6 +17,7 @@ import {
   toRefs,
   toValue,
   unref,
+  useTemplateRef,
 } from 'vue'
 import { type IsAny, type IsUnion, describe, expectType } from './utils'
 
@@ -456,3 +457,10 @@ describe('toRef <-> toValue', () => {
 // unref
 declare const text: ShallowRef<string> | ComputedRef<string> | MaybeRef<string>
 expectType<string>(unref(text))
+
+// useTemplateRef
+const tRef = useTemplateRef('foo')
+expectType<Readonly<ShallowRef<unknown>>>(tRef)
+
+const tRef2 = useTemplateRef<HTMLElement>('bar')
+expectType<Readonly<ShallowRef<HTMLElement | null>>>(tRef2)
diff --git a/packages/runtime-core/__tests__/apiTemplateRef.spec.ts b/packages/runtime-core/__tests__/apiTemplateRef.spec.ts
new file mode 100644
index 000000000..fa69adccb
--- /dev/null
+++ b/packages/runtime-core/__tests__/apiTemplateRef.spec.ts
@@ -0,0 +1,71 @@
+import {
+  h,
+  nextTick,
+  nodeOps,
+  ref,
+  render,
+  useTemplateRef,
+} from '@vue/runtime-test'
+
+describe('useTemplateRef', () => {
+  test('should work', () => {
+    let tRef
+    const key = 'refKey'
+    const Comp = {
+      setup() {
+        tRef = useTemplateRef(key)
+      },
+      render() {
+        return h('div', { ref: key })
+      },
+    }
+    const root = nodeOps.createElement('div')
+    render(h(Comp), root)
+    expect(tRef!.value).toBe(root.children[0])
+  })
+
+  test('should be readonly', () => {
+    let tRef
+    const key = 'refKey'
+    const Comp = {
+      setup() {
+        tRef = useTemplateRef(key)
+      },
+      render() {
+        return h('div', { ref: key })
+      },
+    }
+    const root = nodeOps.createElement('div')
+    render(h(Comp), root)
+
+    // @ts-expect-error
+    tRef.value = 123
+
+    expect(tRef!.value).toBe(root.children[0])
+    expect('target is readonly').toHaveBeenWarned()
+  })
+
+  test('should be updated for ref of dynamic strings', async () => {
+    let t1, t2
+    const key = ref('t1')
+    const Comp = {
+      setup() {
+        t1 = useTemplateRef<HTMLAnchorElement>('t1')
+        t2 = useTemplateRef('t2')
+      },
+      render() {
+        return h('div', { ref: key.value })
+      },
+    }
+    const root = nodeOps.createElement('div')
+    render(h(Comp), root)
+
+    expect(t1!.value).toBe(root.children[0])
+    expect(t2!.value).toBe(null)
+
+    key.value = 't2'
+    await nextTick()
+    expect(t2!.value).toBe(root.children[0])
+    expect(t1!.value).toBe(null)
+  })
+})
diff --git a/packages/runtime-core/src/apiTemplateRef.ts b/packages/runtime-core/src/apiTemplateRef.ts
new file mode 100644
index 000000000..df2651471
--- /dev/null
+++ b/packages/runtime-core/src/apiTemplateRef.ts
@@ -0,0 +1,25 @@
+import { type ShallowRef, readonly, shallowRef } from '@vue/reactivity'
+import { getCurrentInstance } from './component'
+import { warn } from './warning'
+import { EMPTY_OBJ } from '@vue/shared'
+
+export function useTemplateRef<T = unknown>(
+  key: string,
+): Readonly<ShallowRef<T | null>> {
+  const i = getCurrentInstance()
+  const r = shallowRef(null)
+  if (i) {
+    const refs = i.refs === EMPTY_OBJ ? (i.refs = {}) : i.refs
+    Object.defineProperty(refs, key, {
+      enumerable: true,
+      get: () => r.value,
+      set: val => (r.value = val),
+    })
+  } else if (__DEV__) {
+    warn(
+      `useTemplateRef() is called when there is no active component ` +
+        `instance to be associated with.`,
+    )
+  }
+  return (__DEV__ ? readonly(r) : r) as any
+}
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index d5f0e6865..dd8c4eefa 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -62,6 +62,7 @@ export { defineComponent } from './apiDefineComponent'
 export { defineAsyncComponent } from './apiAsyncComponent'
 export { useAttrs, useSlots } from './apiSetupHelpers'
 export { useModel } from './helpers/useModel'
+export { useTemplateRef } from './apiTemplateRef'
 
 // <script setup> API ----------------------------------------------------------