getkyo / kyo

Toolkit for Scala Development
https://getkyo.io
Apache License 2.0
527 stars 44 forks source link

new core design (kyo-prelude) #489

Closed fwbrasil closed 3 months ago

fwbrasil commented 3 months ago

This PR is part of the work to introduce a new overall design for the library core abstractions as explored in https://github.com/getkyo/kyo/issues/462. Making a single large change for all effects would be too challenging so I'm splitting the work into two main changes:

  1. Introduce kyo-prelude: I've been wanting to split code that doesn't rely on IOs into a separate module (see https://github.com/getkyo/kyo/pull/308) so I'm using this opportunity to introduce the new module as an isolated implementation with the new design under a separate package kyo2. Enabling both versions of the core design to be in the codebase will ease migration. For example, we can implement benchmarks for it in kyo-bench. The module name is inspired by ZIO :)
  2. Port kyo-core: I'm planning to work on a separate PR that ports kyo-core to use the new design in kyo-prelude. For the files that this PR duplicates, I'll first delete the duplicate file, move the previous one to the new location, and then update it with the new code. This will allow us to properly keep the git history. Other modules will also need to be ported but they shouldn't require significant changes.

This new core design introduces several major improvements:

1. Computation nesting

Kyo's flat representation of computations doesn't seem a major limitation given that composable effects generally preclude the need for nesting but the flat representation introduces adoption and usability issues. Users can get confused by the approach diverging from other similar libraries and passing the Flat evidence around in user code doesn't seem ideal but is required in generic contexts that perform effect handling.

In this new design, the < type is an AnyVal. This enables efficient unboxed execution in the majority of the scenarios and, when a Kyo computation is passed to a generic scope, it's automatically boxed by the compiler, which enables safe execution of nested computations and avoids the need for the Flat check.

I've also explored defining < as a context function: type <[+T, -S] = Runtime ?=> Kyo[T, S]. Although the compiler desugaring is able to avoid allocations in more scenarios than with AnyVal, the fact that context functions can't be aliased by an opaque type introduces usability issues since it exposes the internal Kyo[T, S] type to the user and the compiler can infer it instead of < in certain situations. I think we can still explore this path further later, especially as the compiler has been evolving significantly lately.

2. Isolated and user-friendly kernel

I've renamed core to kernel since it'd be confusing for the core to be in kyo-prelude instead of kyo-core. The new kernel implementation reorganizes APIs to be more user-friendly so we can start allowing users to create new effects.

Something that has been challenging while evolving Kyo's design is the fact that internal core APIs are used in a few places of the codebase like in IOTask and ZIOs. The new kernel API is designed to fulfill the needs of all effects, including IOTask and ZIOs via new APIs that enable handling multiple effects at once. This will allow us to iterate faster on the core design in isolation. I think the only missing feature we'll need to port kyo-core is how to handle IOs.ensure/Resources and preemptions, which I'm planning to work on as a follow up.

3. Naming improvements

Kyo's naming for effects is based on early versions of the library where most effects had a monad counterpart: Future/Futures, Try/Tries, Option/Options, etc. This hasn't been the case for some time now so there's no need for the pattern using plurals for effect names. In case we have such cases, we can also have internal effects in the companion object (see Local in this PR as an example).

I'm also moving away from T/U/etc for generic type names to A/B/etc

4. Type-safe effect handling

The current lack of type safety when handling effects is a major limitation to allow users to create new effects. In this new design, effects are defined as an abstract polymorphic arrow:

// Note how an effect has to define lambdas for its input and output types
abstract class Effect[-I[_], +O[_]]

// For example, `Sum` is defined as an arrow that takes a constant `V` and returns `Unit`
sealed trait Sum[V] extends Effect[Const[V], Const[Unit]]

// `Choice` that takes a `Seq` of an arbitrary type and returns the type itself (`Id`)
sealed trait Choice extends Effect[Seq, Id]

The effect handling APIs are also more streamlined without requiring Handler classes. This means that handling is allocation-free by default without a need to cache handler instances. The new encoding of effects as arrows also enables fully type-safe effect handling. The best example of it is Var:

// If I pass the wrong type to `cont` for each of the input types, the compilation fails!
Effect.handle.state(tag, state, v) {
    [C] =>
        (input, state, cont) =>
            input match
                case input: Get[?] =>
                    (state, cont(state)) 
                case input: Set[?] =>
                    (input(), cont(())) // fails to compile if I pass `state` to cont for example!
                case input: Update[?] =>
                    (input(state), cont(()))
}

The effect handling is also "pinned" to the result type of the computation being handled, which is something that wasn't possible with the previous design using separate handler classes. I think we don't have effects that can leverage that yet but it opens new possibilities for how to implement effects.

5. Stack safety

Kyo's current design doesn't provide stack safety by default and users need to rely on effect suspensions or Loops to ensure stack safety for recursive computations. Although this isn't a major limitation given that recursive computations without effect suspensions don't make much sense and we provide convenient alternatives, providing stack safety should help with adoption since it follows the behavior of other effect systems.

The execution of the new kernel is stack-safe for recursive computations in general but there's still a limitation with computations that dynamically accumulate a large number map transformations without any effect suspensions since the kernel still leverages the stack to "unfold" the nested computation when resuming from an effect suspension. I've added pending tests showing this limitation. I explored approaches to avoid this limitation as well but this is a central aspect to optimize the execution of Kyo computations via monomorphization. It seems a reasonable limitation and I'm keeping the Loop effect as an alternative.

6. Built-in Defer effect

The current core design attempts to provide the core abstractions to define effects without including any effects in the core itself. This approach has two major downsides:

The new kernel has an internal Defer effect that is used to provide stack safety and encode the use of local values without a separate effect. This approach means that a A < Any is actually a A < Defer internally, which is why kyo.pure is renamed to kyo.eval since it can have pending internal Defer suspensions.

7. New runtime effects

There are effects like Env and Local that only require propagating values into the computation without the need for a proper effect handling with a continuation. The new kernel provides a new kind of effect for these scenarios:

// The `A` type parameter specifies the type of the local value the effect uses
abstract class RuntimeEffect[+A]

// For example, `Env` uses a `TypeMap`
sealed trait Env[+R] extends RuntimeEffect[TypeMap[R]]

// `Local` uses an untyped map (State is its internal effect type)
sealed private trait State extends RuntimeEffect[Map[Local[?], AnyRef]]

Suspension of runtime effects has an interesting characteristic. If the suspension provides a default value, the effect isn't added to the pending set.

// The `Env` effect suspends without a default value so it's added to the pending set
// (this is a simplification of the actual impl that requires erased tags)
def get[R](using tag: Tag[Env[R]]): R < Env[R] =
    RuntimeEffect.suspend(tag)(_.get(using tag))

// But `Local` provides a default value (`Map.empty`) so it doesn't get added to the pending set
def get: T < Any =
    RuntimeEffect.suspend(Tag[Local.State], Map.empty)(_.getOrElse(this, default).asInstanceOf[T]) 

This approach is designed so we can easily propagate runtime effects when forking fibers. The method RuntimeEffect.boundary saves the state of the runtime, which includes runtime effect values and traces, and provides a function to resume the computation using the original runtime state. The method is currently unused and I'll explore it further when porting kyo-core.

// I'm using a match type to remove runtime effects from the pending set
type WithoutRuntimeEffects[S] = S match
    case RuntimeEffect[x] => Any
    case s1 & s2          => WithoutRuntimeEffects[s1] & WithoutRuntimeEffects[s2]
    case _                => S

// Note how `f` is a function that produces a computation without pending runtime effects
// since they're already handled by the "propagated" runtime
def boundary[A, B, S, S2](v: A < S)(
    f: (() => A < WithoutRuntimeEffects[S]) => B < S2
)(using _frame: Frame): B < (S & S2) =

8. Stack traces

The new Runtime mechanism also collects traces of the execution and automatically injects frames into exception stack traces. The boundary method is also designed to propagate the trace to forked fibers. Example output:

java.lang.Exception: test exception
    at kyo2.TraceTest.ex(TraceTest.scala:8)
    at      x.map(_ => throw ex) @ kyo2.TraceTest.boom(TraceTest.scala:10)
    at , y).map(_ + _).map(boom) @ kyo2.TraceTest.withEffects(TraceTest.scala:17)
    at  Kyo.zip(x, y).map(_ + _) @ kyo2.TraceTest.withEffects(TraceTest.scala:17)
    at             Kyo.zip(x, y) @ kyo2.TraceTest.withEffects(TraceTest.scala:17)
    at       Env.use[Int](_ * 2) @ kyo2.TraceTest.loop(TraceTest.scala:16)
    at             Kyo.zip(x, y) @ kyo2.TraceTest.withEffects(TraceTest.scala:17)
    at       Env.use[Int](_ + 1) @ kyo2.TraceTest.loop(TraceTest.scala:15)
    at        Env.run(1)(z).eval @ kyo2.TraceTest.withEffects(TraceTest.scala:18)
    at kyo2.TraceTest.withEffects(TraceTest.scala:18)

