Stacks for Kotlin: Design Notes

[RossTate]

This is a discussion of the design notes for Stacks for Kotlin. These notes summarize ongoing research on a library providing primitives for creating and working with first-class call-stacks, a feature that provides safer and more compositional support for coroutines. Note that this work heavily relies on local lifetimes, so the reader is advised to first familiarize themselves with that research before proceeding with these notes.

Disclaimer: This work is still research in progress. We are sharing it early because we want to collect feedback (hopes, concerns, suggestions, and so on) from the community before determining whether to develop it into a formal proposal.

[kyay10]

Apologies, I haven't fully read the proposal. I'm somewhat struggling to understand the motivation and purpose behind this feature.
How does it compare to shift/reset? It looks like it's effectively doing that, just with non-standard names. It's reminiscent of stack-switching as a technique for implementing delimited continuations.

It's my understanding that, while suspend obviously implements delimited continuations, it doesn't provide a shift/reset style mechanism of interfacing with it (i.e. "composable continuations"). So is this effectively adding that, or something that's a bit lower-level than it?

[kyay10]

(I'm referencing shift/reset terms, which I've laid out in the other discussion)

The arena lifetime seems to me like the lifetime that shift puts you in, and that the Cont provided by shift is accessible in. It's the lifetime outside of reset but that reset has access to. In other words, it's the lifetime parameter of Prompt

The environment sounds to me like the lifetimes of the caller to Cont.resume, i.e. each environment is simply the lifetime of the callsite resuming the stack, and thus it can't be relied on to "stick around" since it's dynamic.

Finally, the execution lifetime sounds like the lifetime inside reset i.e. it's the lifetime of the Prompt.

