From e80609b636b08e54fea487dfa34b35374470723d Mon Sep 17 00:00:00 2001
From: Robby6Strings <rausten93@gmail.com>
Date: Fri, 4 Oct 2024 14:01:16 +1300
Subject: [PATCH] lib - refactor symbol names, implement GenericHMRAcceptor for
 stores

---
 packages/lib/src/constants.ts        | 11 ++++-----
 packages/lib/src/context.ts          |  4 ++--
 packages/lib/src/element.ts          |  4 ++--
 packages/lib/src/error.ts            |  7 +++---
 packages/lib/src/globalContext.ts    | 34 ++++++++++++++++++----------
 packages/lib/src/hmr.ts              | 21 +++++++++++++++++
 packages/lib/src/hooks/useContext.ts |  6 ++---
 packages/lib/src/index.ts            |  2 +-
 packages/lib/src/reconciler.ts       |  8 +++----
 packages/lib/src/renderToString.ts   |  8 ++-----
 packages/lib/src/signal.ts           |  6 ++---
 packages/lib/src/ssr/server.ts       |  8 ++-----
 packages/lib/src/store.ts            | 29 ++++++++++++++++++++++--
 packages/lib/src/types.ts            |  4 ++--
 packages/lib/src/types.utils.ts      |  4 +---
 packages/lib/src/utils.ts            | 16 +++++--------
 sandbox/csr/src/App.tsx              | 13 ++++++++++-
 sandbox/csr/src/countStore.ts        | 10 ++++++++
 18 files changed, 128 insertions(+), 67 deletions(-)
 create mode 100644 packages/lib/src/hmr.ts
 create mode 100644 sandbox/csr/src/countStore.ts

diff --git a/packages/lib/src/constants.ts b/packages/lib/src/constants.ts
index b8155ce4..ca6f1700 100644
--- a/packages/lib/src/constants.ts
+++ b/packages/lib/src/constants.ts
@@ -1,10 +1,9 @@
-export const signalSymbol = Symbol.for("kaioken.signal")
-export const componentSymbol = Symbol.for("kaioken.component")
-export const contextProviderSymbol = Symbol.for("kaioken.contextProvider")
-export const fragmentSymbol = Symbol.for("kaioken.fragment")
-export const kaiokenErrorSymbol = Symbol.for("kaioken.error")
+export const $SIGNAL = Symbol.for("kaioken.signal")
+export const $CONTEXT_PROVIDER = Symbol.for("kaioken.contextProvider")
+export const $FRAGMENT = Symbol.for("kaioken.fragment")
+export const $KAIOKEN_ERROR = Symbol.for("kaioken.error")
+export const $HMR_ACCEPTOR = Symbol.for("kaioken.hrmAcceptor")
 