Note how traces also include a short snippet of the code of the frame. This behavior is on by default but we should provide a way to disable it since I imagine it could break tooling that read stack traces. I think we could also explore providing a separate way to render traces by compacting the information into a base64 string and providing a link to a static page that reads the payload and renders the trace.

9. Remove Options effect

I wasn't sure what to name Options since we're moving away from plurals and decided to remove it instead. I think I've never used it and we can provide the same functionality via Abort directly. I've added a Abort.get(someOption) method instead.

Conclusion

I apologize for the large and intrusive change but I think it's necessary to evolve the library faster. This new design gets us much closer to a stable set of primitives.

fwbrasil commented 3 months ago

I've just pushed an important improvement. The encoding with < defined as opaque type <[+T, -S] >: T didn't allow us to avoid nesting of computations by the compiler since the value was widened directly by the type system. With the new kernel, the lifting is done via an implicit conversion so it's possible to use a NotGiven to avoid lifting nested computations:

implicit inline def lift[A, S](v: A)(using inline ng: NotGiven[A <:< (Any < Nothing)]): A < S = <(v) 

This is an important behavior for usability because mismatch of pending effects trigger lifting. For example:

def test[T](v: T < IO): T < IO = ...
val a: Int < Choice = ...

// Without `NotGiven` the compiler nests the computation if there's a mismatch.
// This fails to compile after my last change
val b: Int < Choice < IO = test(a)

Although nesting is safe with the new encoding, this behavior when there's a mismatch in pending effects is an important usability issue. Note that nesting is still possible indirectly as demonstrated in the tests (see PendingTest.flatten as an example).

fwbrasil commented 3 months ago

Runtime is an overloaded term, especially considering the parallel with ZIO's Runtime, and gives an impression of an impure effect, which isn't the case. I've renamed Runtime to Safepoint, RuntimeEffect to ContextEffect, and Values to Context. It seems to communicate better what those types do.

fwbrasil commented 3 months ago

After multiple battles with the compiler, I've finally been able to make Boundary work to fork fibers! The match type approach wasn't working properly so I tried a few mechanisms similar to Has we have for Abort and Env, which didn't work, and ended landing on a macro to make the type-level transformation: https://github.com/getkyo/kyo/pull/489/commits/d631a99d1980606ea833f0a11d0f4412ea9a43c7

I think there isn't anything major pending in this PR after this change.