The lifetime of StackSuspension is thus that of Cont, while its resumption lifetime parameter is precisely the lifetime parameter of Cont (i.e. it's the lifetime at the call-site of shift).
In some sense, it seems that the StackSuspension is a primitive that a Cont would hold. StackSuspension is also single-shot, of course, but we can make Cont single-shot for the purposes of this discussion.
In fact, it seems my Cont is very close to StackContinuation, except that StackContinuation is maybe more flexible (although you can just have a Cont<Unit, Nothing> to represent a StackContinuation<() -> Nothing>, but I digress).

For instance, I can do ignoreInput too!

fun <T> ignoreInput(local continuation: Cont<Unit, Nothing>): Cont<T, Nothing>_{continuation}
  = Cont { _ -> continuation.resume(Unit) }

I think the analogy maybe breaks down, though, because resumer seems to be only usable inside the coroutine, while Cont's resume is usable outside... I think I'm getting confused here, so I'll carry on reading.

The StackMount looks like a Prompt to me. Its E type parameter is Prompt's type parameter. StackMount.new seems to be like shift, although there are stark differences.
I'm not fully understanding the trick with mounted either. It seems like it effectively restricts when shift/new can be used? I'm not sure I'm fully getting it.
It seems the resumer is somewhat like resumePush, in that it gets to run inside the coroutine, and thus access everything internal to it. That's why it must always have access to mounted, because inside resumePush, obviously you can call shift!
So it seems maybe StackContinuation is a Cont but also allowing you to place things inside that are effectively resumePushed. I guess it's effectively like a fun interface Cont, since you can use other Conts inside and thus resumePush to access their resumption lifetime, just like resumer seems to have access to.

restack feels like an entry-point to this world, so maybe like a runCC. Effectively, it's something that would be stashed away in CoroutineContext if that were an option? Some machinery akin to a Dispatcher or a trampoline or something like that?

StackRestacker.mount seems to be a resumePush of sorts? "Mounting" a StackMount sounds like adding a delimiter onto the current continuation, and then the suspension effectively is the segment of a suspended continuation that we're glueing back on the stack. I don't quite understand what the environment is just yet. Maybe some coroutine context esque thing? It reminds me a lot of delimited dynamic binding.
It also seems to double as the reset primitive if given an empty StackSuspension to work with, which makes sense.
This is really interesting because it mirrors a limitation that optimized implementations of multi-prompt delimited continuations have to contend with: having a Prompt delimiting multiple stacks at the same time seems to ruin performance; recent papers have thus imposed a non-reentrancy restriction on Prompts. Using your terminology, a mount cannot be mounted while it's already mounted!

StackRestacker.dismount is very clearly a shift. It gives you access to the suspension, which is (almost) a Cont, and it brings you up to the arena. The environment again feels like a dynamic binding thing.

I'm completely blanking on what StackRestacker.switchTo is analogous to.

[kyay10]

Why not go for a more familiar shift/reset design? The practical use cases presented are also implementable using shift/reset. Well, I'm not sure about deep recursion, but at least PausingStack can be easily implemented:

local class PausingStack(
    private val block: (local pause: () -> Unit) ->_{this} Unit
) {
    private var continuation: Cont<Unit, Unit>? = null
    init {
        reset {
            fun pause() = shift { continuation = it; false }
            pause()
            block(::pause)
            true
        }
    }
    fun progress(): Boolean  = (continuation ?: return true).also { continuation = null }.resume(Unit)
}

FWIW, I already have a multiplatform implementation of single-shot multi-prompt delimited continuations using suspend. It also supports multi-shot using reflection on JVM/JS, or bytecode transformation on JVM, but the single-shot implementation uses none of that.

[kyay10]

The new background section seems to arrive at what's effectively multi-prompt delimited continuations, with suspend<R> being a Prompt<R>, and the continuations being the "composable" continuations (that I refer to as Cont)

[kyay10]

The entire Flow section should probably be moved to the local lifetimes notes, since it doesn't use anything related to the stacks.

[RossTate]

In a sense you're right, but flows themselves are closely tied to coroutines and suspend. In other words, although the fixes to flows don't themselves directly use stacks, they rely on the fact that stacks in general can be used to support coroutines and suspend.

[kyay10]

It would be nice if lifetime members get introduced in the local lifetime notes. They could be an important point of discussion.

[RossTate]

Really the question is whether or not there should be members or not, regardless of type or locality. This is theoretically an orthogonal decision. Here it makes the API much simpler, so I assumed it was supported for the sake of focusing on the more important aspects of the design. However, there is a real possibility that these primitives will have to be repackaged in a way that does not use members.

[rnett]

+1 on discussing them somewhat separately from stacks, even if it's a dedicated section. It's not clear to me how they are related to lifetime type params or how having a compile time construct (IIRC?) as a member even works.

[kyay10]

For StackMount, how can it have its mounted lifetime member private, yet it's mentioned in a public signature (new)?
How does one call new if they don't know what mounted is? Maybe I'm misunderstanding lifetime members

[RossTate]

The mounted locality is public; what actual locality it represents is private.

[kyay10]

StackRestacker seems to reference this E for an environment type quite often, but without ever declaring it. Is it a type parameter on each function? e.g. is mount actually written like so:

fun <E, local_{resumption} arena, local that> mount(
    mount: StackMount<arena, E>_{local},
    environment: E_{resumption},
    suspension: StackSuspension<that>_{mount.mounted},
    block: local StackRestacker<that>.() ->_{that} Nothing
): Nothing

[RossTate]

Ah, oops! It's supposed to be a type parameter of mount and dismount. Fixed! Thanks.

[kyay10]

This echoes a question I asked on the lifetimes notes. Currently, restack is defined as such

fun restack(once block: local StackRestacker<local>.() -> R): R

but could it also be defined as:

fun restack(once block: local StackRestacker<block>.() -> R): R

and what would be the differences, if any?

[RossTate]

The intent is that it is supposed to capture the locality available at the call site to restack, which is the mental model for what the local locality represents. On the other hand, the mental model for what the block locality represents is the lifetime of the block parameter. In actuality, both are locality parameters that get inferred with various constraints, so in this case it seems like either would work.

[kyay10]

Is all the environment juggling necessary for this design? It seems like environments could simply be a reference that gets set on the mount. Perhaps lifetime safety necessitates such juggling?
A mount can be mounted at most once before being unmounted, so at most there's a single environment available, so it seems a variable on the mount could do the trick. Maybe there's something I'm not seeing with how the environment can actually shift its lifetime, and hence a variable wouldn't reflect that aptly.
Ohhh, is it because previous environments need to stop being accessible? Hence, each environment is given its own lifetime, and you only know that you have access to it upon dismounting, but remounting doesn't guarantee such access. That seems quite clever!

[RossTate]

Heh, you ended up answering your own question.

[kyay10]

Can't dismount be simplified to the following?

fun <local_{resumption} arena> dismount(
    mount: StackMount<arena, E, mounted_{resumption}>_{local},
    block: local^{environment} StackRestacker<environment>.(
        local^{arena} environment: E,
        StackSuspension<resumption>_{mount.mounted}
    ) -> Nothing
): Nothing

i.e. because block already takes the environment locally, it has the lifetime of the environment.

[RossTate]

Ah, there was a typo. The signature should be

fun <local_{resumption} arena, E> dismount(
    mount: StackMount<arena, E, mounted_{resumption}>_{local},
    block: local^{environment} StackRestacker<environment>.(
        local^{arena} environment: E,
        StackSuspension<resumption>_{mount.mounted}
    ) ->_{arena} Nothing
): Nothing

Without that, the declaration local^{arena} environment is invalid. It introduces a locality variable whose lower bound is local and whose upper bound is arena. For the algorithm to work, a the lower and upper bound of a variable must be subtypes (before declaring the variable). The effect of adding the ->_{arena} restriction is to make its local locality parameter be upper bounded by arena, which then implies the subtyping required for local^{arena} environment to be valid.

[kyay10]

Can you elaborate a bit on why restack and finish are necessary? They seem superficial to me. What do they actually do? What is a "restacking session"?

[RossTate]

restack and finish are the start and end of a restacking session. A restacking session essentially lets you directly perform the various operations that other systems perform for you in predetermined sequences behind the scenes. Unfortunately, I haven't had a chance to highlight some of the more advanced operations this enables while maintaining both safety and efficiency, in particular when implementing structured concurrency in the kotlinx.coroutines library. I hope to get around to writing that up some day.

[kyay10]

It seems to me that if you remove all the lifetimes and add suspend everywhere, this stacks library maybe could be implementable in Kotlin right now. Has such an effort been made? I feel it would help in understanding this quite a lot

[kyay10]

Sounds like a challenge to me!

On a very theoretical level, there's likely some Monad that implements the stacks library. Since shift/reset is already expressible using current primitives, we can use Filinski's monadic reflection, so it seems like it should just work? It wouldn't necessarily have the same performance characteristics, but it should otherwise have the same semantics.

[kyay10]

I've taken up the challenge! My implementation uses shift/reset (I call them handle/use in my library), but it doesn't support switchTo. restack and finish do nothing effectively.

Of course, there is no lifetime safety. I might end up trying to use my monadic regions system to make it safe, but it'll be quite ugly!
This helped confirm my suspicions that dismount is shift and mount is resume.

[kyay10]

I'm back with even better news! I stripped away my code fully so that it's just implementing stacks, and I now have an implementation in under 200 LoC! I see what you mean now about switchTo, although I'm still convinced that it's effectively a dismount then a mount.
finish now is responsible for trampolining, so that no stack overflows occur.
I haven't fully implemented the difference between Pending and Active, but it seems like that's not necessary for the implementation to be correct.
It'd be great if you could have a look! I'm hoping to effectively have a reference implementation so that I understand the system better.

[kyay10]

For completeness sake, I also made a lifetime-safe version, using @RestrictsSuspension and type parameters to encode lifetimes, and a standard "flattening" trick to encode higher-kinded types (for things like the environment and the resumer). It helped find a couple mistakes! It also just made the lifetime portion of the code way more clear to me

[RossTate]

Ah, sorry, I should have been clearer about the implicit constraints in my statement. You've done the stack-of-stacks inside a trampoline encoding (though you've done so in an impressively thorough and compact manner!). The stack-of-stacks lets you emulate suspending to the non-closest handler. The trampoline lets you emulate switching. Unfortunately, it requires you to first enter the trampoline, which is why your restacker has to be marked as suspend and requires a Locality parameter. And its correctness relies on there being one trampoline in place at a time, which causes compositionality issues. Nonetheless! This is likely the encoding we'll have to employ if we want to use WebAssembly's poorly designed stack-switching proposal, which only allows suspending up to the nearest handler.

