feat(runtime-vapor): support custom directives on vapor components

packages/runtime-vapor/__tests__/directives/customDirective.spec.ts

@@ -1,6 +1,12 @@
 import { effectScope, ref } from '@vue/reactivity'
-import { type VaporDirective, withVaporDirectives } from '../../src'
+import {
+  type VaporDirective,
+  createComponent,
+  defineVaporComponent,
+  withVaporDirectives,
+} from '../../src'
 import { nextTick, watchEffect } from '@vue/runtime-dom'
+import type { Mock } from 'vitest'
 
 describe('custom directive', () => {
   it('should work', async () => {
@@ -36,4 +42,68 @@ describe('custom directive', () => {
     // should be stopped and not update
     expect(el.textContent).toBe('2')
   })
+
+  it('should work on single root component', async () => {
+    const teardown = vi.fn()
+    const dir: VaporDirective = vi.fn((el, source) => {
+      watchEffect(() => {
+        el.textContent = source()
+      })
+      return teardown
+    })
+    const scope = effectScope()
+    const n = ref(1)
+    const source = () => n.value
+
+    // Child component with single root
+    const Child = defineVaporComponent({
+      render() {
+        const el = document.createElement('div')
+        return el
+      },
+    })
+
+    const root = document.createElement('div')
+
+    scope.run(() => {
+      const instance = createComponent(Child)
+      withVaporDirectives(instance, [[dir, source]])
+      root.appendChild(instance.block as Node)
+    })
+
+    // Should resolve to the div element inside Child
+    expect(dir).toHaveBeenCalled()
+    const el = (dir as unknown as Mock).mock.calls[0][0]
+    expect(el).toBeInstanceOf(HTMLDivElement)
+    expect(el.textContent).toBe('1')
+
+    n.value = 2
+    await nextTick()
+    expect(el.textContent).toBe('2')
+
+    scope.stop()
+    expect(teardown).toHaveBeenCalled()
+  })
+
+  it('should warn on multi-root component', () => {
+    const dir: VaporDirective = vi.fn()
+    const scope = effectScope()
+
+    // Child component with multiple roots
+    const Child = defineVaporComponent({
+      render() {
+        return [document.createElement('div'), document.createElement('span')]
+      },
+    })
+
+    scope.run(() => {
+      const instance = createComponent(Child)
+      withVaporDirectives(instance, [[dir]])
+    })
+
+    expect(dir).not.toHaveBeenCalled()
+    expect(
+      'Runtime directive used on component with non-element root node',
+    ).toHaveBeenWarned()
+  })
 })

packages/runtime-vapor/src/component.ts

@@ -99,7 +99,7 @@ import {
   isLastInsertion,
   resetInsertionState,
 } from './insertionState'
-import { DynamicFragment } from './fragment'
+import { DynamicFragment, isFragment } from './fragment'
 import type { VaporElement } from './apiDefineVaporCustomElement'
 
 export { currentInstance } from '@vue/runtime-dom'
@@ -415,7 +415,7 @@ export function applyFallthroughProps(
   block: Block,
   attrs: Record<string, any>,
 ): void {
-  const el = getRootElement(block)
+  const el = getRootElement(block, false)
   if (el) {
     isApplyingFallthroughProps = true
     setDynamicProps(el, [attrs])
@@ -820,16 +820,24 @@ export function getExposed(
   }
 }
 
-function getRootElement(block: Block): Element | undefined {
+export function getRootElement(
+  block: Block,
+  recurse: boolean = true,
+): Element | undefined {
   if (block instanceof Element) {
     return block
   }
 
-  if (block instanceof DynamicFragment) {
+  if (recurse && isVaporComponent(block)) {
+    return getRootElement(block.block, recurse)
+  }
+
+  if (isFragment(block)) {
     const { nodes } = block
     if (nodes instanceof Element && (nodes as any).$root) {
       return nodes
     }
+    return getRootElement(nodes, recurse)
   }
 
   // The root node contains comments. It is necessary to filter out
@@ -843,7 +851,7 @@ function getRootElement(block: Block): Element | undefined {
         hasComment = true
         continue
       }
-      const thisRoot = getRootElement(b)
+      const thisRoot = getRootElement(b, recurse)
       // only return root if there is exactly one eligible root in the array
       if (!thisRoot || singleRoot) {
         return

packages/runtime-vapor/src/directives/custom.ts

@@ -1,5 +1,9 @@
-import { type DirectiveModifiers, onScopeDispose } from '@vue/runtime-dom'
-import type { VaporComponentInstance } from '../component'
+import { type DirectiveModifiers, onScopeDispose, warn } from '@vue/runtime-dom'
+import {
+  type VaporComponentInstance,
+  getRootElement,
+  isVaporComponent,
+} from '../component'
 
 // !! vapor directive is different from vdom directives
 export type VaporDirective = (
@@ -25,10 +29,20 @@ export function withVaporDirectives(
   node: Element | VaporComponentInstance,
   dirs: VaporDirectiveArguments,
 ): void {
-  // TODO handle custom directive on component
+  const element = isVaporComponent(node) ? getRootElement(node.block) : node
+  if (!element) {
+    if (__DEV__) {
+      warn(
+        `Runtime directive used on component with non-element root node. ` +
+          `The directives will not function as intended.`,
+      )
+    }
+    return
+  }
+
   for (const [dir, value, argument, modifiers] of dirs) {
     if (dir) {
-      const ret = dir(node, value, argument, modifiers)
+      const ret = dir(element, value, argument, modifiers)
       if (ret) onScopeDispose(ret)
     }
   }