From 2db9e7b9e061ef90d788df940ba1fcbfe23a540c Mon Sep 17 00:00:00 2001
From: Tianer Zhou <ezhoureal@gmail.com>
Date: Fri, 16 May 2025 16:15:28 +0800
Subject: [PATCH] optimize state update

Signed-off-by: Tianer Zhou <ezhoureal@gmail.com>
---
 incremental/runtime/src/states/Journal.ts | 10 +++++++++
 incremental/runtime/src/states/State.ts   | 27 +++++++++++++++++------
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/incremental/runtime/src/states/Journal.ts b/incremental/runtime/src/states/Journal.ts
index dc1186694..90d9ccbfa 100644
--- a/incremental/runtime/src/states/Journal.ts
+++ b/incremental/runtime/src/states/Journal.ts
@@ -24,6 +24,8 @@ export interface Changes {
      */
     getChange<Value>(state: Object): AtomicRef<Value> | undefined
 
+    entries(): MapIterator<[Object, Object]> | undefined
+
     /**
      * Removes all changes.
      */
@@ -61,6 +63,10 @@ export class Journal implements Changes {
         return reference.value ? new ChunkChanges(reference) : undefined
     }
 
+    entries() : MapIterator<[Object, Object]> | undefined {
+        return this.current.value?.map.entries()
+    }
+
     /**
      * Finds the latest value change for the given state.
      * @param state - the state to search
@@ -107,6 +113,10 @@ class ChunkChanges implements Changes {
         return find<Value>(state, this.reference.value)
     }
 
+    entries() : MapIterator<[Object, Object]> | undefined {
+        return this.reference.value?.map.entries()
+    }
+
     clear(): void {
         this.reference.value = undefined
     }
diff --git a/incremental/runtime/src/states/State.ts b/incremental/runtime/src/states/State.ts
index aab34f593..113d0ef65 100644
--- a/incremental/runtime/src/states/State.ts
+++ b/incremental/runtime/src/states/State.ts
@@ -567,15 +567,28 @@ class StateManagerImpl implements StateManager {
         // try to update snapshot for every state, except for parameter states
         const changes = this.journal.getChanges()
         const created = this.statesCreated.size as int32 // amount of created states to update
-        if (created > 0) {
-            const it = this.statesCreated.keys()
-            while (true) {
-                const result = it.next()
-                if (result.done) break
-                result.value?.updateStateSnapshot(changes)
-                if (result.value?.modified == true) modified++
+        const entries = changes?.entries()
+        if (entries) {
+            for (const [key, value] of entries!) {
+                if (!(key instanceof StateImpl)) {
+                    continue
+                }
+                let state = key as ManagedState
+                if (this.statesCreated.has(state)) {
+                    state.updateStateSnapshot(changes)
+                    if (state.modified) modified++
+                }
             }
         }
+        // if (created > 0) {
+        //     const it = this.statesCreated.keys()
+        //     while (true) {
+        //         const result = it.next()
+        //         if (result.done) break
+        //         result.value?.updateStateSnapshot(changes)
+        //         if (result.value?.modified == true) modified++
+        //     }
+        // }
         changes?.clear()
         KoalaProfiler.counters?.updateSnapshot(modified, created)
         // recompute dirty scopes only
-- 
Gitee