refactor(runtime-vapor): improve DynamicFragment context restoration

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

@@ -13,7 +13,9 @@ import {
   toDisplayString,
 } from '@vue/runtime-dom'
 import {
+  VaporTransition,
   createComponent,
+  createIf,
   createSlot,
   createTextNode,
   createVaporApp,
@@ -401,6 +403,60 @@ describe('api: provide/inject', () => {
     expect(host.innerHTML).toBe('hello<!--slot-->')
   })
 
+  it('transition out-in deferred branch should keep parent inject context', async () => {
+    const toggle = ref(true)
+
+    const ChildA = defineVaporComponent({
+      setup() {
+        const foo = inject('foo', 'missing')
+        const n0 = template('<div></div>')()
+        setElementText(n0, `A:${foo}`)
+        return n0
+      },
+    })
+
+    const ChildB = defineVaporComponent({
+      setup() {
+        const foo = inject('foo', 'missing')
+        const n0 = template('<div></div>')()
+        setElementText(n0, `B:${foo}`)
+        return n0
+      },
+    })
+
+    const Parent = defineVaporComponent({
+      setup() {
+        provide('foo', 'from-parent')
+        const onLeave = (_: any, done: Function) => setTimeout(done, 0)
+        return createComponent(
+          VaporTransition,
+          {
+            mode: () => 'out-in',
+            onLeave: () => onLeave,
+          },
+          {
+            default: withVaporCtx(() =>
+              createIf(
+                () => toggle.value,
+                () => createComponent(ChildA),
+                () => createComponent(ChildB),
+              ),
+            ),
+          },
+        )
+      },
+    })
+
+    const { host } = define(Parent).render()
+    expect(host.textContent).toContain('A:from-parent')
+
+    toggle.value = false
+    await nextTick()
+    await new Promise(r => setTimeout(r, 0))
+
+    expect(host.textContent).toContain('B:from-parent')
+  })
+
   describe('hasInjectionContext', () => {
     it('should be false outside of setup', () => {
       expect(hasInjectionContext()).toBe(false)

packages/runtime-vapor/src/fragment.ts

@@ -192,24 +192,22 @@ export class DynamicFragment extends VaporFragment {
         this.scope = new EffectScope()
       }
 
-      // restore slot owner
-      const prevOwner = setCurrentSlotOwner(this.slotOwner)
+      const prevSlotOwner = setCurrentSlotOwner(this.slotOwner)
       // set currentKeepAliveCtx so nested DynamicFragments and components can capture it
-      const prevCtx = setCurrentKeepAliveCtx(keepAliveCtx)
-      let prevBranchKey: any
-      if (keepAliveCtx && this.keyed) {
-        prevBranchKey = keepAliveCtx.setCurrentBranchKey(this.current)
+      const prevKeepAliveCtx = setCurrentKeepAliveCtx(keepAliveCtx)
+      const prevBranchKey =
+        keepAliveCtx && this.keyed
+          ? keepAliveCtx.setCurrentBranchKey(this.current)
+          : undefined
+      const prevInstance = setCurrentInstance(instance)
+      try {
+        this.nodes = this.scope.run(render) || []
+      } finally {
+        setCurrentInstance(...prevInstance)
+        if (prevBranchKey) keepAliveCtx!.setCurrentBranchKey(prevBranchKey)
+        setCurrentKeepAliveCtx(prevKeepAliveCtx)
+        setCurrentSlotOwner(prevSlotOwner)
       }
-      // switch current instance to parent instance during update
-      // ensure that the parent instance is correct for nested components
-      const prev = parent && instance ? setCurrentInstance(instance) : undefined
-      this.nodes = this.scope.run(render) || []
-      if (prev !== undefined) setCurrentInstance(...prev)
-      if (keepAliveCtx && this.keyed) {
-        keepAliveCtx.setCurrentBranchKey(prevBranchKey)
-      }
-      setCurrentKeepAliveCtx(prevCtx)
-      setCurrentSlotOwner(prevOwner)
 
       // set key on blocks
       if (this.keyed) setKey(this.nodes, this.current)