Tail-recursion optimization and @tailrec annotation fails in important cases

[LPTK]

The hkmc2/shared/src/test/mlscript-compile/apps/parsing/Lexer.mls file currently has this defintion:

  fun scan(idx: Int, acc: Stack[TokenType]): Stack[TokenType] =
    fun go(idx: Int, tok: TokenType) = if
      options.noWhitespace and tok is
        Token.Comment then scan(idx, acc)
        Token.Space then scan(idx, acc)
      else scan(idx, tok :: acc)
    
    if char(idx) is
      None then reverse of acc
      Some(ch) and ch is...
      Char.Whitespace and whitespace(idx) is idx' then
        go(idx', Token.space(idx, idx'))
      "\"" then go(...string(idx + 1))
      Bracket as b then go(idx + 1, Token.symbol(b, idx))
      "/" then go(...comment(idx + 1))
      Operator as ch then go(...operator(idx + 1, ch))
      Char.Digit as ch then go(...number(idx + 1, ch))
      IdentifierStart as ch then go(...identifier(idx + 1, ch))
      IdentifierQuote as quote
        and char(idx + 1) is Some(IdentifierStart as ch)
        and identifier(idx + 2, quote + ch) is [idx', token] and
          token is Token.Identifier(name, _) then
            go(idx', Token.identifier(name, idx))
          else go(idx + 1, Token.error(idx, idx + 1))
      else
        print("Unrecognized character: '" ~ ch ~ "'")
        go(idx + 1, Token.error(idx, idx + 1))
  
  scan(0, Nil)

This should clearly be rewritten as a loop by tail-call optimization.

Currently, there are a couple of things preventing that:

• The functions are not lifted by default, and tail-call opt only deals with lifted functions.
  Either we should generalize the tail-call optimizer, or we should enable lifting by default, or we should always selectively lift functions that need tail-call optimization.
• The UCS tries to share common branches by introducing labelled block which move some calls to non-tail positions, a known issue that really needs fixing.
  It should be possible to share common branches while keeping tail-calls in tail position.

Compounded with that is the strange issue that annotating the scan and go functions above as @tailred does not report any error (defeating the point of the annotation). How come?

The result is that big-enough inputs break it, as can be seen by using slightly larger inputs in hkmc2/shared/src/test/mlscript/apps/parsing/CamlLightTest.mls.

---

From our discussion in the meeting, we want the tail-rec optimizer to:

• Handle spreads
• Handle not just purely top-level fucntions, but also static module-paths to function

[CAG2Mark]

we should generalize the tail-call optimizer

Being able to compile optimize a mutually recursive parent and inner function into a combined loop subsumes lifting in a sense, since we must necessarily move the two function bodies to the same scope.

[LPTK]

we should generalize the tail-call optimizer

Being able to compile optimize a mutually recursive parent and inner function into a combined loop subsumes lifting in a sense, since we must necessarily move the two function bodies to the same scope.

Good. So what do you think about the other two options? Maybe we could do "soft-lifting" by default, meaning lifting only (mutually) recursive definitions and leaving the rest as is (since it seems JS handles nested function fine)?

[CAG2Mark]

I don't know... "soft lifting" in principle it is fine, but there's the question of whether we should lift them all the way to the top level, and the lifter is only able to do this at the moment. Also, it could cause reduce performance in unintended ways, i.e.

fun f(x) =
  fun g() = x + 1
  fun h() = 
    g()
    f(0)
  h()

If we lift h only, then the lifter will pass a closure containing the function g and the local variable x to h, which could lead to extra memory allocations.

There's also the question of dealing with functions that are mutually recursive across module boundaries; for example,

module M with
  fun f() = g()
fun g() = M.f()

This case is currently not handled, although I do think I can implement it.

[LPTK]

Ah, good point about the potential loss of performance due to soft lifting.

It seems to me we should just lift everything by default, then. And yes, it would be awesome if you could support functions that are accessible through non empty module paths too, rather than just purely top level functions.

[AnsonYeung]

The original example also contain the spread argument which blocks the optimization from happening

[LPTK]

Ouch. Spread arguments really shouldn't be such a headache. We may want to just define some reusable abstraction to deal with them uniformly during IR passes/analyses.

@CAG2Mark do you think you could also handle that?

[CAG2Mark]

Sure. I do recall discussing how it was quite annoying, but I think I can deal with it.

Note that I will be a bit busy in the coming week as I have some things due, and I would like to finish the generalized switch specialization stuff first.