Inlining

We can return to our previous simple example that prints all natural numbers.

// broken, throws OOM
def print_naturals =
  val naturals = LazyList.from(0)
  naturals.foreach(println)

// fixed, constant memory use
def print_naturals =
  LazyList.from(0).foreach(println)

On the right side, the LazyList is not stored as a local variable, so there is no extra reference to the head of the list once we start traversing it, making it ready for garbage collection. This approach works in simple cases, but it is not always possible (or desirable) to inline everything.

Explicit nullification

In some cases, it is possible to modify a mutable variable to null to release the reference.

def trainModel =
  var data = loadDataset()
  val stats = statistics(data)
  val model = fitModel(data, stats)
  data = null // drop the reference
  evaluate(model)

Tail-recursion

The tail-recursive functions case is interesting. These methods receive a specific backend optimization, so that the method parameters are reused instead of being aliased to a mutable variable, leaving no extra reference to the original parameters.

Before the optimization, tail recursive methods are written as a while loop. However, this needs to create a new reference in order to modify it in-place. The original reference remains. The backend has been modified to directly loop on the method parameter and overwrite it (The Scala code would be illegal, method parameters are immutable!).

//original tailrec rewriting
def func(x: T) =
  var x_ = x
  while cond do
    x_ = step(x_)

//optimized: reusing `x`
def func(x: T) =
  while cond do
    x = step(x)

This optimization allows tail-recursive function to avoid keeping an extra reference to the variable, so that it can be dropped by the GC when no other handle remains.

Other cases

It is sometimes more complex to remove a reference. For example, method parameters are immutable so they cannot be nullified. In the code below, we did not find a way to remove the reference to l before traversing it.

//broken
def printEven(l: LazyList[Int]) =
  l
  .filter(isEven)
  .foreach(println)

In other cases where the nullification needs cannot be placed at usual program points, we can also use some tricks. For example, if we want to nullify the variable after passing it to a variable, we can do the following:

//nullify `nat` after passing it to `foreach`
def print_naturals =
  var nat = LazyList.from(0)
  (nat, {nat = null})._1.foreach(println)

By creating a pair, each argument are evalated in order and kept in the operand stack. The second element of the pair allows to nullify l at the right time. Naturally, Lisp fans will have recognized prog1 (prog1 x (setq x nil)).

From the Scala standard library: disallowing further uses

In the implementation of LazyList from the standard library, a comment is trying to disallow the reuse of a variable. However, it is possible to overlook the comment and cause trouble.