[kyay10]

With this and lifetimes, is suspend ever necessary? Does suspend effectively become shorthand for some context(local _: SuspendScope)?

[RossTate]

That's definitely a possibility.

[kyay10]

Can sequence be simplified to the following?

public fun <T> sequence(local_{} block: local SequenceScope<T, block>.() -> Unit): Sequence<T>_{block}

Note the lower-bound on block so that it's not "discharged" where sequence is called.

On the other side,iterator feels unnecessarily limiting. Shouldn't it be defined as follows?

public fun <T> iterator(
    local_{} block: local SequenceScope<T, block>.() -> Unit
): Iterator<T>_{block}

So that one can make an iterator even if they can't do block actions right now.

[RossTate]

Oops, those blocks were apparently not properly updated. Fixed!

[kyay10]

One detail is that, if an exception gets thrown from the suspendingBlock of one of these separate call stacks, then that exception is passed onto and thrown from the last invocation to resume that caused the separate call stack to execute. So, there's a way to throw an exception to the resumption site, but interestingly there's no way to return a value to the resumption site.

I don't think that's true! I'll focus on JVM here, but I'm very sure that this applies to all platforms.
Firstly, coroutines generally treat values and exceptions the same way, which is really important.
If we take a look at how resumeWith is implemented, we see this important snippet:

