Implement memoization

Implement `memo(...)` function which caches its argument on first evaluation
and re-uses the cached value afterwards. The cache is placed next to the
method enclosing the memo(...) call.

`memo` is a member of package `compiletime`.

Author: odersky

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -237,6 +237,7 @@ class Definitions {
     @threadUnsafe lazy val Compiletime_constValue   : SymbolPerRun = perRunSym(CompiletimePackageObject.requiredMethodRef("constValue"))
     @threadUnsafe lazy val Compiletime_constValueOpt: SymbolPerRun = perRunSym(CompiletimePackageObject.requiredMethodRef("constValueOpt"))
     @threadUnsafe lazy val Compiletime_code         : SymbolPerRun = perRunSym(CompiletimePackageObject.requiredMethodRef("code"))
+    @threadUnsafe lazy val Compiletime_memo         : SymbolPerRun = perRunSym(CompiletimePackageObject.requiredMethodRef("memo"))
 
   /** The `scalaShadowing` package is used to safely modify classes and
    *  objects in scala so that they can be used from dotty. They will

compiler/src/dotty/tools/dotc/core/NameKinds.scala

@@ -213,6 +213,9 @@ object NameKinds {
       safePrefix + info.num
     }
 
+    def currentCount(prefix: TermName = EmptyTermName) given (ctx: Context): Int =
+      ctx.freshNames.currentCount(prefix, this)
+
     /** Generate fresh unique term name of this kind with given prefix name */
     def fresh(prefix: TermName = EmptyTermName)(implicit ctx: Context): TermName =
       ctx.freshNames.newName(prefix, this)
@@ -296,6 +299,7 @@ object NameKinds {
   val UniqueInlineName: UniqueNameKind       = new UniqueNameKind("$i")
   val InlineScrutineeName: UniqueNameKind    = new UniqueNameKind("$scrutinee")
   val InlineBinderName: UniqueNameKind       = new UniqueNameKind("$elem")
+  val MemoCacheName: UniqueNameKind          = new UniqueNameKind("memo$")
 
   /** A kind of unique extension methods; Unlike other unique names, these can be
    *  unmangled.

compiler/src/dotty/tools/dotc/typer/Inliner.scala

@@ -15,7 +15,7 @@ import StdNames._
 import transform.SymUtils._
 import Contexts.Context
 import Names.{Name, TermName}
-import NameKinds.{InlineAccessorName, InlineBinderName, InlineScrutineeName}
+import NameKinds.{InlineAccessorName, InlineBinderName, InlineScrutineeName, MemoCacheName}
 import ProtoTypes.selectionProto
 import SymDenotations.SymDenotation
 import Inferencing.fullyDefinedType
@@ -188,6 +188,19 @@ object Inliner {
     if (callSym.is(Macro)) ref(callSym.topLevelClass.owner).select(callSym.topLevelClass.name).withSpan(pos.span)
     else Ident(callSym.topLevelClass.typeRef).withSpan(pos.span)
   }
+
+  /** For every occurrence of a memo cache symbol `memo$N` of type `T_N` in `tree`,
+   *  an assignment `val memo$N: T_N = null`
+   */
+  def memoCacheDefs(tree: Tree) given Context: Set[ValDef] = {
+    val memoCacheSyms = tree.deepFold[Set[TermSymbol]](Set.empty) {
+      (syms, t) => t match {
+        case Assign(lhs, _) if lhs.symbol.name.is(MemoCacheName) => syms + lhs.symbol.asTerm
+        case _ => syms
+      }
+    }
+    memoCacheSyms.map(ValDef(_, Literal(Constant(null))))
+  }
 }
 
 /** Produces an inlined version of `call` via its `inlined` method.
@@ -392,6 +405,36 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
       case _ => EmptyTree
     }
 
+  /** The expansion of `memo(op)` where `op: T` is:
+   *
+   *    { if (memo$N == null) memo$N = op; $memo.asInstanceOf[T] }
+   *
+   *  This creates as a side effect a memo cache symbol $memo$N` of type `T | Null`.
+   *  TODO: Restrict this to non-null types, once nullability checking is in.
+   */
+  def memoized: Tree = {
+    val currentOwner = ctx.owner.skipWeakOwner
+    if (currentOwner.isRealMethod) {
+      val cacheOwner = ctx.owner.effectiveOwner
+      val argType = callTypeArgs.head.tpe
+      val memoVar = ctx.newSymbol(
+        owner = cacheOwner,
+        name = MemoCacheName.fresh(),
+        flags =
+          if (cacheOwner.isTerm) Synthetic | Mutable
+          else Synthetic | Mutable | Private | Local,
+        info = OrType(argType, defn.NullType),
+        coord = call.span)
+      val cond = If(
+        ref(memoVar).select(defn.Any_==).appliedTo(Literal(Constant(null))),
+        ref(memoVar).becomes(callValueArgss.head.head),
+        Literal(Constant(())))
+      val expr = ref(memoVar).cast(argType)
+      Block(cond :: Nil, expr)
+    }
+    else errorTree(call, em"""memo(...) outside method""")
+  }
+
   /** The Inlined node representing the inlined call */
   def inlined(sourcePos: SourcePosition): Tree = {
 
@@ -408,6 +451,8 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
           else New(defn.SomeClass.typeRef.appliedTo(constVal.tpe), constVal :: Nil)
         )
       }
