2 files changed, 101 insertions, 0 deletions

diff --git a/state/memory_backend.go b/state/memory_backend.go
new file mode 100644
index 0000000..6594657
--- /dev/null
+++ b/state/memory_backend.go
@@ -0,0 +1,57 @@
+package state
+
+import (
+        "context"
+        "flux/core"
+)
+
+// TODO: currently using Golang Map for go through the DAG pipline
+// This should be use index map index map[uint64]uint32  Hash(KeyGroup + StateName + Key) -> Offset
+// with 1GB of []byte allocation
+// This should be optimize with ZERO GC Scanning, Memory Alignment
+
+type MemoryStateBackend struct {
+        maxParallelism  int
+        currentKey      string
+        currentGroupKey int
+        currentNamespace any
+        // 4D structure StateName -> KeyGroup -> Key -> Namespace -> Value
+        states map[string]map[int]map[string]map[any]any
+}
+
+func NewMemoryStateBackend(maxParallelism int) *MemoryStateBackend {
+        return &MemoryStateBackend{
+                maxParallelism: maxParallelism,
+                states:         make(map[string]map[int]map[string]map[any]any),
+        }
+}
+
+func (ms *MemoryStateBackend) GetCurrentKey() string {
+        return ms.currentKey
+}
+
+func (ms *MemoryStateBackend) SetCurrentKey(key string) error {
+        ms.currentKey = key
+        ms.currentGroupKey = core.AssignToKeyGroup(key, ms.maxParallelism)
+        return nil
+}
+
+func (ms *MemoryStateBackend) GetValueState(ctx context.Context, key string) (core.ValueState[any], error) {
+        // Assign if state not existed
+
+        if _, ok := ms.states[key]; !ok {
+                ms.states[key] = make(map[int]map[string]map[any]any)
+        }
+
+        return &memoryStateValue{backend: ms, stateName: key}, nil
+}
+
+func (ms *MemoryStateBackend) GetCurrentNamespace() any{
+
+        return ms.currentNamespace
+}
+
+func (ms *MemoryStateBackend) SetCurrentNamespace(namespace any) error {
+        ms.currentNamespace = namespace
+        return nil
+}
diff --git a/state/memory_state.go b/state/memory_state.go
new file mode 100644
index 0000000..95c7318
--- /dev/null
+++ b/state/memory_state.go
@@ -0,0 +1,44 @@
+package state
+
+import (
+"context"
+)
+
+type memoryStateValue struct {
+        backend   *MemoryStateBackend
+        stateName string
+}
+
+func (m *memoryStateValue) Value(ctx context.Context) (any, error) {
+        kg := m.backend.currentGroupKey
+        key := m.backend.currentKey
+        ns := m.backend.currentNamespace
+        if kgMap, exist := m.backend.states[m.stateName][kg][key]; exist {
+                if v, hasKey := kgMap[ns]; hasKey {
+                        return v, nil
+                }
+        }
+
+        return nil, nil
+}
+
+func (m *memoryStateValue) Update(ctx context.Context, newVal any) (any, error) {
+        kg := m.backend.currentGroupKey
+        key := m.backend.currentKey
+        ns := m.backend.currentNamespace
+        if _, exists := m.backend.states[m.stateName][kg][key]; !exists {
+                m.backend.states[m.stateName][kg] = make(map[string]map[any]any)
+        }
+
+        m.backend.states[m.stateName][kg][key][ns] = newVal
+        return newVal, nil
+}
+
+func (m *memoryStateValue) Clear(ctx context.Context) error {
+        kg := m.backend.currentGroupKey
+        key := m.backend.currentKey
+        if kgMap, exists := m.backend.states[m.stateName][kg]; exists {
+                delete(kgMap, key)
+        }
+        return nil
+}