Adds a type argument for State

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## [Unreleased]
 
+### Breaking
+
+* Added type argument to State to allow for the ability to add customer behaviour to your states. This is a breaking
+  change as it will require you to update your state classes to include the type argument.
+
 ## [0.6.0]
 
 ### Breaking

README.md

@@ -27,7 +27,7 @@ The states are a collection of related classes that define a distinct state that
 which states are valid next states.
 
 ```kotlin
-sealed class Color(to: () -> Set<Color>) : app.cash.kfsm.State(to)
+sealed class Color(to: () -> Set<Color>) : app.cash.kfsm.State<Color>(to)
 data object Green : Color({ setOf(Amber) })
 data object Amber : Color({ setOf(Red) })
 data object Red : Color({ setOf(Green) })
@@ -122,7 +122,7 @@ val greenLight: Result<Light> = transitioner.transition(redLight, Go)
 
 ```kotlin
 // The state
-sealed class Color(to: () -> Set<Color>) : app.cash.kfsm.State(to)
+sealed class Color(to: () -> Set<Color>) : app.cash.kfsm.State<Color>(to)
 data object Green : Color({ setOf(Amber) })
 data object Amber : Color({ setOf(Red) })
 data object Red : Color({ setOf(Green) })

lib/src/main/kotlin/app/cash/kfsm/State.kt

@@ -1,24 +1,24 @@
 package app.cash.kfsm
 
