From 5a1a89bd6178cc2f84ba91da7d72aee4c6ec1282 Mon Sep 17 00:00:00 2001
From: Evan You <evan@vuejs.org>
Date: Sat, 3 Aug 2024 14:19:19 +0800
Subject: [PATCH] feat(custom-element): useShadowRoot() helper

close #6113
close #8195
---
 packages/runtime-core/src/component.ts        |  2 +-
 .../__tests__/customElement.spec.ts           | 20 ++++++++++++
 packages/runtime-dom/src/apiCustomElement.ts  | 31 +++++++++++++++++--
 packages/runtime-dom/src/index.ts             |  1 +
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/packages/runtime-core/src/component.ts b/packages/runtime-core/src/component.ts
index 9431ce6c4dd..667fa5e713e 100644
--- a/packages/runtime-core/src/component.ts
+++ b/packages/runtime-core/src/component.ts
@@ -417,7 +417,7 @@ export interface ComponentInternalInstance {
    * is custom element?
    * @internal
    */
-  isCE?: boolean
+  isCE?: Element
   /**
    * custom element specific HMR method
    * @internal
diff --git a/packages/runtime-dom/__tests__/customElement.spec.ts b/packages/runtime-dom/__tests__/customElement.spec.ts
index 1ed7cc25311..1762f0cec64 100644
--- a/packages/runtime-dom/__tests__/customElement.spec.ts
+++ b/packages/runtime-dom/__tests__/customElement.spec.ts
@@ -11,6 +11,7 @@ import {
   ref,
   render,
   renderSlot,
+  useShadowRoot,
 } from '../src'
 
 describe('defineCustomElement', () => {
@@ -861,4 +862,23 @@ describe('defineCustomElement', () => {
       )
     })
   })
+
+  describe('useCustomElementRoot', () => {
+    test('should work for style injection', () => {
+      const Foo = defineCustomElement({
+        setup() {
+          const root = useShadowRoot()!
+          const style = document.createElement('style')
+          style.innerHTML = `div { color: red; }`
+          root.appendChild(style)
+          return () => h('div', 'hello')
+        },
+      })
+      customElements.define('my-el', Foo)
+      container.innerHTML = `<my-el></my-el>`
+      const el = container.childNodes[0] as VueElement
+      const style = el.shadowRoot?.querySelector('style')!
+      expect(style.textContent).toBe(`div { color: red; }`)
+    })
+  })
 })
diff --git a/packages/runtime-dom/src/apiCustomElement.ts b/packages/runtime-dom/src/apiCustomElement.ts
index 47672370791..7c4d8793ba1 100644
--- a/packages/runtime-dom/src/apiCustomElement.ts
+++ b/packages/runtime-dom/src/apiCustomElement.ts
@@ -24,6 +24,7 @@ import {
   type VNodeProps,
   createVNode,
   defineComponent,
+  getCurrentInstance,
   nextTick,
   warn,
 } from '@vue/runtime-core'
@@ -191,7 +192,10 @@ export class VueElement extends BaseClass {
   private _numberProps: Record<string, true> | null = null
   private _styles?: HTMLStyleElement[]
   private _ob?: MutationObserver | null = null
-  private _root: Element | ShadowRoot
+  /**
+   * @internal
+   */
+  public _root: Element | ShadowRoot
   private _slots?: Record<string, Node[]>
 
   constructor(
@@ -247,6 +251,7 @@ export class VueElement extends BaseClass {
           this._ob = null
         }
         render(null, this._root)
+        this._instance!.isCE = undefined
         this._instance = null
       }
     })
@@ -395,7 +400,7 @@ export class VueElement extends BaseClass {
     if (!this._instance) {
       vnode.ce = instance => {
         this._instance = instance
-        instance.isCE = true
+        instance.isCE = this
         // HMR
         if (__DEV__) {
           instance.ceReload = newStyles => {
@@ -508,3 +513,25 @@ export class VueElement extends BaseClass {
     }
   }
 }
+
+/**
+ * Retrieve the shadowRoot of the current custom element. Only usable in setup()
+ * of a `defineCustomElement` component.
+ */
+export function useShadowRoot(): ShadowRoot | null {
+  const instance = getCurrentInstance()
+  const el = instance && instance.isCE
+  if (el) {
+    return el.shadowRoot
+  } else if (__DEV__) {
+    if (!instance) {
+      warn(`useCustomElementRoot called without an active component instance.`)
+    } else {
+      warn(
+        `useCustomElementRoot can only be used in components defined via ` +
+          `defineCustomElement.`,
+      )
+    }
+  }
+  return null
+}
diff --git a/packages/runtime-dom/src/index.ts b/packages/runtime-dom/src/index.ts
index 989a5fd3b80..7acaa98c9d5 100644
--- a/packages/runtime-dom/src/index.ts
+++ b/packages/runtime-dom/src/index.ts
@@ -242,6 +242,7 @@ function normalizeContainer(
 export {
   defineCustomElement,
   defineSSRCustomElement,
+  useShadowRoot,
   VueElement,
   type VueElementConstructor,
 } from './apiCustomElement'