-export const ELEMENT_ID_BASE = 16
 export const CONSECUTIVE_DIRTY_LIMIT = 50
 
 export const FLAG = {
diff --git a/packages/lib/src/context.ts b/packages/lib/src/context.ts
index c213d09e..86ba91ba 100644
--- a/packages/lib/src/context.ts
+++ b/packages/lib/src/context.ts
@@ -1,11 +1,11 @@
-import { contextProviderSymbol } from "./constants.js"
+import { $CONTEXT_PROVIDER } from "./constants.js"
 import { createElement } from "./element.js"
 
 export function createContext<T>(defaultValue: T): Kaioken.Context<T> {
   const ctx: Kaioken.Context<T> = {
     Provider: ({ value, children }: Kaioken.ProviderProps<T>) => {
       return createElement(
-        contextProviderSymbol,
+        $CONTEXT_PROVIDER,
         { value, ctx },
         typeof children === "function" ? children(value) : children
       )
diff --git a/packages/lib/src/element.ts b/packages/lib/src/element.ts
index 218c3db2..57268fe9 100644
--- a/packages/lib/src/element.ts
+++ b/packages/lib/src/element.ts
@@ -1,4 +1,4 @@
-import { fragmentSymbol } from "./constants.js"
+import { $FRAGMENT } from "./constants.js"
 import { isValidElementKeyProp, isValidElementRefProp } from "./props.js"
 
 export function createElement<T extends Kaioken.VNode["type"]>(
@@ -40,5 +40,5 @@ export function Fragment({
   children: JSX.Children
   key?: JSX.ElementKey
 }): Kaioken.VNode {
-  return createElement(fragmentSymbol, key ? { key } : null, children)
+  return createElement($FRAGMENT, key ? { key } : null, children)
 }
diff --git a/packages/lib/src/error.ts b/packages/lib/src/error.ts
index b308700a..81c0ea0c 100644
--- a/packages/lib/src/error.ts
+++ b/packages/lib/src/error.ts
@@ -1,4 +1,4 @@
-import { kaiokenErrorSymbol } from "./constants.js"
+import { $KAIOKEN_ERROR } from "./constants.js"
 import { __DEV__ } from "./env.js"
 import { findParent, noop } from "./utils.js"
 
@@ -13,7 +13,7 @@ type KaiokenErrorOptions =
     }
 
 export class KaiokenError extends Error {
-  [kaiokenErrorSymbol] = true
+  [$KAIOKEN_ERROR] = true
   /** Indicates whether the error is fatal and should crash the application */
   fatal?: boolean
   /** Present if vNode is provided */
@@ -34,8 +34,7 @@ export class KaiokenError extends Error {
 
   static isKaiokenError(error: unknown): error is KaiokenError {
     return (
-      error instanceof Error &&
-      (error as KaiokenError)[kaiokenErrorSymbol] === true
+      error instanceof Error && (error as KaiokenError)[$KAIOKEN_ERROR] === true
     )
   }
 }
diff --git a/packages/lib/src/globalContext.ts b/packages/lib/src/globalContext.ts
index 36de0495..7ab18994 100644
--- a/packages/lib/src/globalContext.ts
+++ b/packages/lib/src/globalContext.ts
@@ -1,5 +1,8 @@
 import type { AppContext } from "./appContext"
+import type { Store } from "./store"
+import { $HMR_ACCEPTOR } from "./constants.js"
 import { __DEV__ } from "./env.js"
+import { isGenericHmrAcceptor } from "./hmr.js"
 import { traverseApply } from "./utils.js"
 
 export { KaiokenGlobalContext, type GlobalKaiokenEvent }
@@ -24,6 +27,8 @@ type Evt =
 
 type GlobalKaiokenEvent = Evt["name"]
 
+type HotVar = Kaioken.FC | Store<any, any>
+
 class KaiokenGlobalContext {
   #contexts: Set<AppContext> = new Set()
   private listeners: Map<
@@ -37,24 +42,29 @@ class KaiokenGlobalContext {
   }
 
   HMRContext = {
-    register: (filePath: string, componentMap: Record<string, Kaioken.FC>) => {
+    register: (filePath: string, hotVars: Record<string, HotVar>) => {
       if (__DEV__) {
-        const components = this.moduleMap.get(filePath)
-        if (!components) {
-          this.moduleMap.set(filePath, new Map(Object.entries(componentMap)))
+        const mod = this.moduleMap.get(filePath)
+        if (!mod) {
+          this.moduleMap.set(filePath, new Map(Object.entries(hotVars)))
           return
         }
-        for (const [name, newFn] of Object.entries(componentMap)) {
-          const oldFn = components.get(name)
-          components.set(name, newFn)
-          if (!oldFn) continue
+        for (const [name, newVal] of Object.entries(hotVars)) {
+          const oldVal = mod.get(name)
+          mod.set(name, newVal)
+          if (!oldVal) continue
+          if (isGenericHmrAcceptor(oldVal) && isGenericHmrAcceptor(newVal)) {
+            newVal[$HMR_ACCEPTOR].inject(oldVal[$HMR_ACCEPTOR].provide())
+            oldVal[$HMR_ACCEPTOR].destroy()
+            continue
+          }
           this.#contexts.forEach((ctx) => {
             if (!ctx.mounted || !ctx.rootNode) return
             traverseApply(ctx.rootNode, (vNode) => {
-              if (vNode.type === oldFn) {
-                vNode.type = newFn
+              if (vNode.type === oldVal) {
+                vNode.type = newVal
                 if (vNode.prev) {
-                  vNode.prev.type = newFn
+                  vNode.prev.type = newVal
                 }
                 ctx.requestUpdate(vNode)
               }
@@ -90,5 +100,5 @@ class KaiokenGlobalContext {
     this.listeners.get(event)!.delete(callback)
   }
 
-  private moduleMap = new Map<string, Map<string, Kaioken.FC>>()
+  private moduleMap = new Map<string, Map<string, HotVar>>()
 }
diff --git a/packages/lib/src/hmr.ts b/packages/lib/src/hmr.ts
new file mode 100644
index 00000000..77b39659
--- /dev/null
+++ b/packages/lib/src/hmr.ts
@@ -0,0 +1,21 @@
+import { $HMR_ACCEPTOR } from "./constants.js"
+
+export type GenericHMRAcceptor<T = {}> = {
+  [$HMR_ACCEPTOR]: {
+    provide: () => T
+    inject: (prev: T) => void
+    destroy: () => void
+  }
+}
+
+export function isGenericHmrAcceptor(
+  thing: unknown
+): thing is GenericHMRAcceptor<any> {
+  return (
+    !!thing &&
+    (typeof thing === "object" || typeof thing === "function") &&
+    $HMR_ACCEPTOR in thing &&
+    typeof thing[$HMR_ACCEPTOR] === "object" &&
+    !!thing[$HMR_ACCEPTOR]
+  )
+}
diff --git a/packages/lib/src/hooks/useContext.ts b/packages/lib/src/hooks/useContext.ts
index a9865dc9..f1a34901 100644
--- a/packages/lib/src/hooks/useContext.ts
+++ b/packages/lib/src/hooks/useContext.ts
@@ -1,9 +1,9 @@
 import { type HookCallbackState, useHook } from "./utils.js"
 import { __DEV__ } from "../env.js"
-import { contextProviderSymbol } from "../constants.js"
+import { $CONTEXT_PROVIDER } from "../constants.js"
 
 type ContextProviderNode<T> = Kaioken.VNode & {
-  type: typeof contextProviderSymbol
+  type: typeof $CONTEXT_PROVIDER
   props: { value: T; ctx: Kaioken.Context<T> }
 }
 
@@ -43,7 +43,7 @@ const useContextCallback = <T>({
 
     let n = vNode.parent
     while (n) {
-      if (n.type === contextProviderSymbol) {
+      if (n.type === $CONTEXT_PROVIDER) {
         const ctxNode = n as ContextProviderNode<T>
         if (ctxNode.props.ctx === hook.context) {
           hook.ctxNode = ctxNode
diff --git a/packages/lib/src/index.ts b/packages/lib/src/index.ts
index 643b0349..52e31413 100644
--- a/packages/lib/src/index.ts
+++ b/packages/lib/src/index.ts
@@ -14,7 +14,7 @@ export * from "./portal.js"
 export * from "./renderToString.js"
 export * from "./router/index.js"
 export * from "./signal.js"
-export * from "./store.js"
+export { createStore, type Store, type MethodFactory } from "./store.js"
 export * from "./transition.js"
 
 if ("window" in globalThis) {
diff --git a/packages/lib/src/reconciler.ts b/packages/lib/src/reconciler.ts
index 2e4c70b3..5b333c24 100644
--- a/packages/lib/src/reconciler.ts
+++ b/packages/lib/src/reconciler.ts
@@ -1,5 +1,5 @@
 import type { AppContext } from "./appContext"
-import { ELEMENT_TYPE, FLAG, fragmentSymbol } from "./constants.js"
+import { ELEMENT_TYPE, FLAG, $FRAGMENT } from "./constants.js"
 import { ctx } from "./globals.js"
 import { isVNode } from "./utils.js"
 import { Signal } from "./signal.js"
@@ -211,7 +211,7 @@ function updateTextNode(parent: VNode, oldNode: VNode | null, content: string) {
 
 function updateNode(parent: VNode, oldNode: VNode | null, newNode: VNode) {
   const nodeType = newNode.type
-  if (nodeType === fragmentSymbol) {
+  if (nodeType === $FRAGMENT) {
     return updateFragment(
       parent,
       oldNode,
@@ -239,8 +239,8 @@ function updateFragment(
   children: unknown[],
   newProps = {}
 ) {
-  if (oldNode === null || oldNode.type !== fragmentSymbol) {
-    const el = createElement(fragmentSymbol, { children, ...newProps })
+  if (oldNode === null || oldNode.type !== $FRAGMENT) {
+    const el = createElement($FRAGMENT, { children, ...newProps })
     el.parent = parent
     el.depth = parent.depth + 1
     return el
diff --git a/packages/lib/src/renderToString.ts b/packages/lib/src/renderToString.ts
index 3beef425..bc62be69 100644
--- a/packages/lib/src/renderToString.ts
+++ b/packages/lib/src/renderToString.ts
@@ -8,11 +8,7 @@ import {
   selfClosingTags,
 } from "./utils.js"
 import { Signal } from "./signal.js"
-import {
-  contextProviderSymbol,
-  ELEMENT_TYPE,
-  fragmentSymbol,
-} from "./constants.js"
+import { $CONTEXT_PROVIDER, ELEMENT_TYPE, $FRAGMENT } from "./constants.js"
 import { assertValidElementProps } from "./props.js"
 
 export function renderToString<T extends Record<string, unknown>>(
@@ -56,7 +52,7 @@ function renderToString_internal(
   const type = el.type
   if (type === ELEMENT_TYPE.text)
     return encodeHtmlEntities(props.nodeValue ?? "")
-  if (type === fragmentSymbol || type === contextProviderSymbol) {
+  if (type === $FRAGMENT || type === $CONTEXT_PROVIDER) {
     if (!Array.isArray(children))
       return renderToString_internal(children, el, idx)
     return children.map((c, i) => renderToString_internal(c, el, i)).join("")
diff --git a/packages/lib/src/signal.ts b/packages/lib/src/signal.ts
index 99bc5569..6b9f8cb4 100644
--- a/packages/lib/src/signal.ts
+++ b/packages/lib/src/signal.ts
@@ -1,4 +1,4 @@
-import { signalSymbol } from "./constants.js"
+import { $SIGNAL } from "./constants.js"
 import { __DEV__ } from "./env.js"
 import { node } from "./globals.js"
 import { useHook } from "./hooks/utils.js"
@@ -91,7 +91,7 @@ export interface SignalLike<T> {
 type SignalSubscriber = Kaioken.VNode | Function
 
 export class Signal<T> {
-  [signalSymbol] = true
+  [$SIGNAL] = true
   #value: T
   #subscribers = new Set<SignalSubscriber>()
   displayName?: string
@@ -148,7 +148,7 @@ export class Signal<T> {
   }
 
   static isSignal(x: any): x is Signal<any> {
-    return typeof x === "object" && !!x && signalSymbol in x
+    return typeof x === "object" && !!x && $SIGNAL in x
   }
 
   static unsubscribe(sub: SignalSubscriber, signal: Signal<any>) {
diff --git a/packages/lib/src/ssr/server.ts b/packages/lib/src/ssr/server.ts
index cb41f137..fccf8e4f 100644
--- a/packages/lib/src/ssr/server.ts
+++ b/packages/lib/src/ssr/server.ts
@@ -9,11 +9,7 @@ import {
   selfClosingTags,
 } from "../utils.js"
 import { Signal } from "../signal.js"
-import {
-  contextProviderSymbol,
-  ELEMENT_TYPE,
-  fragmentSymbol,
-} from "../constants.js"
+import { $CONTEXT_PROVIDER, ELEMENT_TYPE, $FRAGMENT } from "../constants.js"
 import { assertValidElementProps } from "../props.js"
 
 type RequestState = {
@@ -84,7 +80,7 @@ function renderToStream_internal(
     state.stream.push(encodeHtmlEntities(props.nodeValue ?? ""))
     return
   }
-  if (type === fragmentSymbol || type === contextProviderSymbol) {
+  if (type === $FRAGMENT || type === $CONTEXT_PROVIDER) {
     if (!Array.isArray(children))
       return renderToStream_internal(state, children, el, idx)
     return children.forEach((c, i) => renderToStream_internal(state, c, el, i))
diff --git a/packages/lib/src/store.ts b/packages/lib/src/store.ts
index 3eb949ce..6e09d9be 100644
--- a/packages/lib/src/store.ts
+++ b/packages/lib/src/store.ts
@@ -2,6 +2,8 @@ import type { Prettify } from "./types.utils.js"
 import { __DEV__ } from "./env.js"
 import { sideEffectsEnabled, useAppContext, useHook } from "./hooks/utils.js"
 import { getVNodeAppContext, shallowCompare } from "./utils.js"
+import { $HMR_ACCEPTOR } from "./constants.js"
+import { GenericHMRAcceptor } from "./hmr.js"
 
 export { createStore }
 export type { Store, MethodFactory }
@@ -45,7 +47,7 @@ function createStore<T, U extends MethodFactory<T>>(
   let state = initial
   let stateIteration = 0
   const subscribers = new Set<Kaioken.VNode | Function>()
-  const nodeToSliceComputeMap = new WeakMap<Kaioken.VNode, NodeSliceCompute[]>()
+  let nodeToSliceComputeMap = new WeakMap<Kaioken.VNode, NodeSliceCompute[]>()
 
   const getState = () => state
   const setState = (setter: Kaioken.StateSetter<T>) => {
@@ -126,7 +128,7 @@ function createStore<T, U extends MethodFactory<T>>(
     )
   }
 
-  return Object.assign(useStore, {
+  const store = Object.assign(useStore, {
     getState,
     setState,
     methods,
@@ -135,4 +137,27 @@ function createStore<T, U extends MethodFactory<T>>(
       return (() => (subscribers.delete(fn), void 0)) as () => void
     },
   })
+
+  if (__DEV__) {
+    return Object.assign(store, {
+      [$HMR_ACCEPTOR]: {
+        provide: () => ({ state, subscribers, nodeToSliceComputeMap }),
+        inject: (prev) => {
+          prev.subscribers.forEach((sub) => subscribers.add(sub))
+          nodeToSliceComputeMap = prev.nodeToSliceComputeMap
+          setState(prev.state)
+        },
+        destroy: () => {
+          subscribers.clear()
+          nodeToSliceComputeMap = new WeakMap()
+        },
+      },
+    } satisfies GenericHMRAcceptor<{
+      state: T
+      subscribers: Set<Kaioken.VNode | Function>
+      nodeToSliceComputeMap: WeakMap<Kaioken.VNode, NodeSliceCompute[]>
+    }>)
+  }
+
+  return store
 }
diff --git a/packages/lib/src/types.ts b/packages/lib/src/types.ts
index d3bdc570..df4d6e49 100644
--- a/packages/lib/src/types.ts
+++ b/packages/lib/src/types.ts
@@ -3,7 +3,7 @@ import type {
   Signal as SignalClass,
   SignalLike,
 } from "./signal"
-import type { contextProviderSymbol, fragmentSymbol } from "./constants"
+import type { $CONTEXT_PROVIDER, $FRAGMENT } from "./constants"
 import type { KaiokenGlobalContext } from "./globalContext"
 import type {
   EventAttributes,
@@ -171,7 +171,7 @@ declare global {
 
     type Signal<T> = SignalClass<T> | ReadonlySignal<T>
 
-    type ExoticSymbol = typeof fragmentSymbol | typeof contextProviderSymbol
+    type ExoticSymbol = typeof $FRAGMENT | typeof $CONTEXT_PROVIDER
 
     type VNode = {
       type: string | Function | ExoticSymbol
diff --git a/packages/lib/src/types.utils.ts b/packages/lib/src/types.utils.ts
index 9ce55229..09e8e625 100644
--- a/packages/lib/src/types.utils.ts
+++ b/packages/lib/src/types.utils.ts
@@ -1,5 +1,3 @@
-import { contextProviderSymbol, fragmentSymbol } from "./constants"
-
 export type SomeElement = HTMLElement | SVGElement
 export type SomeDom = HTMLElement | SVGElement | Text
 export type MaybeDom = SomeDom | undefined
@@ -8,7 +6,7 @@ type VNode = Kaioken.VNode
 
 export type FunctionVNode = VNode & { type: (...args: any) => any }
 export type ExoticVNode = VNode & {
-  type: typeof contextProviderSymbol | typeof fragmentSymbol
+  type: Kaioken.ExoticSymbol
 }
 export type ElementVNode = VNode & { dom: SomeElement }
 export type DomVNode = VNode & { dom: SomeDom }
diff --git a/packages/lib/src/utils.ts b/packages/lib/src/utils.ts
index bc12297c..7e134995 100644
--- a/packages/lib/src/utils.ts
+++ b/packages/lib/src/utils.ts
@@ -1,9 +1,5 @@
 import { node, nodeToCtxMap, renderMode } from "./globals.js"
-import {
-  contextProviderSymbol,
-  fragmentSymbol,
-  REGEX_UNIT,
-} from "./constants.js"
+import { $CONTEXT_PROVIDER, $FRAGMENT, REGEX_UNIT } from "./constants.js"
 import { unwrap } from "./signal.js"
 import { KaiokenError } from "./error.js"
 import type { AppContext } from "./appContext"
@@ -50,20 +46,20 @@ function isVNode(thing: unknown): thing is VNode {
 function isExoticVNode(thing: unknown): thing is ExoticVNode {
   return (
     isVNode(thing) &&
-    (thing.type === fragmentSymbol || thing.type === contextProviderSymbol)
+    (thing.type === $FRAGMENT || thing.type === $CONTEXT_PROVIDER)
   )
 }
 
 function isFragment(
   thing: unknown
-): thing is VNode & { type: typeof fragmentSymbol } {
-  return isVNode(thing) && thing.type === fragmentSymbol
+): thing is VNode & { type: typeof $FRAGMENT } {
+  return isVNode(thing) && thing.type === $FRAGMENT
 }
 
 function isContextProvider(
   thing: unknown
-): thing is VNode & { type: typeof contextProviderSymbol } {
-  return isVNode(thing) && thing.type === contextProviderSymbol
+): thing is VNode & { type: typeof $CONTEXT_PROVIDER } {
+  return isVNode(thing) && thing.type === $CONTEXT_PROVIDER
 }
 
 function getCurrentVNode(): VNode | undefined {
diff --git a/sandbox/csr/src/App.tsx b/sandbox/csr/src/App.tsx
index d87d4c27..3056ee1e 100644
--- a/sandbox/csr/src/App.tsx
+++ b/sandbox/csr/src/App.tsx
@@ -1,4 +1,5 @@
 import { Router, Route, Link, lazy } from "kaioken"
+import { countStore } from "./countStore"
 
 type AppRoute = {
   title: string
@@ -7,7 +8,17 @@ type AppRoute = {
 }
 
 const Home: Kaioken.FC = () => {
-  return <h1 style={{ color: 123 }}>Home</h1>
+  const { value: count, increment, decrement, double, triple } = countStore()
+  return (
+    <div>
+      <h1 style={{ color: "red" }}>Home</h1>
+      <p>Count: {count}</p>
+      <p>Double: {double()}</p>
+      <p>Triple: {triple()}</p>
+      <button onclick={() => increment()}>Increment</button>
+      <button onclick={() => decrement()}>Decrement</button>
+    </div>
+  )
 }
 
 const ROUTES: Record<string, AppRoute> = {
diff --git a/sandbox/csr/src/countStore.ts b/sandbox/csr/src/countStore.ts
new file mode 100644
index 00000000..e525c139
--- /dev/null
+++ b/sandbox/csr/src/countStore.ts
@@ -0,0 +1,10 @@
+import { createStore } from "kaioken"
+
+export const countStore = createStore(0, (set, get) => ({
+  increment: () => {
+    set((count) => count + 1)
+  },
+  decrement: () => set((count) => count - 1),
+  double: () => get() * 2,
+  triple: () => get() * 3,
+}))