private def filterImpl[A](ll: LazyList[A], p: A => Boolean, isFlipped: Boolean): LazyList[A] = {
  // DO NOT REFERENCE `ll` ANYWHERE ELSE, OR IT WILL LEAK THE HEAD
  var restRef = ll
  newLL { ...

A @lastUse annotation could easily enforce the comment, making it safer.

var restRef = ll: @lastUse

This makes any use of ll illegal after this line. An error will be thrown if it is not respected.

Nullifying immutable method parameters

When a method parameter needs to be nullified, our annotation comes to the rescue. Even though ll is immutable, adding an annotation removes its reference before giving up control to foreach.

def printEven(ll: LazyList[Int]): Unit =
  (ll: @lastUse).filter(isEven).foreach(println)

@main def main =
  printEven(LazyList.from(0))

Nullifying at a fine-grained scale

The introduction problem can easily be solved. Here, the reference to sine will be removed just before entering play. Hence, play contains the only reference to the list. Once it starts traversing it, the beginning of the list will get collected.

@main def main = {
   val sine = loop(sinePulse(440))
   play(sine: @lastUse)

Aliasing to a mutable var

In other cases, it is possible to ensure better memory usage when aliasing:

def simulate(grid: Grid): Grid =
  var grid_ = grid: @lastUse //removes original reference

  while !converged(grid_) do
    grid_ = step(grid_)
  grid_

As grid_ evolves, the original grid reference is not pinned anymore, and can be collected if no one else is keeping it. note that the JIT compiler can easily remove this reference when executing the loop, but only outside of the step function. Once inside of step, the JVM does not perform optimizations on the simulate method, so the reference to grid remains.

overview of the design

We want the compiler to remove the reference contained in a local variable, once it is annotated with @lastUse. For this, we need to add a new type annotation to the standard Scala library. Then, the backend of the compiler is modified to catch the annotation and nullify the reference right after loading it.

But this alone will not work: The Erasure phase, in the middle of the transform phases, will remove our annotation. To allow our information to remain until the last phase (the backend), we first need to preserve the annotation as a Scala attachment. These structures can be attached on any Tree and carry any information. This process of capturing annotations and transforming them into attachments is performed in the PostTyper phase. We also take advantage of catching all the annotated variables to store them in a Set for easier access in the future.

Finally, we can make the use of the @lastUse annotation safer by verifying that it is indeed not used later. This phase is placed early in the transform phases, to better match the user’s code. The verifying process runs a simple control flow analysis, taking care of most illegal uses of the annotation.

PostTyper: catching annotations, switch to attachments

override def transform(tree: Tree)(using Context): Tree =
  try tree match {
    [...]
    case Typed(t, tpt: TypeTree) if tpt.tpe.hasAnnotation(defn.LastUseAnnot) =>
      t match
        case _: Ident =>
          t.putAttachment(PostTyper.lastUseAttachment, ())
          val enclosing = ctx.property(enclosingMethod).get
          val others = enclosing.getAttachment(methodLastUses).getOrElse(Set.empty)
          //attach the symbol to the method
          enclosing.putAttachment(methodLastUses, others + t.symbol)
          Typed(t, tpt)
        case _ =>
          report.error("`@lastUse` annotation can only be applied on local variables", tree.srcPos)
          Typed(t, tpt)

In this phase, we capture all @lastUse annotations on local variables, replace them by an attachment on the tree and store the annotated symbol in the methodLastUses attached set. This set is stored as an attachment to the enclosing method. This will allow for easy access and efficient verification.

Backend: bytecode modification

def genLoadTo(tree: Tree, expectedType: BType, dest: LoadDestination): Unit =
  tree match
    [...]
    case t @ Ident(name) =>
      [...]
      locals.load(sym) // local variable loading
      if t.hasAttachment(lastUseAttachment) then
        emit(asm.Opcodes.ACONST_NULL)
        val idx = locals.getOrMakeLocal(sym).idx
        bc.store(idx, tk)

Here we can check for the attachment when loading a local symbol. If the attachment exists, we add the following lines in the final bytecode:

1: aconst_null
2: astore_<n>

This loads the value null in the operand stack, and overwrites our variable with it. This effectively removes the reference.

However, the variable could still be used, which would potentially cause an error. To ensure a safe operation, a verification phase is necessary. We describe it in the following part.

VerifyLastUseAnnotations: checking validity with control flow analysis

/** Phases dealing with the transformation from pickled trees to backend trees */
 protected def transformPhases: List[List[Phase]] =
   List(new InstrumentCoverage) ::
   List(new CrossVersionChecks,
        new FirstTransform,
        new VerifyLastUseAnnotations, // added miniphase
        new CheckReentrant,
        new ElimPackagePrefixes,
        [...]

The phase is placed early in the transformPhases. Hence, it can throw an error as early as possible, with the code being relatively close to the source code.

Each annotated symbol will be tracked, for its uses and its annotation, to ensure a valid use of the annotation. In this section, we will call them "lastUse variables".

First, a few types and helper methods are created.

case class VariableStatus(allowVar: Boolean, allowAnnot: Boolean, isAfterValDef: Boolean){
  def afterAnnot = VariableStatus(false, allowAnnot, isAfterValDef)
  def afterValDef = VariableStatus(allowVar, allowAnnot, true)
  def enterLoop = VariableStatus(allowVar, !isAfterValDef, isAfterValDef)
  def merge(other: VariableStatus) = VariableStatus(allowVar && other.allowVar, allowAnnot && other.allowAnnot, isAfterValDef ||  other.isAfterValDef)
}

type VariableStatusMap = Map[Symbol, VariableStatus]

extension(ss: VariableStatusMap)
  def merge(other: VariableStatusMap): VariableStatusMap =
    ss.map((sym, state) => (sym -> state.merge(other(sym))))

  def enterLoop: VariableStatusMap =
    ss.map((sym, state) => (sym -> state.enterLoop))

VariableStatus holds the verification state for a single variable. merge is useful to connect exclusive branches, such as if-then-else branches. enterLoop is useful to disallow annotating a variable in a loop (if the variable is created before the loop)

override def prepareForDefDef(tree: DefDef)(using Context): Context =
  val freeVars = ctx.property(LastUseVariables).getOrElse(Set.empty[Symbol]) ++ tree.getAttachment(methodLastUses).getOrElse(Set.empty[Symbol])
  ctx.withProperty(LastUseVariables, Some(freeVars))

When recursively entering methods (DefDef), we build an environment to know of potential variables to verify from enclosing methods. Using this, we can efficiently detect if a nested method is using a variable.

override def transformDefDef(mdef: DefDef)(using Context): Tree =
  LastUseVerifier(mdef).verify
  mdef

The phase traverses the children trees first. So in our case, the most deeply nested method will be verified first. While verifying, We can mark any use of a lastUse variable coming from an enclosing method, and mark the current method as using it. Next, we call the verify method of the LastUseVerifier class, which is defined below.

private class LastUseVerifier(method: DefDef) extends TreeAccumulator[VariableStatusMap]{
  def verify(using Context): Unit =
    method.getAttachment(methodLastUses) match
      case None if ctx.property(LastUseVariables).get.nonEmpty =>
        // enclosing method has lastUse variables => check if it uses any
        traverse(method.rhs)(Map.empty)
      case None => ()
      case Some(symbols) =>
        val ss = symbols.map(s => (s -> VariableStatus.start)).toMap
        traverse(method.rhs)(ss)

The verifier is implemented as a TreeAccumulator, which allows for greater flexibility in the data flow. Unfortunately, it cannot be run in parallel with other miniphases.

verify is the entry method and performs some filtering and initialisation before traversing the method for verification:

• if the method has no lastUse parameters but is nested in a method which does, traverse and note any use of these captured variables

• if the method has lastUsed symbols, verify them

• else: do not enter the method. It cannot contain any error.

Below, we describe different interesting cases of the traverse method, the main recursive verification method:

private def traverse(tree: Tree)(accInfo: VariableStatusMap)(using Context): VariableStatusMap = {
  var info = accInfo//used to mutate and pass the information
  val sym = tree.symbol
  tree match
    case vd: ValDef if info.contains(sym) => // definition of a lastUse variable
      val info1 = info(sym)
      info = info.updated(sym, info1.afterValDef)
      foldOver(info, vd)

We need to catch when a lastUse variable is defined, to know if it was created in a loop or before.

case _: Ident if info.contains(sym) => //lastUse variable in current method
  val info1 = info(sym)
  if !info1.allowVar then
    report.error(s"reuse of variable ${tree.symbol.name} after @lastUse", tree.sourcePos) //reuse after annot
  if tree.hasAttachment(PostTyper.lastUseAttachment) then
    checkVarValidity(tree)
    if info1.allowAnnot then info.updated(sym, info1.afterAnnot) else //annot here
      report.error(s"cannot use @lastUse annotation in a loop", tree.sourcePos) //annot here but illegal
      tree.removeAttachment(PostTyper.lastUseAttachment)
      info
  else
    info

This catches a lastUse variable, which was created in the scope of the current method.

• if it was annotated before (!allowVar): report an error

• if it has the attachment:

  • and the attachment is legal: update the symbol info

  • and the attachment is illegal (currently in a loop: !allowAnnot): report an error

case _: Ident if ctx.property(LastUseVariables).get.contains(sym) =>//lastUse variable from an enclosing method
  val others = capturedVariables.getOrElse(method.symbol, Set.empty[Symbol])
  capturedVariables.update(method.symbol, others + sym)
  info

This catches a lastUse variable, which was created in an enclosing method. We simply update the capturedVariables table to reflect this. As children are traversed first, the table will be ready when the enclosing method is verified.

case _: Ident if capturedVariables.contains(sym) =>// calling a nested method which uses a lastUse variable
  val freeVars = capturedVariables.apply(sym)
  freeVars.foreach(fv =>
    info.get(fv) match
      case Some(value) => if !value.allowVar then report.error(s"calling function ${sym.name}, using variable ${fv.name} previously annotated with @lastUse", tree.sourcePos)
      case None => //lastUse variable is defined in an enclosing method. mark the current method as using it
        val others = capturedVariables.getOrElse(method.symbol, Set.empty[Symbol])
        capturedVariables.update(method.symbol, others + fv)
  )
  info

This catches a function call, which is using a lastUse variable.

• if the variable is local to the current function: report an error if after lastUse.

• else: the lastUse variable comes from an enclosing method, mark the current method as using it.

case If(cond, thenp, elsep) =>
  info = traverse(cond)(info)
  val info_branch1 = traverse(thenp)(info)//both branches are exclusive
  val info_branch2 = traverse(elsep)(info)//so we combine the info at the end
  info_branch1.merge(info_branch2)

case WhileDo(cond, body) =>
  info = info.enterLoop
  info = traverse(cond)(info)
  // cannot be annotated or valdef'd inside the loop => no information to propagate further.
  traverse(body)(info)

case Match(selector, cases) =>
  info = traverse(selector)(info)
  traverseCaseDefs(cases)(info)

case Try(expr, cases, finalizer) =>
  info = traverse(expr)(info)
  info = traverseCaseDefs(cases)(info)
  traverse(finalizer)(info)

case DefDef(name, paramss, tpt, preRhs) => info//nested method already checked, do not enter it

case _ => foldOver(info, tree)//visit other trees linearly

These cases catch different control flow scenarios:

• If branches: check both branches separately, then merge

• While: variables that were defined before the loop cannot be annotated inside

• Match, Try: CaseDefs are handled in a special way, to allow linear reading of the patterns and guards, but an exclusive reading of the bodies.

• DefDef: do not verify the child method. It has already been done

• Everything else is verified linearly, using the recursive foldOver method.

In summary, the VerifyLastUseAnnotations phase allows then to perform efficient (at almost no cost if there is no annotation) verification of @lastUse annotated variables.