refactor hydration

packages/lib/src/appContext.ts

@@ -36,17 +36,18 @@ export class AppContext<T extends Record<string, unknown> = {}> {
   }
 
   mount() {
-    this.scheduler = new Scheduler(this, this.options?.maxFrameMs ?? 50)
-    this.rootNode = createElement(
-      this.root!.nodeName.toLowerCase(),
-      {},
-      createElement(this.appFunc, this.appProps)
-    )
-    this.rootNode.dom = this.root
-    this.scheduler.queueUpdate(this.rootNode)
-    this.scheduler.wake()
     return new Promise<AppContext<T>>((resolve) => {
-      this.scheduler!.nextIdle(() => {
+      if (this.mounted) return resolve(this)
+      this.scheduler = new Scheduler(this, this.options?.maxFrameMs ?? 50)
+      this.rootNode = createElement(
+        this.root!.nodeName.toLowerCase(),
+        {},
+        createElement(this.appFunc, this.appProps)
+      )
+      this.rootNode.dom = this.root
+      this.scheduler.queueUpdate(this.rootNode)
+      this.scheduler.wake()
+      this.scheduler.nextIdle(() => {
         this.mounted = true
         window.__kaioken?.emit("mount", this as AppContext<any>)
         resolve(this)
@@ -55,8 +56,8 @@ export class AppContext<T extends Record<string, unknown> = {}> {
   }
 
   unmount() {
-    if (!this.mounted) return this
     return new Promise<AppContext<T>>((resolve) => {
+      if (!this.mounted) return resolve(this)
       if (!this.rootNode?.child) return resolve(this)
       this.requestDelete(this.rootNode.child)
 
@@ -74,7 +75,7 @@ export class AppContext<T extends Record<string, unknown> = {}> {
     const rootChild = this.rootNode?.child
     const scheduler = this.scheduler
     if (!this.mounted || !rootChild || !scheduler)
-      return console.error(
+      throw new Error(
         "[kaioken]: failed to apply new props - ensure the app is mounted"
       )
     return new Promise<AppContext<T>>((resolve) => {

packages/lib/src/dom.ts

@@ -10,8 +10,9 @@ import { EffectTag, elementTypes } from "./constants.js"
 import { Component } from "./component.js"
 import { Signal } from "./signal.js"
 import { renderMode } from "./globals.js"
+import { hydrationStack } from "./hydration.js"
 
-export { commitWork, createDom, updateDom }
+export { commitWork, createDom, updateDom, hydrateDom }
 
 type VNode = Kaioken.VNode
 type DomParentSearchResult = {
@@ -23,15 +24,68 @@ type CommitStackItem = [VNode, MaybeDom, DomParentSearchResult | undefined]
 
 function createDom(vNode: VNode): HTMLElement | SVGElement | Text {
   const t = vNode.type as string
-  let dom =
-    t == elementTypes.text
-      ? document.createTextNode(vNode.props?.nodeValue ?? "")
-      : svgTags.includes(t)
-        ? document.createElementNS("http://www.w3.org/2000/svg", t)
-        : document.createElement(t)
-
-  vNode.dom = updateDom(vNode, dom)
-  return dom
+  return t == elementTypes.text
+    ? document.createTextNode(vNode.props.nodeValue ?? "")
+    : svgTags.includes(t)
+      ? document.createElementNS("http://www.w3.org/2000/svg", t)
+      : document.createElement(t)
+}
+
+function updateDom(node: VNode) {
+  if (node.instance?.doNotModifyDom) return
+  const dom = node.dom as HTMLElement | SVGElement | Text
+  const prevProps: Record<string, any> = node.prev?.props ?? {}
+  const nextProps: Record<string, any> = node.props ?? {}
+
+  const keys = new Set([...Object.keys(prevProps), ...Object.keys(nextProps)])
+
+  keys.forEach((key) => {
+    if (key === "innerHTML") {
+      return setInnerHTML(node.dom as any, nextProps[key], prevProps[key])
+    }
+
+    if (propFilters.internalProps.includes(key)) return
+
+    if (
+      propFilters.isEvent(key) &&
+      (prevProps[key] !== nextProps[key] || renderMode.current === "hydrate")
+    ) {
+      const eventType = key.toLowerCase().substring(2)
+      if (key in prevProps) dom.removeEventListener(eventType, prevProps[key])
+      if (key in nextProps) dom.addEventListener(eventType, nextProps[key])
+      return
+    }
+
+    if (!(dom instanceof Text) && prevProps[key] !== nextProps[key]) {
+      setProp(dom, key, nextProps[key], prevProps[key])
+      return
+    }
+    if (node.prev?.props && prevProps.nodeValue !== nextProps.nodeValue) {
+      dom.nodeValue = nextProps.nodeValue
+    }
+  })
+}
+
+function hydrateDom(vNode: VNode) {
+  const dom = hydrationStack.nextChild()
+  const nodeName = dom?.nodeName.toLowerCase()
+  if ((vNode.type as string) !== nodeName) {
+    throw new Error(
+      `[kaioken]: Hydration mismatch - expected node of type ${vNode.type} but received ${nodeName}`
+    )
+  }
+  vNode.dom = dom
+  if (vNode.type !== elementTypes.text) {
+    updateDom(vNode)
+    return
+  }
+  let prev = vNode
+  let sibling = vNode.sibling
+  while (sibling && sibling.type === elementTypes.text) {
+    sibling.dom = (prev.dom as Text)!.splitText(prev.props.nodeValue.length)
+    prev = sibling
+    sibling = sibling.sibling
+  }
 }
 
 function handleAttributeRemoval(
@@ -133,42 +187,6 @@ function setStyleProp(
   }
 }
 
-function updateDom(node: VNode, dom: HTMLElement | SVGElement | Text) {
-  if (node.instance?.doNotModifyDom) return node.dom
-  const prevProps: Record<string, any> = node.prev?.props ?? {}
-  const nextProps: Record<string, any> = node.props ?? {}
-
-  const keys = new Set([...Object.keys(prevProps), ...Object.keys(nextProps)])
-
-  keys.forEach((key) => {
-    if (key === "innerHTML") {
-      return setInnerHTML(dom as any, nextProps[key], prevProps[key])
-    }
-
-    if (propFilters.internalProps.includes(key)) return
-
-    if (
-      propFilters.isEvent(key) &&
-      (prevProps[key] !== nextProps[key] || renderMode.current === "hydrate")
-    ) {
-      const eventType = key.toLowerCase().substring(2)
-      if (key in prevProps) dom.removeEventListener(eventType, prevProps[key])
-      if (key in nextProps) dom.addEventListener(eventType, nextProps[key])
-      return
-    }
-
-    if (!(dom instanceof Text) && prevProps[key] !== nextProps[key]) {
-      setProp(dom, key, nextProps[key], prevProps[key])
-      return
-    }
-    if (node.prev?.props && prevProps.nodeValue !== nextProps.nodeValue) {
-      dom.nodeValue = nextProps.nodeValue
-    }
-  })
-
-  return dom
-}
-
 function getDomParent(node: VNode): DomParentSearchResult {
   let domParentNode: VNode | undefined = node.parent ?? node.prev?.parent
   let domParent = domParentNode?.dom
@@ -192,10 +210,10 @@ function getDomParent(node: VNode): DomParentSearchResult {
 
 function placeDom(
   vNode: VNode,
-  dom: HTMLElement | SVGElement | Text,
   prevSiblingDom: MaybeDom,
   mntParent: DomParentSearchResult
 ) {
+  const dom = vNode.dom as HTMLElement | SVGElement | Text
   if (prevSiblingDom) {
     prevSiblingDom.after(dom)
     return
@@ -248,7 +266,7 @@ function commitWork(ctx: AppContext, vNode: VNode) {
     const dom = n.dom
 
     if (dom) {
-      mntParent = commitDom(n, dom, prevSiblingDom, mntParent) || mntParent
+      mntParent = commitDom(n, prevSiblingDom, mntParent) || mntParent
     } else if (n.effectTag === EffectTag.PLACEMENT) {
       // propagate the effect to children
       let c = n.child
@@ -296,10 +314,10 @@ function commitWork(ctx: AppContext, vNode: VNode) {
 
 function commitDom(
   n: VNode,
-  dom: HTMLElement | SVGElement | Text,
   prevSiblingDom: MaybeDom,
   mntParent: DomParentSearchResult | undefined
 ) {
+  const dom = n.dom as HTMLElement | SVGElement | Text
   if (renderMode.current === "hydrate") {
     if (n.props.ref) {
       n.props.ref.current = dom
@@ -309,10 +327,10 @@ function commitDom(
   if (n.instance?.doNotModifyDom) return
   if (!dom.isConnected || n.effectTag === EffectTag.PLACEMENT) {
     const p = mntParent ?? getDomParent(n)
-    placeDom(n, dom, prevSiblingDom, p)
+    placeDom(n, prevSiblingDom, p)
     return p
   } else if (n.effectTag === EffectTag.UPDATE) {
-    updateDom(n, dom)
+    updateDom(n)
   }
   return
 }

packages/lib/src/globals.ts

@@ -1,14 +1,6 @@
 import type { AppContext } from "./appContext"
 
-export {
-  ctx,
-  node,
-  nodeToCtxMap,
-  contexts,
-  renderMode,
-  hydrationStack,
-  childIndexStack,
-}
+export { ctx, node, nodeToCtxMap, contexts, renderMode }
 
 const nodeToCtxMap = new WeakMap<Kaioken.VNode, AppContext>()
 const contexts: Array<AppContext<any>> = []
@@ -24,6 +16,3 @@ const ctx = {
 const renderMode = {
   current: "dom" as "dom" | "hydrate" | "string" | "stream",
 }
-
-const hydrationStack = [] as Array<HTMLElement | SVGElement | Text>
-const childIndexStack = [] as Array<number>

packages/lib/src/hydration.ts

@@ -0,0 +1,22 @@
+const parentStack = [] as Array<HTMLElement | SVGElement | Text>
+const childIdxStack = [] as Array<number>
+
+export const hydrationStack = {
+  clear: () => {
+    parentStack.length = 0
+    childIdxStack.length = 0
+  },
+  pop: () => {
+    parentStack.pop()
+    childIdxStack.pop()
+  },
+  push: (el: HTMLElement | SVGElement | Text) => {
+    parentStack.push(el)
+    childIdxStack.push(0)
+  },
+  nextChild: () => {
+    return parentStack[parentStack.length - 1].childNodes[
+      childIdxStack[childIdxStack.length - 1]++
+    ] as HTMLElement | SVGElement | Text | undefined
+  },
+}

packages/lib/src/scheduler.ts

@@ -1,14 +1,9 @@
 import type { AppContext } from "./appContext"
 import { Component } from "./component.js"
 import { EffectTag, elementTypes as et } from "./constants.js"
-import { commitWork, createDom, updateDom } from "./dom.js"
-import {
-  childIndexStack,
-  ctx,
-  hydrationStack,
-  node,
-  renderMode,
-} from "./globals.js"
+import { commitWork, createDom, hydrateDom, updateDom } from "./dom.js"
+import { ctx, node, renderMode } from "./globals.js"
+import { hydrationStack } from "./hydration.js"
 import { assertValidElementProps } from "./props.js"
 import { reconcileChildren } from "./reconciler.js"
 import { vNodeContains } from "./utils.js"
@@ -226,7 +221,6 @@ export class Scheduler {
       if (vNode.child) {
         if (renderMode.current === "hydrate" && vNode.dom) {
           hydrationStack.push(vNode.dom)
-          childIndexStack.push(0)
         }
         return vNode.child
       }
@@ -242,7 +236,6 @@ export class Scheduler {
       nextNode = nextNode.parent
       if (renderMode.current === "hydrate" && nextNode?.dom) {
         hydrationStack.pop()
-        childIndexStack.pop()
       }
     }
   }
@@ -285,20 +278,10 @@ export class Scheduler {
     node.current = vNode
     if (!vNode.dom) {
       if (renderMode.current === "hydrate") {
-        const dom = currentDom()!
-        if ((vNode.type as string) !== dom.nodeName.toLowerCase()) {
-          throw new Error(
-            `[kaioken]: Expected node of type ${vNode.type} but received ${dom.nodeName}`
-          )
-        }
-        vNode.dom = dom
-        if (vNode.type === et.text) {
-          handleTextNodeSplitting(vNode)
-        } else {
-          updateDom(vNode, vNode.dom)
-        }
+        hydrateDom(vNode)
       } else {
         vNode.dom = createDom(vNode)
+        updateDom(vNode)
       }
     }
     if (vNode.props.ref) {

packages/lib/src/ssr/client.ts

@@ -1,4 +1,5 @@
 import type { AppContext, AppContextOptions } from "../appContext"
+import { hydrationStack } from "../hydration.js"
 import { renderMode } from "../globals.js"
 import { mount } from "../index.js"
 
@@ -19,6 +20,7 @@ export function hydrate<T extends Record<string, unknown>>(
   optionsOrRoot: HTMLElement | AppContextOptions,
   appProps = {} as T
 ) {
+  hydrationStack.clear()
   const prevRenderMode = renderMode.current
   renderMode.current = "hydrate"
   return new Promise((resolve) => {