From 8450be57e6cf7144432992ba83736ae22a6642c0 Mon Sep 17 00:00:00 2001
From: Sergey Malenkov <sergey.malenkov@huawei.com>
Date: Mon, 4 Aug 2025 14:14:07 +0300
Subject: [PATCH] Use WorkerLocal for ArkTS and add compatibility for TS

---
 incremental/common/src/index.ts            |  2 +-
 incremental/compat/src/arkts/utils.ts      | 14 +------
 incremental/compat/src/index.ts            |  2 +-
 incremental/compat/src/ohos/index.ts       |  2 +-
 incremental/compat/src/typescript/utils.ts | 48 +++++++++++++++++++---
 incremental/runtime/src/states/State.ts    |  4 +-
 6 files changed, 48 insertions(+), 24 deletions(-)

diff --git a/incremental/common/src/index.ts b/incremental/common/src/index.ts
index 45bd3d454c..fa91fbf70c 100644
--- a/incremental/common/src/index.ts
+++ b/incremental/common/src/index.ts
@@ -42,7 +42,7 @@ export {
     int8Array,
     errorAsString,
     unsafeCast,
-    CoroutineLocalValue,
+    WorkerLocalValue,
     scheduleCoroutine,
     memoryStats,
     launchJob
diff --git a/incremental/compat/src/arkts/utils.ts b/incremental/compat/src/arkts/utils.ts
index ec855ec898..d21116f6ce 100644
--- a/incremental/compat/src/arkts/utils.ts
+++ b/incremental/compat/src/arkts/utils.ts
@@ -39,16 +39,4 @@ export function launchJob(task: () => void): Promise<Any> {
     return taskpool.execute(task)
 }
 
-export class CoroutineLocalValue<T> {
-    private map = new containers.ConcurrentHashMap<int, T>
-    get(): T | undefined {
-        return this.map.get(CoroutineExtras.getWorkerId())
-    }
-    set(value: T | undefined) {
-        if (value) {
-            this.map.set(CoroutineExtras.getWorkerId(), value)
-        } else {
-            this.map.delete(CoroutineExtras.getWorkerId())
-        }
-    }
-}
+export type WorkerLocalValue<T> = WorkerLocal<T>
diff --git a/incremental/compat/src/index.ts b/incremental/compat/src/index.ts
index d068f170af..eb7a6130a7 100644
--- a/incremental/compat/src/index.ts
+++ b/incremental/compat/src/index.ts
@@ -55,7 +55,7 @@ export {
     int8Array,
     errorAsString,
     unsafeCast,
-    CoroutineLocalValue,
+    WorkerLocalValue,
     scheduleCoroutine,
     memoryStats,
     launchJob
diff --git a/incremental/compat/src/ohos/index.ts b/incremental/compat/src/ohos/index.ts
index 892265e669..4d454e8bdd 100644
--- a/incremental/compat/src/ohos/index.ts
+++ b/incremental/compat/src/ohos/index.ts
@@ -53,7 +53,7 @@ export {
     int8Array,
     errorAsString,
     unsafeCast,
-    CoroutineLocalValue,
+    WorkerLocalValue,
     scheduleCoroutine,
     memoryStats,
     launchJob
diff --git a/incremental/compat/src/typescript/utils.ts b/incremental/compat/src/typescript/utils.ts
index 1169707fa5..b6c0cfcfe8 100644
--- a/incremental/compat/src/typescript/utils.ts
+++ b/incremental/compat/src/typescript/utils.ts
@@ -13,6 +13,8 @@
  * limitations under the License.
  */
 
+import { AtomicRef } from "./atomic"
+
 export function errorAsString(error: any): string {
     if (error instanceof Error) {
         return error.stack ?? error.toString()
@@ -41,12 +43,46 @@ export function launchJob(task: () => void): Promise<any> {
     })
 }
 
-export class CoroutineLocalValue<T> {
-    private value: T | undefined = undefined
-    get(): T | undefined {
-        return this.value
+export class WorkerLocalValue<T> {
+    private ref?: AtomicRef<T>
+
+    /**
+     * @param init - a factory function to provide initial worker-local value if needed
+     */
+    constructor(private init?: () => T) {
+    }
+
+    /**
+     * @returns the worker-local value for current worker
+     * @throws `Error` when value not initialized and no `init` function provided
+     */
+    get(): T {
+        const ref = this.ref
+        if (ref) return ref.value
+        const init = this.init
+        if (!init) throw new Error("WorkerLocalValue not initialized: call set() first or provide init() function.")
+        const value = init()
+        this.ref = new AtomicRef<T>(value)
+        return value
     }
-    set(value: T | undefined) {
-        this.value = value
+
+    /**
+     * Updates the worker-local value for current worker.
+     * @param value - new value to store
+     */
+    set(value: T) {
+        const ref = this.ref
+        if (ref) {
+            ref.value = value
+        } else {
+            this.ref = new AtomicRef<T>(value)
+        }
+    }
+
+    /**
+     * Deletes the worker-local value for current worker.
+     */
+    delete() {
+        this.ref = undefined
     }
 }
diff --git a/incremental/runtime/src/states/State.ts b/incremental/runtime/src/states/State.ts
index e094481714..e01615d290 100644
--- a/incremental/runtime/src/states/State.ts
+++ b/incremental/runtime/src/states/State.ts
@@ -16,7 +16,7 @@
 import {
     Array_from_set,
     className,
-    CoroutineLocalValue,
+    WorkerLocalValue,
     float64toInt32,
     int32,
     KoalaCallsiteKey,
@@ -54,7 +54,7 @@ export function createStateManager(): StateManager {
     return new StateManagerImpl()
 }
 
-export const StateManagerLocal = new CoroutineLocalValue<StateManager>()
+export const StateManagerLocal = new WorkerLocalValue<StateManager | undefined>(() => undefined)
 
 /**
  * State manager, core of incremental runtime engine.
-- 
Gitee