Wednesday, January 1, 2025

Calling Conventions in the Interpreter

C# is not tail recursive. It could be. The IL that it compiles to supports tail recursion, but the C# compiler doesn’t generate the tail call instruction. It would be a simple thing to add: when the compiler emits a call instruction, it could check if the next instruction is a return, and if so, emit a tail call instruction. This could be controlled by a compiler flag so only us weirdos who want this feature would enable it.

But until the C# compiler adds this feature, we have to resort to other methods. I chose to use a trampoline at each call site. This is a small segment of code that awaits the result of the function call. If the callee wants to tail call, it returns the tail call target to the caller, which performs the call on the callee’s behalf. This requires a modification to the calling conventions.

EvalStep is the virtual method that all S-code objects implement to perform an evaluation. Its signature is this:

abstract class Control : SchemeObject
{
     public abstract TailRecursionFlag EvalStep (out object answer, 
                                                 ref Control expression, 
                                                 ref Environment environment);
}

The result of the evaluation is returned in the answer parameter. This is an out parameter, so the answer is allocated in the caller and a pointer to it is passed to the callee. The callee returns the answer by modifying it in the callers stack frame.

The expression and environment parameters are the expected parameters for a call to Eval. They, too, are allocated in the caller’s frame and references to them are passed to the callee. The callee is allowed to modify the caller’s values of these variables.

The returned value is a TailRecursionFlag. This is either 1, indicating that a value has been returned in the answer, or 0, indicating that the caller should perform another EvalStep. To return a value, the callee modifies the answer. To perform a tail call, the callee modifies the expression and environment references and returns 0.

Any caller must call EvalStep as follows: The caller allocates an answer variable to receive the answer of the call. It also allocates an expression, and environment variable to pass to the callee. It then calls EvalStep in a loop until the callee returns a TailRecursionFlag of 1, indicating that the answer has been set to the return value.

In the EvalStep for an S-code Conditional we see an example of the calling convention:

  object ev;
  Control unev = predicate;
  Environment env = environment;

  while (unev.EvalStep (out ev, ref unev, ref env) == TailRecursionFlag.TailCall) { };

We are making a recursive call to evaluate the predicate. We set up ev to receive the result of the evaluation. We set up unev and env to hold the expression and environment to pass to EvalStep. unev.EvalStep does the eval dispatch via virtual function dispatch.

If the predicate returns a TailRecursionFlag of ReturnValue, the loop will exit. The predicate is assumed to have put the return value in the ev variable.

If the predicate wants to tail call, it will modify the values of unev and env to the new expression and new environment, and return a TailRecursionFlag of TailCall. The loop will iterate, using the new value of unev and env to again dispatch a call to EvalStep.

When the while loop exits, the ev variable will contain the return value of the predicate. Control may be returned to the while loop several times before the loop exits. This is the trampoline in action.

Conditional expressions don’t return a value. They either tail call the consequent or the alternative. The EvalStep for a conditional ends like this:

  answer = null;
  expression = (ev is bool evb && evb == false) ? alternative :
  return TailRecursionFlag.TailCall;
}

The answer variable in the caller is set to null. out parameters must always be assigned to before the function exits, so this just keeps the compiler happy. If the return value of calling EvalStep on the predicate is the boolean false, we set the expression in the caller to the alternative, otherwise the consequent. This is the target of our tail call to EvalStep. For the scode for a conditional, we leave the environment alone — the tail call uses the same environment unchanged. We finally return TailRecursionFlag.TailCall so that the caller’s trampoline makes another iteration around its while. It will call EvalStep on the alternative or consequent that we stuffed in the caller’s expression.

This little song and dance is performed at every recursive call to EvalStep making EvalStep behave as a tail-recursive function. This calling convention is about half the speed of a normal C# method call. It is the cost of using a trampoline for tail recursion.

First Class Continuations

There is one more obscure reason that the control might return to us when evaluating the predicate. If some function further down the call chain invokes call-with-current-continuation, we need to copy the stack. The callee indicates this by returning a magic return value of Special.UnwindStack. The callee sets the caller’s environment to an UnwinderState that will accumulate the stack frames as we unwind the stack. So our calling convention says we need to check the return value of EvalStep, and if it is Special.UnwindStack, we allocate a ConditionalFrame on the heap that will contain the state of the current stack frame. We AddFrame to the UnwinderState. We propagate this up the stack by putting it in the caller’s environment, setting the caller’s value of answer to Special.UnwindStack and returning TailRecursionFlag.ReturnValue to stop the caller’s trampoline loop.

