lib, vpk - improve HMR for signals, context and FCs

packages/lib/src/constants.ts

@@ -1,4 +1,5 @@
 export const $SIGNAL = Symbol.for("kaioken.signal")
+export const $CONTEXT = Symbol.for("kaioken.context")
 export const $CONTEXT_PROVIDER = Symbol.for("kaioken.contextProvider")
 export const $FRAGMENT = Symbol.for("kaioken.fragment")
 export const $KAIOKEN_ERROR = Symbol.for("kaioken.error")

packages/lib/src/context.ts

@@ -1,8 +1,12 @@
-import { $CONTEXT_PROVIDER } from "./constants.js"
+import { $CONTEXT, $CONTEXT_PROVIDER, $HMR_ACCEPT } from "./constants.js"
 import { createElement } from "./element.js"
+import { __DEV__ } from "./env.js"
+import { GenericHMRAcceptor } from "./hmr.js"
+import { traverseApply } from "./utils.js"
 
 export function createContext<T>(defaultValue: T): Kaioken.Context<T> {
   const ctx: Kaioken.Context<T> = {
+    [$CONTEXT]: true,
     Provider: ({ value, children }: Kaioken.ProviderProps<T>) => {
       return createElement(
         $CONTEXT_PROVIDER,
@@ -18,5 +22,33 @@ export function createContext<T>(defaultValue: T): Kaioken.Context<T> {
       return this.Provider.displayName || "Anonymous Context"
     },
   }
+  if (__DEV__) {
+    const asHmrAcceptor = ctx as any as GenericHMRAcceptor<Kaioken.Context<T>>
+    asHmrAcceptor[$HMR_ACCEPT] = {
+      inject: (prev) => {
+        const newProvider = ctx.Provider
+        window.__kaioken!.apps.forEach((ctx) => {
+          if (!ctx.mounted || !ctx.rootNode) return
+          traverseApply(ctx.rootNode, (vNode) => {
+            if (vNode.type === prev.Provider) {
+              vNode.type = newProvider
+              vNode.hmrUpdated = true
+              if (vNode.prev) {
+                vNode.prev.type = newProvider
+              }
+              ctx.requestUpdate(vNode)
+            }
+          })
+        })
+      },
+      destroy: () => {},
+      provide: () => ctx,
+    }
+  }
+
   return ctx
 }
+
+export function isContext<T>(thing: unknown): thing is Kaioken.Context<T> {
+  return typeof thing === "object" && !!thing && $CONTEXT in thing
+}

packages/lib/src/hmr.ts

@@ -15,7 +15,7 @@ export type GenericHMRAcceptor<T = {}> = {
   [$HMR_ACCEPT]: HMRAccept<T>
 }
 
-type HotVar = Kaioken.FC | Store<any, any> | Signal<any>
+type HotVar = Kaioken.FC | Store<any, any> | Signal<any> | Kaioken.Context<any>
 
 export function isGenericHmrAcceptor(
   thing: unknown
@@ -38,7 +38,6 @@ type ModuleMemory = {
 export function createHMRContext() {
   type FilePath = string
   const moduleMap = new Map<FilePath, ModuleMemory>()
-
   let currentModuleMemory: ModuleMemory | null = null
   let isModuleReplacementExecution = false
   const isReplacement = () => isModuleReplacementExecution
@@ -65,8 +64,18 @@ export function createHMRContext() {
 
     for (const [name, newVar] of Object.entries(hotVars)) {
       const oldVar = currentModuleMemory.hotVars.get(name)
-      // @ts-ignore
-      newVar.__devtoolsFileLink = currentModuleMemory.fileLink + ":0"
+      if (typeof newVar === "function") {
+        // @ts-ignore - this is how we tell devtools what file the component is from
+        newVar.__devtoolsFileLink = currentModuleMemory.fileLink + ":0"
+        if (oldVar) {
+          /**
+           * this is how, when the previous function has been stored somewhere else (eg. by Vike),
+           * we can trace it to its latest version
+           */
+          // @ts-ignore
+          oldVar.__next = newVar
+        }
+      }
       currentModuleMemory.hotVars.set(name, newVar)
       if (!oldVar) continue
       if (isGenericHmrAcceptor(oldVar) && isGenericHmrAcceptor(newVar)) {

packages/lib/src/reconciler.ts

@@ -1,7 +1,7 @@
 import type { AppContext } from "./appContext"
 import { ELEMENT_TYPE, FLAG, $FRAGMENT } from "./constants.js"
 import { ctx } from "./globals.js"
-import { isVNode } from "./utils.js"
+import { isVNode, latest } from "./utils.js"
 import { Signal } from "./signals/base.js"
 import { __DEV__ } from "./env.js"
 import { createElement, Fragment } from "./element.js"
@@ -215,7 +215,10 @@ function updateTextNode(
 }
 
 function updateNode(parent: VNode, oldNode: VNode | null, newNode: VNode) {
-  const nodeType = newNode.type
+  let nodeType = newNode.type
+  if (typeof nodeType === "function") {
+    nodeType = latest(nodeType)
+  }
   if (nodeType === $FRAGMENT) {
     return updateFragment(
       parent,

packages/lib/src/signals/base.ts

@@ -3,6 +3,7 @@ import { __DEV__ } from "../env.js"
 import type { HMRAccept } from "../hmr.js"
 import {
   getVNodeAppContext,
+  latest,
   safeStringify,
   sideEffectsEnabled,
 } from "../utils.js"
@@ -18,6 +19,7 @@ export class Signal<T> {
   protected $id: string
   protected $value: T
   protected $initialValue?: string
+  protected __next?: Signal<T>
   constructor(initial: T, displayName?: string) {
     this.$id = crypto.randomUUID()
     signalSubsMap.set(this.$id, new Set())
@@ -38,34 +40,42 @@ export class Signal<T> {
           signalSubsMap.get(this.$id)?.clear?.()
           signalSubsMap.delete(this.$id)
           this.$id = prev.$id
+          // @ts-ignore - this handles scenarios where a reference to the prev has been encapsulated
+          // and we need to be able to refer to the latest version of the signal.
+          prev.__next = this
         },
         destroy: () => {},
       } satisfies HMRAccept<Signal<any>>
     }
   }
 
   get value() {
-    Signal.entangle(this)
-    return this.$value
+    const tgt = latest(this)
+    Signal.entangle(tgt)
+    return tgt.$value
   }
 
   set value(next: T) {
-    if (Object.is(this.$value, next)) return
-    this.$value = next
-    this.notify()
+    const tgt = latest(this)
+    if (Object.is(tgt.$value, next)) return
+    tgt.$value = next
+    tgt.notify()
   }
 
   peek() {
-    return this.$value
+    const tgt = latest(this)
+    return tgt.$value
   }
 
   sneak(newValue: T) {
-    this.$value = newValue
+    const tgt = latest(this)
+    tgt.$value = newValue
   }
 
   toString() {
-    Signal.entangle(this)
-    return `${this.$value}`
+    const tgt = latest(this)
+    Signal.entangle(tgt)
+    return `${tgt.$value}`
   }
 
   subscribe(cb: (state: T) => void): () => void {
@@ -75,10 +85,11 @@ export class Signal<T> {
   }
 
   notify(options?: { filter?: (sub: Function | Kaioken.VNode) => boolean }) {
+    const tgt = latest(this)
     signalSubsMap.get(this.$id)?.forEach((sub) => {
       if (options?.filter && !options.filter(sub)) return
       if (typeof sub === "function") {
-        return sub(this.$value)
+        return sub(tgt.$value)
       }
       getVNodeAppContext(sub).requestUpdate(sub)
     })
@@ -173,6 +184,7 @@ export const useSignal = <T>(initial: T, displayName?: string) => {
             },
           }
           if (hook.signal && vNode.hmrUpdated) {
+            console.log("signal hook hmr updated (initial changed)")
             hook.signal.value = initial
           }
         }

packages/lib/src/types.ts

@@ -3,7 +3,7 @@ import type {
   Signal as SignalClass,
   SignalLike,
 } from "./signals"
-import type { $CONTEXT_PROVIDER, $FRAGMENT } from "./constants"
+import type { $CONTEXT, $CONTEXT_PROVIDER, $FRAGMENT } from "./constants"
 import type { KaiokenGlobalContext } from "./globalContext"
 import type {
   EventAttributes,
@@ -130,6 +130,7 @@ declare global {
       children?: JSX.Children | ((value: T) => JSX.Element)
     }
     type Context<T> = {
+      [$CONTEXT]: true
       Provider: (({ value, children }: ProviderProps<T>) => JSX.Element) & {
         displayName?: string
       }

packages/lib/src/utils.ts

@@ -4,6 +4,7 @@ import { unwrap } from "./signals/utils.js"
 import { KaiokenError } from "./error.js"
 import type { AppContext } from "./appContext"
 import type { ExoticVNode } from "./types.utils"
+import { __DEV__ } from "./env.js"
 
 export {
   isVNode,
@@ -25,6 +26,7 @@ export {
   sideEffectsEnabled,
   encodeHtmlEntities,
   noop,
+  latest,
   propFilters,
   selfClosingTags,
   svgTags,
@@ -36,6 +38,20 @@ type VNode = Kaioken.VNode
 
 const noop: () => void = Object.freeze(() => {})
 
+/**
+ * This is a no-op in production. It is used to get the latest
+ * iteration of a component or signal after HMR has happened.
+ */
+function latest<T>(thing: T): T {
+  let tgt: any = thing
+  if (__DEV__) {
+    while ("__next" in tgt) {
+      tgt = tgt.__next as typeof tgt
+    }
+  }
+  return tgt
+}
+
 /**
  * Returns true if called during DOM or hydration render mode.
  */

packages/vite-plugin-kaioken/src/index.ts

@@ -226,13 +226,18 @@ function findHotVars(nodes: AstNode[], _id: string): string[] {
     createAliasBuilder("kaioken", "computed")
   const { addAliases: addWatchAliases, nodeContainsAliasCall: isWatch } =
     createAliasBuilder("kaioken", "watch")
+  const {
+    addAliases: addCreateContextAliases,
+    nodeContainsAliasCall: isContext,
+  } = createAliasBuilder("kaioken", "createContext")
 
   for (const node of nodes) {
     if (node.type === "ImportDeclaration") {
       addCreateStoreAliases(node)
       addSignalAliases(node)
       addComputedAliases(node)
       addWatchAliases(node)
+      addCreateContextAliases(node)
       continue
     }
 
@@ -245,6 +250,7 @@ function findHotVars(nodes: AstNode[], _id: string): string[] {
       addHotVarNames(node, hotVarNames)
       continue
     }
+
     if (findNode(node, isSignal)) {
       addHotVarNames(node, hotVarNames)
       continue
@@ -258,6 +264,11 @@ function findHotVars(nodes: AstNode[], _id: string): string[] {
     if (findNode(node, isWatch)) {
       addHotVarNames(node, hotVarNames)
     }
+
+    if (findNode(node, isContext)) {
+      addHotVarNames(node, hotVarNames)
+      continue
+    }
   }
   return Array.from(hotVarNames)
 }

sandbox/csr/src/App.tsx

@@ -7,7 +7,6 @@ import {
   useCallback,
   useComputed,
   useSignal,
-  useWatch,
   signal,
   watch,
 } from "kaioken"

sandbox/ssr/src/components/Counter.tsx

@@ -1,8 +1,8 @@
-import { signal } from "kaioken"
+import { signal, useSignal } from "kaioken"
 
 export default function Counter() {
   console.log("Counter")
-  const count = signal(0)
+  const count = useSignal(0)
 
   return (
     <>

sandbox/ssr/src/pages/counter/+Page.tsx

@@ -1,14 +1,15 @@
 import Counter from "$/components/Counter"
 import { PageTitle } from "$/components/PageTitle"
-import { computed, signal, useEffect } from "kaioken"
+import { computed, signal, useComputed, useEffect, useSignal } from "kaioken"
 
 export { Page }
 
+const id = signal(23)
+
 function Page() {
-  const id = signal(0)
-  const divId = computed(() => id.value.toString(), "divId")
+  const divId = useComputed(() => `counter-${id}`, "divId")
   useEffect(() => {
-    const interval = setInterval(() => (id.value += 2), 1000)
+    const interval = setInterval(() => (id.value += 1), 1000)
     return () => clearInterval(interval)
   }, [])
 

sandbox/ssr/src/pages/index/+Page.tsx

@@ -1,7 +1,10 @@
+import { usePageContext } from "$/context/pageContext"
+
 export function Page() {
+  const { isClient } = usePageContext()
   return (
     <div className="p-6">
-      <h1>Hello, world!</h1>
+      <h1>Hello, world! isClient: {`${isClient}`}</h1>
     </div>
   )
 }

sandbox/ssr/src/pages/products/+Page.tsx

@@ -4,7 +4,7 @@ import type { ServerProps } from "./+data"
 export function Page({ products }: ServerProps) {
   return (
     <>
-      <PageTitle>Products</PageTitle>
+      <PageTitle>Product 123s</PageTitle>
       <div>
         {products.map((product) => (
           <ProductCard product={product} />

sandbox/ssr/src/renderer/+onRenderClient.tsx

@@ -5,16 +5,20 @@ import type { AppContext } from "kaioken"
 import { getTitle } from "./utils"
 import { App } from "./App"
 
-let appContext: AppContext<{ pageContext: PageContextClient }> | undefined
+declare global {
+  interface Window {
+    __appContext: AppContext<{ pageContext: PageContextClient }> | undefined
+  }
+}
 
 export const onRenderClient: OnRenderClientAsync = async (pageContext) => {
   const container = document.getElementById("page-root")!
 
-  if (pageContext.isHydration || !appContext) {
-    appContext = await hydrate(App, container, { pageContext })
+  if (pageContext.isHydration || !window.__appContext) {
+    window.__appContext = await hydrate(App, container, { pageContext })
     return
   }
 
   document.title = getTitle(pageContext)
-  await appContext.setProps(() => ({ pageContext }))
+  await window.__appContext.setProps(() => ({ pageContext }))
 }