-open class State(transitionsFn: () -> Set<State>) {
+open class State<S: State<S>>(transitionsFn: () -> Set<S>) {
 
   /** all states that can be transitioned to directly from this state */
-  val subsequentStates: Set<State> by lazy { transitionsFn() }
+  val subsequentStates: Set<S> by lazy { transitionsFn() }
 
   /** all states that are reachable from this state */
-  val reachableStates: Set<State> by lazy { expand() }
+  val reachableStates: Set<S> by lazy { expand() }
 
   /**
    * Whether this state can transition to the given other state.
    */
-  open fun canDirectlyTransitionTo(other: State): Boolean = subsequentStates.contains(other)
+  open fun canDirectlyTransitionTo(other: S): Boolean = subsequentStates.contains(other)
 
   /**
    * Whether this state could directly or indirectly transition to the given state.
    */
-  open fun canEventuallyTransitionTo(other: State): Boolean = reachableStates.contains(other)
+  open fun canEventuallyTransitionTo(other: S): Boolean = reachableStates.contains(other)
 
-  private fun expand(found: Set<State> = emptySet()): Set<State> =
+  private fun expand(found: Set<S> = emptySet()): Set<S> =
     subsequentStates.minus(found).flatMap {
       it.expand(subsequentStates + found) + it
     }.toSet().plus(found)

lib/src/main/kotlin/app/cash/kfsm/StateMachine.kt

@@ -7,18 +7,18 @@ import kotlin.reflect.full.superclasses
 object StateMachine {
 
   /** Check your state machine covers all subtypes */
-  fun <S : State> verify(head: S): Result<Set<State>> = verify(head, baseType(head))
+  fun <S : State<S>> verify(head: S): Result<Set<State<S>>> = verify(head, baseType(head))
 
   /** Render a state machine in Mermaid markdown */
-  fun <S : State> mermaid(head: S): Result<String> = walkTree(head).map { states ->
+  fun <S : State<S>> mermaid(head: S): Result<String> = walkTree(head).map { states ->
     listOf("stateDiagram-v2", "[*] --> ${head::class.simpleName}").plus(
       states.toSet().flatMap { from ->
         from.subsequentStates.map { to -> "${from::class.simpleName} --> ${to::class.simpleName}" }
       }.toList().sorted()
     ).joinToString("\n    ")
   }
 
-  private fun <S : State> verify(head: S, type: KClass<out S>): Result<Set<State>> =
+  private fun <S : State<S>> verify(head: S, type: KClass<out S>): Result<Set<State<S>>> =
     walkTree(head).mapCatching { seen ->
       val notSeen = type.sealedSubclasses.minus(seen.map { it::class }.toSet()).toList().sortedBy { it.simpleName }
       when {
@@ -29,10 +29,10 @@ object StateMachine {
       }
     }
 
-  private fun walkTree(
-    current: State,
-    statesSeen: Set<State> = emptySet()
-  ): Result<Set<State>> = runCatching {
+  private fun <S : State<S>> walkTree(
+    current: S,
+    statesSeen: Set<S> = emptySet()
+  ): Result<Set<S>> = runCatching {
     when {
       statesSeen.contains(current) -> statesSeen
       current.subsequentStates.isEmpty() -> statesSeen.plus(current)
@@ -42,9 +42,8 @@ object StateMachine {
     }
   }
 
-  @Suppress("UNCHECKED_CAST") private fun <S : State> baseType(s: S): KClass<out S> = s::class.allSuperclasses
+  @Suppress("UNCHECKED_CAST") private fun <S : State<S>> baseType(s: S): KClass<out S> = s::class.allSuperclasses
     .find { it.superclasses.contains(State::class) }!! as KClass<out S>
-
 }
 
 data class InvalidStateMachine(override val message: String) : Exception(message)

lib/src/main/kotlin/app/cash/kfsm/States.kt

@@ -1,13 +1,13 @@
 package app.cash.kfsm
 
 /** A collection of states that is guaranteed to be non-empty. */
-data class States<S : State>(val a: S, val other: Set<S>) {
+data class States<S : State<S>>(val a: S, val other: Set<S>) {
   constructor(first: S, vararg others: S) : this(first, others.toSet())
 
   val set: Set<S> = other + a
 
   companion object {
-    fun <S: State> Set<S>.toStates(): States<S> = when {
+    fun <S: State<S>> Set<S>.toStates(): States<S> = when {
       isEmpty() -> throw IllegalArgumentException("Cannot create States from empty set")
       else -> toList().let { States(it.first(), it.drop(1).toSet()) }
     }

lib/src/main/kotlin/app/cash/kfsm/Transition.kt

@@ -1,6 +1,6 @@
 package app.cash.kfsm
 
-open class Transition<V: Value<V, S>, S : State>(val from: States<S>, val to: S) {
+open class Transition<V: Value<V, S>, S : State<S>>(val from: States<S>, val to: S) {
 
   init {
     from.set.filterNot { it.canDirectlyTransitionTo(to) }.let {

lib/src/main/kotlin/app/cash/kfsm/Transitioner.kt

@@ -1,6 +1,6 @@
 package app.cash.kfsm
 
-abstract class Transitioner<T : Transition<V, S>, V : Value<V, S>, S : State> {
+abstract class Transitioner<T : Transition<V, S>, V : Value<V, S>, S : State<S>> {
 
   open fun preHook(value: V, via: T): Result<Unit> = Result.success(Unit)
 

lib/src/main/kotlin/app/cash/kfsm/TransitionerAsync.kt

@@ -1,6 +1,6 @@
 package app.cash.kfsm
 
-abstract class TransitionerAsync<T : Transition<V, S>, V : Value<V, S>, S : State> {
+abstract class TransitionerAsync<T : Transition<V, S>, V : Value<V, S>, S : State<S>> {
 
   open suspend fun preHook(value: V, via: T): Result<Unit> = Result.success(Unit)
 

lib/src/main/kotlin/app/cash/kfsm/Value.kt

@@ -1,6 +1,6 @@
 package app.cash.kfsm
 
-interface Value<V: Value<V, S>, S : State> {
+interface Value<V: Value<V, S>, S : State<S>> {
   val state: S
   fun update(newState: S): V
 }

lib/src/test/kotlin/app/cash/kfsm/Letter.kt

@@ -5,7 +5,11 @@ data class Letter(override val state: Char) : Value<Letter, Char> {
   override fun update(newState: Char): Letter = copy(state = newState)
 }
 
-sealed class Char(to: () -> Set<Char>) : State(to)
+sealed class Char(to: () -> Set<Char>) : State<Char>(to) {
+  fun next(count: Int): List<Char> =
+    if (count <= 0) emptyList()
+    else subsequentStates.filterNot { it == this }.firstOrNull()?.let { listOf(it) + it.next(count - 1) } ?: emptyList()
+}
 
 data object A : Char(to = { setOf(B) })
 data object B : Char(to = { setOf(B, C, D) })

lib/src/test/kotlin/app/cash/kfsm/StateMachineTest.kt

@@ -55,21 +55,21 @@ class StateMachineTest : StringSpec({
   }
 })
 
-sealed class ValidState(to: () -> Set<ValidState>) : State(to)
+sealed class ValidState(to: () -> Set<ValidState>) : State<ValidState>(to)
 data object Valid1 : ValidState({ setOf(Valid2, Valid3) })
 data object Valid2 : ValidState({ setOf(Valid3) })
 data object Valid3 : ValidState({ setOf(Valid4, Valid5) })
 data object Valid4 : ValidState({ setOf() })
 data object Valid5 : ValidState({ setOf() })
 
-sealed class UniCycleState(to: () -> Set<UniCycleState>) : State(to)
+sealed class UniCycleState(to: () -> Set<UniCycleState>) : State<UniCycleState>(to)
 data object UniCycle1 : UniCycleState({ setOf(UniCycle1) })
 
-sealed class BiCycleState(to: () -> Set<BiCycleState>) : State(to)
+sealed class BiCycleState(to: () -> Set<BiCycleState>) : State<BiCycleState>(to)
 data object BiCycle1 : BiCycleState({ setOf(BiCycle2) })
 data object BiCycle2 : BiCycleState({ setOf(BiCycle1) })
 
-sealed class TriCycleState(to: () -> Set<TriCycleState>) : State(to)
+sealed class TriCycleState(to: () -> Set<TriCycleState>) : State<TriCycleState>(to)
 data object TriCycle1 : TriCycleState({ setOf(TriCycle2) })
 data object TriCycle2 : TriCycleState({ setOf(TriCycle3) })
 data object TriCycle3 : TriCycleState({ setOf(TriCycle1) })

lib/src/test/kotlin/app/cash/kfsm/StateTest.kt

@@ -55,4 +55,9 @@ class StateTest : StringSpec({
     E.canEventuallyTransitionTo(E) shouldBe false
   }
 
+  "can define and use custom methods on a state" {
+    A.next(2) shouldBe listOf(B, C)
+    E.next(2) shouldBe emptyList()
+  }
+
 })

lib/src/test/kotlin/app/cash/kfsm/exemplar/Hamster.kt

@@ -16,7 +16,7 @@ data class Hamster(
     println("◟(`･ｪ･)  ╥━╥   (goes to bed)")
   }
 
-  sealed class State(to: () -> Set<State>) : app.cash.kfsm.State(to)
+  sealed class State(to: () -> Set<State>) : app.cash.kfsm.State<State>(to)
 
   /** Hamster is awake... and hungry! */
   data object Awake : State({ setOf(Eating) })