val outcome: Result<Any?> =
    try {
        val outcome = invokeSuspend(param)
        if (outcome === COROUTINE_SUSPENDED) return
        Result.success(outcome)
    } catch (exception: Throwable) {
        Result.failure(exception)
    }

invokeSuspend is where the code inside the coroutine lies. Thus, any exception thrown within suspendBlock is actually caught, ends up propagating up its stack, and then it eventually will reach the completion provided to createCoroutineUnintercepted

In fact, your createCoroutineUnintercepted is missing a key component: completion! It should actually read like so:

public fun <T> createCoroutineUnintercepted(suspendingBlock: suspend (T) -> Unit, completion: Continuation<Unit>): Continuation<T>

Technically, you can throw an exception at the site of resume by implementing your own Continuation, but for normal suspend code, that doesn't happen. A lot of the coroutine machinery really doesn't like if you do that, btw, and end up with unreported errors.

[RossTate]

Technically, you can throw an exception at the site of resume by implementing your own Continuation

Yes, so you can do what I said. The primitive I gave is equivalent to the one in the standard library, as in either can implement the other. I'm just trying to keep things simple for the sake of explanation. Regardless, the point I'm trying to make here isn't really about exceptions; it's about returning values.

[kyay10]

Fair enough. Apologies if I came off too strongly!

FWIW, such an exception doesn't even have to originate from inside the block. It can be from the completion too! Needless to say, resumeWith is messy to work with it you have custom Continuations that want to throw exceptions to it.

[RossTate]

No worries! Thanks for deeply engaging with the design!

[kyay10]

On the other hand, suppose suspendingBlock calls suspendCoroutineUnintercepted(handlingBlock).
Then a few steps happen:

1. The separate call stack executing suspendingBlock gets suspended.
2. We switch back to wherever the call to snapshot.resume(...) was made.
3. We call handlingBlock

That's not true! I'm pretty sure 3 and 4 happen before 2. If you throw an exception within suspendCoroutineUnintercepted, it actually resumes the continuation. Similarly, you stay on the same thread and everything that the call to suspendCoroutineUnintercepted is on. The switchback eventually happens, but only after 4.

[RossTate]

There are many slight variations of these primitives that are all encodable within each other. I chose variations that made for simpler explanations. The point of this is to give enough background to discuss the key issues with the current primitives, which are unaffected by these small differences.

[kyay10]

A few tiny nits for the code (I found these by removing lifetimes and trying to compile the code). Some of them do affect understanding, while others are just compiler technicalities:

new has an extraneous input parameter, and its call to dismount is accidentally destructuring its arguments. It should read like so:

