refactor: wrap the ctx when initializing the slots

packages/runtime-vapor/__tests__/componentSlots.spec.ts

@@ -8,7 +8,6 @@ import {
   nextTick,
   ref,
   template,
-  withCtx,
 } from '../src'
 import { makeRender } from './_utils'
 
@@ -35,21 +34,6 @@ function renderWithSlots(slots: any): any {
 }
 
 describe('component: slots', () => {
-  test('initSlots: instance.slots should be set correctly', () => {
-    const { slots } = renderWithSlots({ _: 1 })
-    expect(slots).toMatchObject({ _: 1 })
-  })
-
-  // NOTE: slot normalization is not supported
-  test.todo(
-    'initSlots: should normalize object slots (when value is null, string, array)',
-    () => {},
-  )
-  test.todo(
-    'initSlots: should normalize object slots (when value is function)',
-    () => {},
-  )
-
   test('initSlots: instance.slots should be set correctly', () => {
     let instance: any
     const Comp = defineComponent({
@@ -74,6 +58,16 @@ describe('component: slots', () => {
     )
   })
 
+  // NOTE: slot normalization is not supported
+  test.todo(
+    'initSlots: should normalize object slots (when value is null, string, array)',
+    () => {},
+  )
+  test.todo(
+    'initSlots: should normalize object slots (when value is function)',
+    () => {},
+  )
+
   // runtime-core's "initSlots: instance.slots should be set correctly (when vnode.shapeFlag is not SLOTS_CHILDREN)"
   test('initSlots: instance.slots should be set correctly', () => {
     const { slots } = renderWithSlots({
@@ -154,12 +148,16 @@ describe('component: slots', () => {
   })
 
   test('the current instance should be kept in the slot', async () => {
-    let instanceInSlot: any
+    let instanceInDefaultSlot: any
+    let instanceInVForSlot: any
+    let instanceInVIfSlot: any
 
     const Comp = defineComponent({
       render() {
         const instance = getCurrentInstance()
         instance!.slots.default!()
+        instance!.slots.inVFor!()
+        instance!.slots.inVIf!()
         return template('<div></div>')()
       },
     })
@@ -170,16 +168,37 @@ describe('component: slots', () => {
           Comp,
           {},
           {
-            default: withCtx(() => {
-              instanceInSlot = getCurrentInstance()
+            default: () => {
+              instanceInDefaultSlot = getCurrentInstance()
               return template('content')()
-            }),
+            },
           },
+          () => [
+            [
+              {
+                name: 'inVFor',
+                fn: () => {
+                  instanceInVForSlot = getCurrentInstance()
+                  return template('content')()
+                },
+              },
+            ],
+            {
+              name: 'inVIf',
+              key: '1',
+              fn: () => {
+                instanceInVIfSlot = getCurrentInstance()
+                return template('content')()
+              },
+            },
+          ],
         )
       },
     }).render()
 
-    expect(instanceInSlot).toBe(instance)
+    expect(instanceInDefaultSlot).toBe(instance)
+    expect(instanceInVForSlot).toBe(instance)
+    expect(instanceInVIfSlot).toBe(instance)
   })
 
   test.todo('should respect $stable flag', async () => {

packages/runtime-vapor/src/componentSlots.ts

@@ -1,6 +1,6 @@
-import { type IfAny, extend, isArray } from '@vue/shared'
+import { type IfAny, isArray } from '@vue/shared'
 import { baseWatch } from '@vue/reactivity'
-import type { ComponentInternalInstance } from './component'
+import { type ComponentInternalInstance, setCurrentInstance } from './component'
 import type { Block } from './apiRender'
 import { createVaporPreScheduler } from './scheduler'
 
@@ -29,7 +29,14 @@ export const initSlots = (
   rawSlots: InternalSlots | null = null,
   dynamicSlots: DynamicSlots | null = null,
 ) => {
-  const slots: InternalSlots = extend({}, rawSlots)
+  const slots: InternalSlots = {}
+
+  for (const key in rawSlots) {
+    const slot = rawSlots[key]
+    if (slot) {
+      slots[key] = withCtx(slot)
+    }
+  }
 
   if (dynamicSlots) {
     const dynamicSlotKeys: Record<string, true> = {}
@@ -41,20 +48,22 @@ export const initSlots = (
           // array of dynamic slot generated by <template v-for="..." #[...]>
           if (isArray(slot)) {
             for (let j = 0; j < slot.length; j++) {
-              slots[slot[j].name] = slot[j].fn
+              slots[slot[j].name] = withCtx(slot[j].fn)
               dynamicSlotKeys[slot[j].name] = true
             }
           } else if (slot) {
             // conditional single slot generated by <template v-if="..." #foo>
-            slots[slot.name] = slot.key
-              ? (...args: any[]) => {
-                  const res = slot.fn(...args)
-                  // attach branch key so each conditional branch is considered a
-                  // different fragment
-                  if (res) (res as any).key = slot.key
-                  return res
-                }
-              : slot.fn
+            slots[slot.name] = withCtx(
+              slot.key
+                ? (...args: any[]) => {
+                    const res = slot.fn(...args)
+                    // attach branch key so each conditional branch is considered a
+                    // different fragment
+                    if (res) (res as any).key = slot.key
+                    return res
+                  }
+                : slot.fn,
+            )
             dynamicSlotKeys[slot.name] = true
           }
         }
@@ -77,4 +86,15 @@ export const initSlots = (
   }
 
   instance.slots = slots
+
+  function withCtx(fn: Slot): Slot {
+    return (...args: any[]) => {
+      const reset = setCurrentInstance(instance.parent!)
+      try {
+        return fn(...args)
+      } finally {
+        reset()
+      }
+    }
+  }
 }

packages/runtime-vapor/src/index.ts

@@ -50,7 +50,6 @@ export {
   type FunctionalComponent,
   type SetupFn,
 } from './component'
-export { withCtx } from './componentRenderContext'
 export { renderEffect } from './renderEffect'
 export {
   watch,