Dual purpose inliner

The new inliner inlines both inline methods with fully-typed expansions
and transparent methods with partially untyped extensions.

Author: odersky

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

@@ -231,13 +231,125 @@ object Inliner {
           val inlineCtx = ctx
           inlined.updateAnnotation(LazyBodyAnnotation { _ =>
             implicit val ctx = inlineCtx
-            val body = treeExpr(ctx)
-            if (ctx.reporter.hasErrors) body else ctx.compilationUnit.inlineAccessors.makeInlineable(body)
+            val rawBody = treeExpr(ctx)
+            val typedBody =
+              if (ctx.reporter.hasErrors) rawBody
+              else ctx.compilationUnit.inlineAccessors.makeInlineable(rawBody)
+            val inlineableBody =
+              if (inlined.isInlinedMethod) typedBody
+              else addReferences(inlined, originalBody, typedBody)
+            inlining.println(i"Body to inline for $inlined: $inlineableBody")
+            inlineableBody
           })
         }
     }
   }
 
+  /** Tweak untyped tree `original` so that all external references are typed
+   *  and it reflects the changes in the corresponding typed tree `typed` that
+   *  make `typed` inlineable. Concretely:
+   *
+   *   - all external references via identifiers or this-references are converted
+   *     to typed splices,
+   *   - if X gets an inline accessor in `typed`, references to X in `original`
+   *     are converted to the inline accessor name.
+   */
+  def addReferences(inlineMethod: Symbol,
+      original: untpd.Tree, typed: tpd.Tree)(implicit ctx: Context): tpd.Tree = {
+
+    def isExternal(sym: Symbol) = sym.exists && !isLocal(sym, inlineMethod)
+
+    // Maps from positions to external reference types and inline selector names.
+    object referenced extends TreeTraverser {
+      val typeAtPos = mutable.Map[Position, Type]()
+      val nameAtPos = mutable.Map[Position, Name]()
+      val implicitSyms = mutable.Set[Symbol]()
+      val implicitRefs = new mutable.ListBuffer[Tree]
+      def registerIfContextualImplicit(tree: Tree) = tree match {
+        case tree: RefTree
+        if tree.removeAttachment(ContextualImplicit).isDefined &&
+           isExternal(tree.symbol) &&
+           !implicitSyms.contains(tree.symbol) =>
+          if (tree.existsSubTree(t => isLocal(tree.symbol, inlineMethod)))
+            ctx.warning("implicit reference $tree is dropped at inline site because it refers to local symbol(s)", tree.pos)
+          else {
+            implicitSyms += tree.symbol
+            implicitRefs += tree
+          }
+        case _ =>
+      }
+      def traverse(tree: Tree)(implicit ctx: Context): Unit = {
+        tree match {
+          case _: Ident | _: This =>
+            //println(i"leaf: $tree at ${tree.pos}")
+            if (isExternal(tree.symbol)) {
+              inlining.println(i"type at pos ${tree.pos.toSynthetic} = ${tree.tpe}")
+              typeAtPos(tree.pos.toSynthetic) = tree.tpe
+            }
+          case _: Select if tree.symbol.name.is(InlineAccessorName) =>
+            inlining.println(i"accessor: $tree at ${tree.pos}")
+            nameAtPos(tree.pos.toSynthetic) = tree.symbol.name
+          case _ =>
+        }
+        registerIfContextualImplicit(tree)
+        traverseChildren(tree)
+      }
+    }
+    referenced.traverse(typed)
+
+    // The untyped tree transform that applies the tweaks
+    object addRefs extends untpd.UntypedTreeMap {
+      override def transform(tree: untpd.Tree)(implicit ctx: Context): untpd.Tree = {
+        def adjustLeaf(tree: untpd.Tree): untpd.Tree = referenced.typeAtPos.get(tree.pos.toSynthetic) match {
+          case Some(tpe) => untpd.TypedSplice(tree.withType(tpe))
+          case none => tree
+        }
+        def adjustName(name: Name) = referenced.nameAtPos.get(tree.pos.toSynthetic) match {
+          case Some(n) => n
+          case none => name
+        }
+        def adjustQualifier(tree: untpd.Tree): untpd.Tree = tree match {
+          case tree @ Ident(name1) =>
+            referenced.typeAtPos.get(tree.pos.startPos) match {
+              case Some(tp: ThisType) =>
+                val qual = untpd.TypedSplice(This(tp.cls).withPos(tree.pos.startPos))
+                cpy.Select(tree)(qual, name1)
+              case none =>
+                tree
+            }
+          case tree => tree
+        }
+        val tree1 = super.transform(tree)
+        tree1 match {
+          case This(_) =>
+            adjustLeaf(tree1)
+          case Ident(name) =>
+            adjustQualifier(adjustLeaf(cpy.Ident(tree1)(adjustName(name))))
+          case Select(pre, name) =>
+            cpy.Select(tree1)(pre, adjustName(name))
+          case tree: untpd.DerivedTypeTree =>
+            inlining.println(i"inlining derived $tree --> ${ctx.typer.typed(tree)}")
+            untpd.TypedSplice(ctx.typer.typed(tree))
+          case _ =>
+            tree1
+        }
+      }
+    }
+    val implicitBindings =
+      for (iref <- referenced.implicitRefs.toList) yield {
+        val localImplicit = iref.symbol.asTerm.copy(
+          owner = inlineMethod,
+          name = UniqueName.fresh(iref.symbol.name.asTermName),
+          flags = Implicit | Method | Stable,
+          info = iref.symbol.info.ensureMethodic,
+          coord = inlineMethod.pos).asTerm
+        polyDefDef(localImplicit, tps => vrefss =>
+            iref.appliedToTypes(tps).appliedToArgss(vrefss))
+      }
+    val untpdSplice = tpd.UntypedSplice(addRefs.transform(original)).withType(typed.tpe)
+    seq(implicitBindings, untpdSplice)
+  }
+
   /** `sym` has an inline method with a known body to inline (note: definitions coming
    *  from Scala2x class files might be `@inline`, but still lack that body.
    */
