From ca728a54bce180e7b8243c355634b5cefbd56539 Mon Sep 17 00:00:00 2001
From: Mikhail Fedotov <nosik90@gmail.com>
Date: Wed, 6 Dec 2023 13:51:40 +0300
Subject: [PATCH] Add choiceDataState pseudo state

---
 .../ru/nsk/kstatemachine/DefaultState.kt      | 22 +++++++++---
 .../kotlin/ru/nsk/kstatemachine/IState.kt     |  5 +++
 .../visitors/ExportPlantUmlVisitor.kt         |  6 ++--
 .../ru/nsk/kstatemachine/ChoiceStateTest.kt   | 35 +++++++++++++++++++
 4 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/DefaultState.kt b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/DefaultState.kt
index 3c2760ea..b34febe3 100644
--- a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/DefaultState.kt
+++ b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/DefaultState.kt
@@ -10,14 +10,26 @@ open class DefaultState(name: String? = null, childMode: ChildMode = EXCLUSIVE)
 
 open class DefaultFinalState(name: String? = null) : DefaultState(name), FinalState
 
-/**
- * Currently it does not allow to target [DataState]
- */
-open class DefaultChoiceState(name: String? = null, private val choiceAction: suspend EventAndArgument<*>.() -> State) :
-    BasePseudoState(name), RedirectPseudoState {
+open class DefaultChoiceState(
+    name: String? = null,
+    private val choiceAction: suspend EventAndArgument<*>.() -> State
+) : BasePseudoState(name), RedirectPseudoState {
+
+    override suspend fun resolveTargetState(eventAndArgument: EventAndArgument<*>) =
+        eventAndArgument.choiceAction().also { log { "$this resolved to $it" } }
+}
+
+open class DefaultChoiceDataState<D : Any>(
+    name: String? = null,
+    private val choiceAction: suspend EventAndArgument<*>.() -> DataState<D>,
+) : DataState<D>, BasePseudoState(name), RedirectPseudoState {
 
     override suspend fun resolveTargetState(eventAndArgument: EventAndArgument<*>) =
         eventAndArgument.choiceAction().also { log { "$this resolved to $it" } }
+
+    override val defaultData: D? = null
+    override val data: D get() = error("PseudoState $this can not have data")
+    override val lastData: D get() = error("PseudoState $this can not have lastData")
 }
 
 open class BasePseudoState(name: String?) : BaseStateImpl(name, EXCLUSIVE), PseudoState {
diff --git a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/IState.kt b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/IState.kt
index aa6310a6..eb8c14e7 100644
--- a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/IState.kt
+++ b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/IState.kt
@@ -260,6 +260,11 @@ inline fun <reified D : Any> IState.finalDataState(
 fun IState.choiceState(name: String? = null, choiceAction: suspend EventAndArgument<*>.() -> State) =
     addState(DefaultChoiceState(name, choiceAction))
 
+fun <D : Any> IState.choiceDataState(
+    name: String? = null,
+    choiceAction: suspend EventAndArgument<*>.() -> DataState<D>
+) = addState(DefaultChoiceDataState(name, choiceAction))
+
 fun IState.historyState(
     name: String? = null,
     defaultState: IState? = null,
diff --git a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/visitors/ExportPlantUmlVisitor.kt b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/visitors/ExportPlantUmlVisitor.kt
index 98559a96..f7e73950 100644
--- a/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/visitors/ExportPlantUmlVisitor.kt
+++ b/kstatemachine/src/commonMain/kotlin/ru/nsk/kstatemachine/visitors/ExportPlantUmlVisitor.kt
@@ -103,8 +103,10 @@ internal class ExportPlantUmlVisitor(private val showEventLabels: Boolean) : CoV
     private fun line(text: String) = builder.appendLine(SINGLE_INDENT.repeat(indent) + text)
 
     private fun transitionLabel(transition: Transition<*>): String {
-        val eventName = if (showEventLabels) transition.eventMatcher.eventClass.simpleName else null
-        val entries = listOfNotNull(transition.name, eventName)
+        val entries = listOfNotNull(
+            transition.name,
+            transition.eventMatcher.eventClass.simpleName.takeIf { showEventLabels },
+        )
         return label(entries.joinToString())
     }
 
diff --git a/tests/src/commonTest/kotlin/ru/nsk/kstatemachine/ChoiceStateTest.kt b/tests/src/commonTest/kotlin/ru/nsk/kstatemachine/ChoiceStateTest.kt
index 7fc56ad9..7cb1e1a7 100644
--- a/tests/src/commonTest/kotlin/ru/nsk/kstatemachine/ChoiceStateTest.kt
+++ b/tests/src/commonTest/kotlin/ru/nsk/kstatemachine/ChoiceStateTest.kt
@@ -68,5 +68,40 @@ class ChoiceStateTest : StringSpec({
                 }
             }
         }
+
+        "redirecting choice data state" {
+            val callbacks = mockkCallbacks()
+
+            class IntEvent(override val data: Int) : DataEvent<Int>
+
+            lateinit var intState1: DataState<Int>
+            lateinit var intState2: DataState<Int>
+
+            val machine = createTestStateMachine(coroutineStarterType) {
+                logger = StateMachine.Logger { println(it()) }
+
+                addInitialState(State1)
+
+                val choice = choiceDataState("data choice") {
+                    log { "$event $argument" }
+                    val intEvent = event as? IntEvent // cast is necessary as we don't know event type here
+                    if (intEvent?.data == 42) intState1 else intState2
+                }
+
+                dataTransition<IntEvent, Int> { targetState = choice }
+
+                intState1 = dataState<Int>("intState1") { callbacks.listen(this) }
+                intState2 = dataState<Int>("intState2") { callbacks.listen(this) }
+                onTransitionTriggered { log { it.toString() } }
+            }
+
+            machine.processEvent(IntEvent(42), true)
+            verifySequenceAndClear(callbacks) { callbacks.onStateEntry(intState1) }
+            machine.processEvent(IntEvent(66), false)
+            verifySequenceAndClear(callbacks) {
+                callbacks.onStateExit(intState1)
+                callbacks.onStateEntry(intState2)
+            }
+        }
     }
 })
\ No newline at end of file