From a2627baffc6acc926012fbddf1a9795e649d6084 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Thu, 9 Feb 2023 10:58:24 +0100
Subject: [PATCH] [statemachine] remove some type casts

By creating a statemachine with a template we can
create a functions of the correct type, removing
the need to cast inside the functions.
---
 codex/utils/asyncstatemachine.nim           | 42 +++++++++++----------
 tests/codex/utils/testasyncstatemachine.nim | 41 ++++++++++----------
 2 files changed, 43 insertions(+), 40 deletions(-)

diff --git a/codex/utils/asyncstatemachine.nim b/codex/utils/asyncstatemachine.nim
index a9fd8a2d..b60c5068 100644
--- a/codex/utils/asyncstatemachine.nim
+++ b/codex/utils/asyncstatemachine.nim
@@ -1,27 +1,29 @@
 import pkg/questionable
 import pkg/chronos
 
-type
-  AsyncStateMachine* = ref object of RootObj
-    state: AsyncState
-    running: Future[?AsyncState]
-  AsyncState* = ref object of RootObj
-  Event* = proc(state: AsyncState): ?AsyncState
+template makeStateMachine*(MachineType, StateType) =
 
-method run*(state: AsyncState): Future[?AsyncState] {.base.} = 
-  discard
+  type
+    MachineType* = ref object of RootObj
+      state: StateType
+      running: Future[?StateType]
+    StateType* = ref object of RootObj
+    Event* = proc(state: StateType): ?StateType
 
-proc runState(machine: AsyncStateMachine, state: AsyncState): Future[void] {.async.} =
-  if not machine.running.isNil:
-    await machine.running.cancelAndWait()
-  machine.state = state
-  machine.running = state.run()
-  if next =? await machine.running:
-    await machine.runState(next)
+  method run*(state: StateType): Future[?StateType] {.base.} =
+    discard
 
-proc start*(stateMachine: AsyncStateMachine, initialState: AsyncState) =
-  asyncSpawn stateMachine.runState(initialState)
+  proc runState*(machine: MachineType, state: StateType) {.async.} =
+    if not machine.running.isNil:
+      await machine.running.cancelAndWait()
+    machine.state = state
+    machine.running = state.run()
+    if next =? await machine.running:
+      await machine.runState(next)
 
-proc schedule*(stateMachine: AsyncStateMachine, event: Event) =
-  if next =? stateMachine.state.event():
-    asyncSpawn stateMachine.runState(next)
\ No newline at end of file
+  proc start*(stateMachine: MachineType, initialState: StateType) =
+    asyncSpawn runState(stateMachine, initialState)
+
+  proc schedule*(stateMachine: MachineType, event: Event) =
+    if next =? event(stateMachine.state):
+      asyncSpawn runState(stateMachine, next)
diff --git a/tests/codex/utils/testasyncstatemachine.nim b/tests/codex/utils/testasyncstatemachine.nim
index f16ef26f..b3727a32 100644
--- a/tests/codex/utils/testasyncstatemachine.nim
+++ b/tests/codex/utils/testasyncstatemachine.nim
@@ -4,42 +4,43 @@ import pkg/chronos
 import codex/utils/asyncstatemachine
 import ../helpers/eventually
 
+makeStateMachine(Machine, State)
+
 type
-  TestState = ref object of AsyncState
-  State1 = ref object of TestState
-  State2 = ref object of TestState
-  State3 = ref object of TestState
+  State1 = ref object of State
+  State2 = ref object of State
+  State3 = ref object of State
 
 var runs, cancellations = [0, 0, 0]
 
-method onMoveToNextStateEvent*(state: TestState): ?AsyncState {.base.} =
+method onMoveToNextStateEvent*(state: State): ?State {.base.} =
   discard
 
-method run(state: State1): Future[?AsyncState] {.async.} =
+method run(state: State1): Future[?State] {.async.} =
   inc runs[0]
-  return some AsyncState(State2.new())
+  return some State(State2.new())
 
-method run(state: State2): Future[?AsyncState] {.async.} =
+method run(state: State2): Future[?State] {.async.} =
   inc runs[1]
   try:
     await sleepAsync(1.hours)
   except CancelledError:
     inc cancellations[1]
 
-method onMoveToNextStateEvent(state: State2): ?AsyncState =
-  return some AsyncState(State3.new())
+method onMoveToNextStateEvent(state: State2): ?State =
+  return some State(State3.new())
 
-method run(state: State3): Future[?AsyncState] {.async.} =
+method run(state: State3): Future[?State] {.async.} =
   inc runs[2]
 
 suite "async state machines":
-  var machine: AsyncStateMachine
-  var state1, state2: AsyncState
+  var machine: Machine
+  var state1, state2: State
 
   setup:
     runs = [0, 0, 0]
     cancellations = [0, 0, 0]
-    machine = AsyncStateMachine.new()
+    machine = Machine.new()
     state1 = State1.new()
     state2 = State2.new()
 
@@ -54,19 +55,19 @@ suite "async state machines":
   test "state2 moves to state3 on event":
     machine.start(state2)
 
-    proc moveToNextStateEvent(state: AsyncState): ?AsyncState =
-      TestState(state).onMoveToNextStateEvent()
+    proc moveToNextStateEvent(state: State): ?State =
+      state.onMoveToNextStateEvent()
 
-    machine.schedule(Event(moveToNextStateEvent))
+    machine.schedule(moveToNextStateEvent)
 
     check eventually runs == [0, 1, 1]
 
   test "state transition will cancel the running state":
     machine.start(state2)
 
-    proc moveToNextStateEvent(state: AsyncState): ?AsyncState =
-      TestState(state).onMoveToNextStateEvent()
+    proc moveToNextStateEvent(state: State): ?State =
+      state.onMoveToNextStateEvent()
 
-    machine.schedule(Event(moveToNextStateEvent))
+    machine.schedule(moveToNextStateEvent)
 
     check eventually cancellations == [0, 1, 0]