@@ -499,17 +611,38 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
       case _ => tree
     }}
 
-    val inlineCtx = inlineContext(call)
+    val inlineTyper = if (meth.isTransparentMethod) new InlineTyper else new InlineReTyper
+    val inlineCtx = inlineContext(call).fresh.setTyper(inlineTyper).setNewScope
+
     // The complete translation maps references to `this` and parameters to
     // corresponding arguments or proxies on the type and term level. It also changes
     // the owner from the inlined method to the current owner.
     val inliner = new TreeTypeMap(typeMap, treeMap, meth :: Nil, ctx.owner :: Nil)(inlineCtx)
 
-    val expansion = inliner(rhsToInline.withPos(call.pos))
+    val expansion = inliner.transform(rhsToInline.withPos(call.pos)) match {
+      case Block(implicits, tpd.UntypedSplice(expansion)) =>
+        val prevOwners = implicits.map(_.symbol.owner).distinct
+        val localizer = new TreeTypeMap(oldOwners = prevOwners, newOwners = prevOwners.map(_ => ctx.owner))
+        val (_, implicits1) = localizer.transformDefs(implicits)
+        for (idef <- implicits1) {
+          bindingsBuf += idef.withType(idef.symbol.typeRef).asInstanceOf[ValOrDefDef]
+            // Note: Substituting new symbols does not automatically lead to good prefixes
+            // if the previous symbol was owned by a class. That's why we need to set the type
+            // of `idef` explicitly. It would be nice if substituters were smarter, but
+            // it seems non-trivial to come up with rules that work in
+          inlineCtx.enter(idef.symbol)
+        }
+        expansion
+      case tpd.UntypedSplice(expansion) =>
+        expansion
+      case expansion =>
+        expansion
+    }
+
     trace(i"inlining $call\n, BINDINGS =\n${bindingsBuf.toList}%\n%\nEXPANSION =\n$expansion", inlining, show = true) {
 
       // The final expansion runs a typing pass over the inlined tree. See InlineTyper for details.
-      val expansion1 = InlineTyper.typed(expansion, pt)(inlineCtx)
+      val expansion1 = inlineTyper.typed(expansion, pt)(inlineCtx)
 
       /** All bindings in `bindingsBuf` except bindings of inlineable closures */
       val bindings = bindingsBuf.toList.map(_.withPos(call.pos))
@@ -529,8 +662,8 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
    */
   private object InlineableArg {
     lazy val paramProxies = paramProxy.values.toSet
-    def unapply(tree: Ident)(implicit ctx: Context): Option[Tree] =
-      if (paramProxies.contains(tree.tpe))
+    def unapply(tree: Trees.Ident[_])(implicit ctx: Context): Option[Tree] =
+      if (paramProxies.contains(tree.typeOpt))
         bindingsBuf.find(_.name == tree.name) match {
           case Some(vdef: ValDef) if vdef.symbol.is(Inline) =>
             Some(vdef.rhs.changeOwner(vdef.symbol, ctx.owner))
@@ -541,46 +674,24 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
       else None
   }
 
-  /** A typer for inlined code. Its purpose is:
+  /** A base trait of InlineTyper and InlineReTyper containing operations that
+   *  work the same way for both. Beyond typing or retyping, an inline typer performs
+   *  the following functions:
+   *
    *  1. Implement constant folding over inlined code
    *  2. Selectively expand ifs with constant conditions
    *  3. Inline arguments that are by-name closures
    *  4. Make sure inlined code is type-correct.
    *  5. Make sure that the tree's typing is idempotent (so that future -Ycheck passes succeed)
    */
-  private object InlineTyper extends ReTyper {
-
-    override def ensureAccessible(tpe: Type, superAccess: Boolean, pos: Position)(implicit ctx: Context): Type = {
-      tpe match {
-        case tpe @ TypeRef(pre, _) if !tpe.symbol.isAccessibleFrom(pre, superAccess) =>
-          tpe.info match {
-            case TypeAlias(alias) => return ensureAccessible(alias, superAccess, pos)
-            case _ =>
-          }
-        case _ =>
-      }
-      super.ensureAccessible(tpe, superAccess, pos)
-    }
-
-    override def typedIdent(tree: untpd.Ident, pt: Type)(implicit ctx: Context) =
-      tree.asInstanceOf[tpd.Tree] match {
-        case InlineableArg(rhs) => inlining.println(i"inline arg $tree -> $rhs"); rhs
-        case _ => super.typedIdent(tree, pt)
-      }
-
-    override def typedSelect(tree: untpd.Select, pt: Type)(implicit ctx: Context): Tree = {
-      assert(tree.hasType, tree)
-      val qual1 = typed(tree.qualifier, selectionProto(tree.name, pt, this))
-      val res = untpd.cpy.Select(tree)(qual1, tree.name).withType(tree.typeOpt)
-      ensureAccessible(res.tpe, tree.qualifier.isInstanceOf[untpd.Super], tree.pos)
-      res
-    }
+  trait InlineTyping extends Typer {
 
     override def typedIf(tree: untpd.If, pt: Type)(implicit ctx: Context) = {
       val cond1 = typed(tree.cond, defn.BooleanType)
       cond1.tpe.widenTermRefExpr match {
         case ConstantType(Constant(condVal: Boolean)) =>
-          val selected = typed(if (condVal) tree.thenp else tree.elsep, pt)
+          var selected = typed(if (condVal) tree.thenp else tree.elsep, pt)
+          if (selected.isEmpty) selected = tpd.Literal(Constant(()))
           if (isIdempotentExpr(cond1)) selected
           else Block(cond1 :: Nil, selected)
         case _ =>
@@ -628,6 +739,48 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
     }
   }
 
+  /** A full typer used for transparent methods */
+  private class InlineTyper extends Typer with InlineTyping {
+
+    override def typedTypedSplice(tree: untpd.TypedSplice)(implicit ctx: Context): Tree =
+      tree.splice match {
+        case InlineableArg(rhs) => inlining.println(i"inline arg $tree -> $rhs"); rhs
+        case _ => super.typedTypedSplice(tree)
+      }
+  }
+
+  /** A re-typer used for inlined methods */
+  private class InlineReTyper extends ReTyper with InlineTyping {
+
+    override def ensureAccessible(tpe: Type, superAccess: Boolean, pos: Position)(implicit ctx: Context): Type = {
+      tpe match {
+        case tpe @ TypeRef(pre, _) if !tpe.symbol.isAccessibleFrom(pre, superAccess) =>
+          tpe.info match {
+            case TypeAlias(alias) => return ensureAccessible(alias, superAccess, pos)
+            case _ =>
+          }
+        case _ =>
+      }
+      super.ensureAccessible(tpe, superAccess, pos)
+    }
+
+    override def typedIdent(tree: untpd.Ident, pt: Type)(implicit ctx: Context) =
+      tree.asInstanceOf[tpd.Tree] match {
+        case InlineableArg(rhs) => inlining.println(i"inline arg $tree -> $rhs"); rhs
+        case _ => super.typedIdent(tree, pt)
+      }
+
+    override def typedSelect(tree: untpd.Select, pt: Type)(implicit ctx: Context): Tree = {
+      assert(tree.hasType, tree)
+      val qual1 = typed(tree.qualifier, selectionProto(tree.name, pt, this))
+      val res = untpd.cpy.Select(tree)(qual1, tree.name).withType(tree.typeOpt)
+      ensureAccessible(res.tpe, tree.qualifier.isInstanceOf[untpd.Super], tree.pos)
+      res
+    }
+
+    override def newLikeThis: Typer = new InlineReTyper
+  }
+
   /** Drop any side-effect-free bindings that are unused in expansion or other reachable bindings.
    *  Inline def bindings that are used only once.
    */

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

@@ -1910,7 +1910,7 @@ class Typer extends Namer
                     // from separately compiled files - the original BodyAnnotation is not kept.
                 }
                 else {
-                  assert(mdef.symbol.isTransparentMethod)
+                  assert(mdef1.symbol.isTransparentMethod, mdef.symbol)
                   Inliner.bodyToInline(mdef1.symbol) // just make sure accessors are computed,
                   buf += mdef1                       // but keep original definition, since inline-expanded code
                                                      // is pickled in this case.

tests/run/typelevel.scala

@@ -17,14 +17,14 @@ object Test extends App {
     }
 
   val x0 = ToNat(0)
-  //val y0: Z = x0.value
-  //val z0: x0.Result = y0
+  val y0: Z = x0.value
+  val z0: x0.Result = y0
   val x1 = ToNat(1)
-  //val y1: S[Z] = x1.value
-  //val z1: x1.Result = y1
+  val y1: S[Z] = x1.value
+  val z1: x1.Result = y1
   val x2 = ToNat(2)
-  //val y2: S[S[Z]] = x2.value
-  //val z2: x2.Result = y2
+  val y2: S[S[Z]] = x2.value
+  val z2: x2.Result = y2
   println(x0)
   println(x1)
   println(x2)