monad.kt

sealed class Maybe<out T> {
    infix fun <R> flatMap(f: (T) -> Maybe<R>): Maybe<R> {
        when (this) {
            is Just -> return f(this.value)
            None -> return None
        }
    }

    infix fun <R> map(f: (T) -> R): Maybe<R> {
        when (this) {
            is Just -> return Just(f(this.value))
            None -> return None
        }
    }

    companion object {
        fun <T> unit(value: T): Maybe<T> {
            return Just(value)
        }
    }

    class Just<out T>(val value: T) : Maybe<T>()
    object None : Maybe<Nothing>()
}

class State<S, out T>(val runState: (S) -> Pair<T, S>) {
    infix fun <R> bind(f: (T) -> State<S, R>): State<S, R> {
        return State {
            it ->
            val (value, state) = this.runState(it)
            f(value).runState(state)
        }
    }

    infix fun <R> map(f: (T) -> R): State<S, R> {
        return State {
            state ->
            val (value, prevState) = runState(state)
            val newValue = f(value)
            Pair(newValue, prevState)
        }
    }

    companion object {
        fun <S, T> unit(value: T) = State<S, T> { Pair(value, it) }
        fun <S> get(): State<S, S> = State { Pair(it, it) }
        fun <S> put(state: S): State<S, Unit> = State { Pair(Unit, state) }
    }
}

class Reader<E, out T>(val runReader: (E) -> T) {
    infix fun <R> bind(f: (T) -> Reader<E, R>): Reader<E, R> {
        return Reader {
            val value = this.runReader(it)
            f(value).runReader(it)
        }
    }

    infix fun <R> map(f: (T) -> R): Reader<E, R> {
        return Reader {
            val value = runReader(it)
            f(value)
        }
    }

    companion object {
        fun <E, T> unit(value: T) = Reader<E, T> { value }
        fun <E> ask(): Reader<E, E> = Reader { it }
    }
}

class Writer<out V, L>(val runWriter: () -> Pair<V, List<L>>) {
    infix fun <R> bind(f: (V) -> Writer<R, L>): Writer<R, L> {
        return Writer({
            val (value, log) = runWriter()
            val (newValue, newLog) = f(value).runWriter()
            Pair(newValue, log + newLog)
        })
    }

    infix fun <R> map(f: (V) -> R): Writer<R, L> {
        return Writer {
            val (value, log) = runWriter()
            Pair(f(value), log)
        }
    }

    companion object {
        fun <V, L> unit(value: V) = Writer { Pair(value, emptyList<L>()) }
        fun <V, L> tell(value: V, log: List<L>): Writer<V, L> = Writer { Pair(value, log) }
        fun <L> tell(log: L) = Writer<Unit, L> { Pair(Unit, listOf(log)) }
    }
}

Test.kt

class MonadTest {
    @Test
    fun testMaybeMonad() {
        val result = Maybe.unit(1) flatMap {
            if (it > 0) Maybe.unit("Hello") else Maybe.None
        } map {
            "Hi, $it"
        } map {
            "$it ..."
        } map {
            10
        } map {
            it + 10
        }

        when (result) {
            is Maybe.Just -> println("value = ${result.value}")
            Maybe.None -> println("None")
        }
    }

    @Test
    fun testStateMonad() {
        val state = State.unit<Int, String>("Hello") bind {
            value ->
            State.get<Int>() bind {
                state ->
                if (state > 3) {
                    State.put(100) bind {
                        State.unit<Int, String>(value)
                    }
                } else {
                    State.unit<Int, String>("Not greater than 3")
                }
            }
        } map {
            it + "..."
        }

        val (v, s) = state.runState(1)
        println("v = $v, s = $s")

        val (vv, ss) = state.runState(10)
        println("v = $vv, s = $ss")
    }

    @Test
    fun testReaderMonad() {
        val program = Reader.unit<Int, Int>(10) bind {
            Reader.unit<Int, Int>(it + 10)
        } bind {
            value ->
            Reader.ask<Int>() bind {
                env ->
                Reader.unit<Int, Int>(value + env)
            }
        } map {
            it + 20
        }

        val result = program.runReader(1000)

        println("Result = $result")
    }

    @Test
    fun testWriterMonad() {
        val p = Writer.unit<Int, String>(1) map {
            it + 1
        } map {
            it + 2
        } bind {
            Writer.tell(it, listOf("Hello, $it"))
        } map {
            it + 10
        } bind {
            Writer.tell(it, listOf("Yes, now, it's $it"))
        }

        val (value, log) = p.runWriter()
        println("value = $value, log = $log")
    }
}