+      else if (inlinedMethod == defn.Compiletime_memo)
+        return memoized
 
     // Compute bindings for all parameters, appending them to bindingsBuf
     computeParamBindings(inlinedMethod.info, callTypeArgs, callValueArgss)

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -2121,6 +2121,7 @@ class Typer extends Namer
           case Some(xtree) =>
             traverse(xtree :: rest)
           case none =>
+            val memoCacheCount = MemoCacheName.currentCount()
             typed(mdef) match {
               case mdef1: DefDef if Inliner.hasBodyToInline(mdef1.symbol) =>
                 buf += inlineExpansion(mdef1)
@@ -2131,6 +2132,8 @@ class Typer extends Namer
                 mdef match {
                   case mdef: untpd.TypeDef if mdef.mods.isEnumClass =>
                     enumContexts(mdef1.symbol) = ctx
+                  case _: untpd.DefDef if MemoCacheName.currentCount() != memoCacheCount =>
+                    buf ++= Inliner.memoCacheDefs(mdef1)
                   case _ =>
                 }
                 if (!mdef1.isEmpty) // clashing synthetic case methods are converted to empty trees

compiler/src/dotty/tools/dotc/util/FreshNameCreator.scala

@@ -9,20 +9,27 @@ import core.StdNames.str
 
 abstract class FreshNameCreator {
   def newName(prefix: TermName, unique: UniqueNameKind): TermName
+  def currentCount(prefix: TermName, unique: UniqueNameKind): Int
 }
 
 object FreshNameCreator {
   class Default extends FreshNameCreator {
-    protected var counter: Int = 0
     protected val counters: mutable.Map[String, Int] = mutable.AnyRefMap() withDefaultValue 0
 
+    private def keyFor(prefix: TermName, unique: UniqueNameKind) =
+      str.sanitize(prefix.toString) + unique.separator
+
+    /** The current counter for the given combination of `prefix` and `unique` */
+    def currentCount(prefix: TermName, unique: UniqueNameKind): Int =
+      counters(keyFor(prefix, unique))
+
     /**
      * Create a fresh name with the given prefix. It is guaranteed
      * that the returned name has never been returned by a previous
      * call to this function (provided the prefix does not end in a digit).
      */
     def newName(prefix: TermName, unique: UniqueNameKind): TermName = {
-      val key = str.sanitize(prefix.toString) + unique.separator
+      val key = keyFor(prefix, unique)
       counters(key) += 1
       prefix.derived(unique.NumberedInfo(counters(key)))
     }

library/src/scala/compiletime/package.scala

@@ -38,4 +38,6 @@ package object compiletime {
   inline def constValue[T]: T = ???
 
   type S[X <: Int] <: Int
+
+  inline def memo[T](op: => T): T = ???
 }

tests/run/memoTest.scala

@@ -1,4 +1,5 @@
 object Test extends App {
+  import compiletime.memo
 
   var opCache: Int | Null = null
 
@@ -7,6 +8,8 @@ object Test extends App {
     opCache.asInstanceOf[Int] + 1
   }
 
-  assert(foo(1) + foo(2) == 4)
+  def bar(x: Int) = memo(x * x) + 1
 
+  assert(foo(1) + foo(2) == 4)
+  assert(bar(1) + bar(2) == 4)
 }