Make memo work in traits

... and also make it work on several nested levels together.

Author: odersky

compiler/src/dotty/tools/dotc/ast/Desugar.scala

@@ -144,21 +144,17 @@ object desugar {
 
 // ----- Desugar methods -------------------------------------------------
 
+  def setterNeeded(flags: FlagSet, owner: Symbol) given Context =
+    flags.is(Mutable) && owner.isClass && (!flags.isAllOf(PrivateLocal) || owner.is(Trait))
+
   /**   var x: Int = expr
    *  ==>
    *    def x: Int = expr
    *    def x_=($1: <TypeTree()>): Unit = ()
    */
   def valDef(vdef0: ValDef)(implicit ctx: Context): Tree = {
     val vdef @ ValDef(name, tpt, rhs) = transformQuotedPatternName(vdef0)
-    val mods = vdef.mods
-    val setterNeeded =
-      mods.is(Mutable) && ctx.owner.isClass && (!mods.isAllOf(PrivateLocal) || ctx.owner.is(Trait))
-    if (setterNeeded) {
-      // TODO: copy of vdef as getter needed?
-      // val getter = ValDef(mods, name, tpt, rhs) withPos vdef.pos?
-      // right now vdef maps via expandedTree to a thicket which concerns itself.
-      // I don't see a problem with that but if there is one we can avoid it by making a copy here.
+    if (setterNeeded(vdef.mods.flags, ctx.owner)) {
       val setterParam = makeSyntheticParameter(tpt = SetterParamTree().watching(vdef))
       // The rhs gets filled in later, when field is generated and getter has parameters (see Memoize miniphase)
       val setterRhs = if (vdef.rhs.isEmpty) EmptyTree else unitLiteral
@@ -168,7 +164,7 @@ object desugar {
         vparamss = (setterParam :: Nil) :: Nil,
         tpt      = TypeTree(defn.UnitType),
         rhs      = setterRhs
-      ).withMods((mods | Accessor) &~ (CaseAccessor | GivenOrImplicit | Lazy))
+      ).withMods((vdef.mods | Accessor) &~ (CaseAccessor | GivenOrImplicit | Lazy))
       Thicket(vdef, setter)
     }
     else vdef

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

@@ -299,7 +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$")
+  val MemoCacheName: UniqueNameKind          = new UniqueNameKind("$cache")
 
   /** A kind of unique extension methods; Unlike other unique names, these can be
    *  unmangled.

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

@@ -482,6 +482,7 @@ object StdNames {
     val materializeClassTag: N  = "materializeClassTag"
     val materializeWeakTypeTag: N = "materializeWeakTypeTag"
     val materializeTypeTag: N   = "materializeTypeTag"
+    val memo: N                 = "memo"
     val mirror : N              = "mirror"
     val moduleClass : N         = "moduleClass"
     val name: N                 = "name"

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

@@ -553,9 +553,10 @@ object Symbols {
 
     /** This symbol entered into owner's scope (owner must be a class). */
     final def entered(implicit ctx: Context): this.type = {
-      assert(this.owner.isClass, s"symbol ($this) entered the scope of non-class owner ${this.owner}") // !!! DEBUG
-      this.owner.asClass.enter(this)
-      if (this.is(Module)) this.owner.asClass.enter(this.moduleClass)
+      if (this.owner.isClass) {
+        this.owner.asClass.enter(this)
+        if (this.is(Module)) this.owner.asClass.enter(this.moduleClass)
+      }
       this
     }
 

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

@@ -193,14 +193,22 @@ object Inliner {
   /** 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
+  def memoCacheDefs(tree: Tree) given Context: List[ValOrDefDef] = {
+    object memoRefs extends TreeTraverser {
+      val syms = new mutable.LinkedHashSet[TermSymbol]
+      def traverse(tree: Tree) given Context = tree match {
+        case tree: RefTree if tree.symbol.name.is(MemoCacheName) =>
+          syms += tree.symbol.asTerm
+        case _: DefDef =>
+          // don't traverse deeper; nested memo caches go next to nested method
+        case _ =>
+          traverseChildren(tree)
       }
     }
-    memoCacheSyms.map(sym => ValDef(sym, Literal(Constant(null))).withSpan(sym.span))
+    memoRefs.traverse(tree)
+    for sym <- memoRefs.syms.toList yield
+      (if (sym.isSetter) DefDef(sym, _ => Literal(Constant(())))
+       else ValDef(sym, Literal(Constant(null)))).withSpan(sym.span)
   }
 }
 
@@ -420,19 +428,22 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
       val argType = callTypeArgs.head.tpe
       val memoVar = ctx.newSymbol(
         owner = cacheOwner,
-        name = MemoCacheName.fresh(),
+        name = MemoCacheName.fresh(nme.memo),
         flags =
           if (cacheOwner.isTerm) Synthetic | Mutable
           else Synthetic | Mutable | Private | Local,
         info = OrType(argType, defn.NullType),
-        coord = call.span)
-      val memoSetter = ctx.newSymbol(
-        owner = cacheOwner,
-        name = memoVar.name.setterName,
-        flags = memoVar.flags | Method | Accessor,
-        info = MethodType(argType :: Nil, defn.UnitType),
-        coord = call.span
-      )
+        coord = call.span).entered
+      val memoSetter =
+        if (desugar.setterNeeded(memoVar.flags, cacheOwner))
+          ctx.newSymbol(
+            owner = cacheOwner,
+            name = memoVar.name.setterName,
+            flags = memoVar.flags | Method | Accessor,
+            info = MethodType(argType :: Nil, defn.UnitType),
+            coord = call.span
+          ).entered
+        else memoVar
       val memoRef = ref(memoVar).withSpan(call.span)
       val cond = If(
         memoRef.select(defn.Any_==).appliedTo(Literal(Constant(null))),

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

@@ -2122,19 +2122,18 @@ class Typer extends Namer
           case Some(xtree) =>
             traverse(xtree :: rest)
           case none =>
-            val memoCacheCount = MemoCacheName.currentCount()
+            val memoCacheCount = MemoCacheName.currentCount(nme.memo)
             typed(mdef) match {
               case mdef1: DefDef if Inliner.hasBodyToInline(mdef1.symbol) =>
                 buf += inlineExpansion(mdef1)
                   // replace body with expansion, because it will be used as inlined body
                   // from separately compiled files - the original BodyAnnotation is not kept.
               case mdef1 =>
-                import untpd.modsDeco
-                mdef match {
-                  case mdef: untpd.TypeDef if mdef.mods.isEnumClass =>
+                mdef1 match {
+                  case mdef1: TypeDef if mdef1.symbol.flags.is(Enum, butNot = Case) =>
                     enumContexts(mdef1.symbol) = ctx
-                  case _: untpd.DefDef if MemoCacheName.currentCount() != memoCacheCount =>
-                    buf ++= Inliner.memoCacheDefs(mdef1)
+                  case mdef1: DefDef if MemoCacheName.currentCount(nme.memo) != memoCacheCount =>
+                    buf ++= Inliner.memoCacheDefs(mdef1.rhs)
                   case _ =>
                 }
                 if (!mdef1.isEmpty) // clashing synthetic case methods are converted to empty trees

tests/run/memoTest.check

@@ -0,0 +1,7 @@
+computing inner
+computing f
+1
+1
+computing f
+1
+1

tests/run/memoTest.scala

@@ -13,6 +13,22 @@ object Test extends App {
   assert(foo(1) + foo(2) == 4)
   assert(bar(1) + bar(2) == 4)
 
+  trait T {
+    def x: Int
+    def y: Int = memo {
+      def inner = memo {
+        println("computing inner");
+        x * x
+      }
+      inner + inner
+    }
+  }
+  val t = new T {
+    def x = 3
+    assert(y == 18)
+  }
+  assert(t.y == 18)
+
   class Context(val n: Int)
   def f(c: Context): Context = {
     println("computing f")