= new {
    val output = block()
    restack {
        dismount(this@new) { environment, _ ->

resume and suspend have a useless type parameter T. Its call to mount forgets to pass the environment, and it also refers to this@new, which is likely intended to be this@resume. It should read like so:

fun <local_{local} arena, E, R> StackMount<arena, E>.resume(
    ...
): Nothing {
    restack {
        mount(environment, this@resume, continuation.suspension) {

Similarly for suspend:

fun_{mounted} <local arena, E, R> StackMount<arena, E>.suspend(
    ...
): Nothing {
    restack {
        dismount(this@suspend) { environment, suspension ->

PausingStack has similar issues, but also an extra }, and its call to resume is missing a lambda:

local class PausingStack(
    private val block: (local pause: () -> Unit) ->_{this} Unit
) {
    private[this] val mount: StackMount<this, (Boolean) -> Nothing>
        = StackMount()
    private[this] var continuation: StackContinuation<() -> Nothing>_{mount.mounted}?
        = mount.new({ exit, _ -> exit(true) }) {
            block pause@{
                mount.suspend({ return@pause }) { exit, continuation ->
                    this.continuation = continuation
                    exit(false)
                }
            }
        }
    fun progress(): Boolean {
        mount.resume(
            { return@progress it },
            (continuation ?: return true).also { continuation = null }
        ) {
            it()
        }
    }
}

onFreshStack creates a StackMount without specifying its type parameter, but that's likely because outference can figure it out! However, the type suggested for it in the comment is wrong. It should instead be StackMount<local, Unit> because that's the environment passed to it.

Also, onFreshStack { function.block(value) } doesn't compile, instead requiring onFreshStack { function.block(this, value) }.

Sequence is missing the fun modifier for its interface. There's also a copy-paste error in the implementation for sequence, it should actually be = Sequence { iterator(block) }.

There's also a little bit of strange formatting, like splitting extension function types with empty parameters (see e.g. finish), but that could be intentional.

[RossTate]

Wow, thanks! I fixed most of them, though there were a few I either misunderstood or found to not be broken:

Also, onFreshStack { function.block(value) } doesn't compile, instead requiring onFreshStack { function.block(this, value) }.

The type of function.block is a function that itself takes a receiver, which gets implicitly bound to this.

Sequence is missing the fun modifier for its interface. There's also a copy-paste error in the implementation for sequence, it should actually be = Sequence { iterator(block) }.

There's no such modifier. Rather, the standard library provides a function named Sequence. Perhaps that's also related the supposed copy-paste error?

There's also a little bit of strange formatting, like splitting extension function types with empty parameters (see e.g. finish), but that could be intentional.

Yeah, I was trying to keep things consistent while dealing with the fact that some signatures are really complex...

[kyay10]

new seems to create a StackSuspension out of nowhere. I'm guessing this is an "empty" StackSuspension? Is this related to switchTo somehow? My hypothesis is that StackMount can technically have multiple stacks that refer to it, but only one can be active at a single time (i.e. when the stack is mounted), and switchTo lets you swap between them without needing to dismount and mount again after.

Do you have any examples of why that can be useful? I don't think an analogy exists with other systems that I've seen.

[RossTate]

You are thinking a level too high. It just makes the target stack become the currently executing stack, like setting the stack-pointer register in assembly.

[kyay10]

so is the suspension it gives you just the stack frames up to the nearest dynamically-enclosing StackMount?
Where is the target stack mounted, then? It can't be at the call to switchTo, since we also get a suspension for that.

[RossTate]

In a sense. When you do StackMount.new, you create a new stack, and all the StackRemounter methods are explicit about what becomes the current stack. (There are also the stacks that come from the runtime environment. For example, the stack the program is started on is a stack, though it doesn't have an associated stack mount.)

The target stack's lifetime only satisfies the requirement for switchTo if all of the stack mounts it transitively uses are mounted.

[kyay10]

Ohhhh, right; apologies! So the mount is already mounted, hence the target stack can just become the current stack (this is in contrast to mount, which explicitly takes an Available mount).

all the StackRemounter methods are explicit about what becomes the current stack

Sure, but when you call restack, the current stack is just some dynamically-determined stack that you happen to be on. This feels like the only dynamic part in this otherwise lexically-determined system, so I'm curious why that's the case. In what situation would I want access to this dynamic stack suspension? I guess I might switchTo and then switch back?

(I think the way the API is designed, you never have to worry about the stack the program starts on. It doesn't have a StackMount, as you said, so dismount can't target it, and switchTo cannot target it either since, to call switchTo, you'd need to have a mounted StackMount that you're under, meaning that your current stack is definitely not the initial program stack)

[kyay10]

You are thinking a level too high. It just makes the target stack become the currently executing stack, like setting the stack-pointer register in assembly.

I'm thinking about it more, and I still don't get it. I don't see the difference between switchTo and a dismount followed by a mount.
I think I'm wrong about the dynamic part I said above. It seems like switchTo gives you a suspension in a manner similar to what dismount does, thus it's static.

[kyay10]

Does dismount also dismount all the mounts between the current stack and the target mount?
Maybe an example helps:

// convenience function to create a mount, create a new continuation from it, and resume it
fun newMount(block: (mount: StackMount<local, Unit>) ->_{local&mount.mounted} Unit): Unit = with(StackMount()) {
  resume(Unit, new(Unit)) { return block() }
}

newMount { m1 ->
  newMount { m2 ->
    m1.suspend(Unit) { _, _ ->
      m1.resume(Unit, new(Unit)) { // setting it up again so that the arena for `m2` is satisfied
        m2.resume(Unit, new(Unit)) { ... } // does this result in an exception?
      }
    }
  }
}

If true, I don't see how the dynamic checks for mount would work while keeping mount and dismount constant-time operations.
If false, then I'm somewhat struggling to understand the point of having multiple stacks for the same StackMount.

[RossTate]

It does not. But I don't see how that's related to having multiple stacks for one StackMount.

[kyay10]

StackMount has an out local arena lifetime parameter. I think it should actually be invariant. Because it's out, it means that one can simply not meet the arena requirements, and yet still call mount because StackMount can be upcast to have effectively no arena requirement.

[RossTate]

Oh wow, I'm not sure how I let that one slip past me. Thanks for that critical catch! I've fixed the write up and confirmed that the change doesn't affect any of the code I have using StackMount.

[kyay10]

fun <local arena, E, O> StackMount<arena, E>.new(
   after: (local E, O) ->_{arena} Nothing,
   block: () ->_{mounted} O
): StackContinuation<() -> Nothing>_{mounted}

StackMount is not given a lifetime here. Does that make it global? That seems unnecessarily restrictive. Same thing applies for the other utility stack functions

[RossTate]

An older design had a notion of limited-lifetime StackMount, but it didn't prove to be particularly useful, so I simplified and now treat StackMount as if all its objects have global lifetime.

[kyay10]

fun <local_{local} arena, E, R, T> StackMount<arena, E>.resume(
    environment: E_{local},
    continuation: StackContinuation<R>_{mounted},
    block: (local R) ->_{mounted} Nothing
): Nothing {
    restack {
        mount(this@new, continuation.suspension) {
            finish {
                block(continuation.resumer)
            }
        }
    }
}

I don't think this is valid. mount doesn't actually place you inside mount.mounted, so block cannot be called inside!
Maybe mount's that needs to have a mount.mounted upper bound?

[RossTate]

Oh shoot, I see what I did. There's an invariant that isn't captured in the design for StackContinuation. I've now fixed it. Thanks for catching this!

[rnett]

This is somewhat tangential to the proposal, but I've been thinking about ways to have pausable and resumable computations that survive process boundaries (i.e. can be serialized and persisted). This doesn't seem like it would help directly, although I'd probably want to use the "another kind of suspend" composability support and being able to inspect suspend calls for captures would be helpful for writing the compiler transforms (or doing JVM hacks). But I'm curious if you've thought about this at all or there's anything with the current proposal that I've missed that would help make it possible.

[RossTate]

It could help, but you'd still be responsible for the bulk of the work. In the StackContinuation class, I show how you can have a suspended stack with an associated "resumer" object that's valid to use within the suspension point's resumption locality. While typically this object will be used to resume the suspended stack, it can also be used to inspect the suspended stack. You can even inspect the suspended stack without switching to it: when you mount it, the restacker session then runs within the suspended stack's resumption locality without actually switching to the stack yet. That means that your "resumer" object could have a serialize method that you could invoke while inspecting the stack. That still allows the object to have a limited lifetime, and as such it could be allocated on the stack rather than the heap.

All that means that this design enables you to efficiently "decorate" the stack with functionality for serializing it on demand. But it's still on you to provide that decoration, which unfortunately would likely be a lot of boilerplate.