The full code of EvalStep for an S-code if expression is this:

 public override TailRecursionFlag EvalStep (out object answer, 
                                             ref Control expression,
                                             ref Environment environment)
{
    object ev;
    Control unev = predicate;
    Environment env = environment;

    // Tail recursion trampoline.
    while (unev.EvalStep (out ev, ref unev, ref env) == TailRecursionFlag.TailCall) { };
    // Support for first class continuations.
    if (ev == Special.UnwindStack)
    {
        ((UnwinderState) env).AddFrame (new ConditionalFrame (this, environment));
        environment = env;
        answer = Special.UnwindStack;

        return TailRecursionFlag.ReturnValue;
    }

    // Tail call EvalStep on the consequent or alternative.
    answer = null;
    expression = (ev is bool evb && evb == false) ? alternative : consequent;
    return TailRecursionFlag.TailCall;
}

First class continuations allow you unload and reload the pending call chain. We see that at each call site, we must check the return value and, if it is Special.UnwindStack, we create a new Frame on the heap and add it to the unwinder state befor we propagate the Special.UnwindStack up the call chain.

At the very top of the call chain, we have the outermost call to EvalStep. If the Special.UnwindStack value is returned to this call, the stack has been unwound and the UnwinderState is sitting in the environment variable. We need to rewind the stack and put the stack frames back on the stack. We create a RewindState from the UnwinderState. Each time we PopFrame from the RewindState, we get a deeper frame. We reload the stack by getting the outermost frame from the RewindState and calling EvalStep on it. The EvalStep for a Frame sets up the trampoline loop, calls PopFrame to get the next frame, and calls EvalStep on it. When we run out of stack frames to reload, the stack is reloaded and we return control the innermost frame so it can continue where it left off. This is the rewind loop.

The EvalStep for a Frame, after making the recursive call to EvalStep on the next frame, continues with code that is a duplicate of the code in the original frame before the cotinuation was captured. A specific example will make this clear. If an if expression is on the stack when it is uwound, a ConditionalFrame is created. A ConditionalFrame is a subclass of SubproblemFrame which has this EvalStep method:

public override TailRecursionFlag EvalStep (out object answer,
                                            ref Control expression,
                                            ref Environment environment)
{
    object temp;
    Control expr = ((RewindState) environment).PopFrame ();
    Environment env = environment;
    while (expr.EvalStep (out temp, ref expr, ref env) == TailRecursionFlag.TailCall) { };
    if (temp == Special.UnwindStack)
    {
        ((UnwinderState) env).AppendContinuationFrames (continuation);
        environment = env;
        answer = Special.UnwindStack;

        return TailRecursionFlag.ReturnValue;
    }
    expression = this.expression;
    environment = this.environment;
    return Continue (out answer, ref expression, ref environment, temp);
}

public abstract TailRecursionFlag Continue (out object answer,
                                            ref Control expression,
                                            ref Environment environment,
                                            object value);

That is, the EvalStep of the SubproblemFrame establishes a trampoline, pops the next frame from the RewindState, and invokes its EvalStep method. When an answer is returned, the SubproblemFrame calls its Continue method.

The Continue method is a virtual method that is implemented by each subclass of SubproblemFrame. It finishes the work of the frame. In the case of a ConditionalFrame, the Continue method is this:

public override TailRecursionFlag Continue (out object answer,
                                            ref Control expression,
                                            ref Environment environment,
                                            object value)
{
    answer = null;
    expression = value is bool bvalue && bvalue == false
      ? SCode.EnsureSCode (this.expression.Alternative)
      : SCode.EnsureSCode (this.expression.Consequent);
    return TailRecursionFlag.TailCall;
}
compare this to the code in the original Conditional:
    // Tail call EvalStep on the consequent or alternative.
    answer = null;
    expression = (ev is bool evb && evb == false) ? alternative : consequent;
    return TailRecursionFlag.TailCall;

There are only superficial differences: the Continue method gets the value returned by the predicate in an argument rather than in a local variable. It type checks the alternative and consequent components of the if expression by calling SCode.EnsureSCode. Otherwise, the code does the same thing.

It is not possible to actually rewind the stack with the original set of pending methods. What we do instead is rewind the stack with methods that do the same thing as the original pending methods. It is close enough. The same values will be computed.

There is one place you can see the difference. If you look at the stack trace in the debugger before you capture a continuation, you will see the pending recursive calls to the S-code EvalStep methods. If you look at the stack trace in the debugger after you capture a continuation, you will instead see pending calls to the EvalStep methods of a set of frames. The pending frames are in the same order and have names similar to the original pending methods. They compute the same values, too. But the debugger can notice that these are not the originals.

No comments: