Changed boxing of free mutable variables to be flexible wrt when liftcode takes place.

A major redesign that unifies the various different approaches to boxing of free variables. Free variables are marked with CAPTURED and eliminated by LambdaLift. I also added some hooks in MacroContext that a reifier needs to use.

8 files changed, 105 insertions, 122 deletions

diff --git a/src/compiler/scala/tools/nsc/Global.scala b/src/compiler/scala/tools/nsc/Global.scala
index f70aa60309..2dd32e355b 100644
--- a/src/compiler/scala/tools/nsc/Global.scala
+++ b/src/compiler/scala/tools/nsc/Global.scala
@@ -39,6 +39,7 @@ class Global(var currentSettings: Settings, var reporter: Reporter) extends Symb
                                                                       with Trees
                                                                       with TreePrinters
                                                                       with DocComments
+                                                                      with MacroContext
                                                                       with symtab.Positions {
 
   override def settings = currentSettings
@@ -151,7 +152,7 @@ class Global(var currentSettings: Settings, var reporter: Reporter) extends Symb
   /** Register top level class (called on entering the class)
    */
   def registerTopLevelSym(sym: Symbol) {}
-
+  
 // ------------------ Reporting -------------------------------------
 
   // not deprecated yet, but a method called "error" imported into
diff --git a/src/compiler/scala/tools/nsc/MacroContext.scala b/src/compiler/scala/tools/nsc/MacroContext.scala
new file mode 100644
index 0000000000..e739eade3a
--- /dev/null
+++ b/src/compiler/scala/tools/nsc/MacroContext.scala
@@ -0,0 +1,10 @@
+package scala.tools.nsc
+
+import symtab.Flags._
+
+trait MacroContext extends reflect.api.MacroContext { self: Global =>
+  
+  def captureVariable(vble: Symbol): Unit = vble setFlag CAPTURED
+  
+  def referenceCapturedVariable(id: Ident): Tree = ReferenceToBoxed(id)
+}
diff --git a/src/compiler/scala/tools/nsc/ast/Trees.scala b/src/compiler/scala/tools/nsc/ast/Trees.scala
index 30ee7fc885..88a9b5e18b 100644
--- a/src/compiler/scala/tools/nsc/ast/Trees.scala
+++ b/src/compiler/scala/tools/nsc/ast/Trees.scala
@@ -44,6 +44,15 @@ trait Trees extends reflect.internal.Trees { self: Global =>
 
   /** emitted by typer, eliminated by refchecks */
   case class TypeTreeWithDeferredRefCheck()(val check: () => TypeTree) extends TypTree
+  
+  /** Marks underlying reference to id as boxed. 
+   *  @pre: id must refer to a captured variable
+   *  A reference such marked will refer to the boxed entity, no dereferencing
+   *  with `.elem` is done on it.
+   *  This tree node can be emitted by macros such as reify that call markBoxedReference.
+   *  It is eliminated in LambdaLift, where the boxing conversion takes place.
+   */
+  case class ReferenceToBoxed(idt: Ident) extends TermTree
 
   // --- factory methods ----------------------------------------------------------
 
@@ -152,6 +161,8 @@ trait Trees extends reflect.internal.Trees { self: Global =>
       traverser.traverse(lhs); traverser.traverse(rhs)
     case SelectFromArray(qualifier, selector, erasure) =>
       traverser.traverse(qualifier)
+    case ReferenceToBoxed(idt) =>
+      traverser.traverse(idt)
     case TypeTreeWithDeferredRefCheck() => // TODO: should we traverse the wrapped tree?
       // (and rewrap the result? how to update the deferred check? would need to store wrapped tree instead of returning it from check)
     case _ => super.xtraverse(traverser, tree)
@@ -161,6 +172,7 @@ trait Trees extends reflect.internal.Trees { self: Global =>
     def DocDef(tree: Tree, comment: DocComment, definition: Tree): DocDef
     def AssignOrNamedArg(tree: Tree, lhs: Tree, rhs: Tree): AssignOrNamedArg
     def SelectFromArray(tree: Tree, qualifier: Tree, selector: Name, erasure: Type): SelectFromArray
+    def ReferenceToBoxed(tree: Tree, idt: Ident): ReferenceToBoxed
     def TypeTreeWithDeferredRefCheck(tree: Tree): TypeTreeWithDeferredRefCheck
   }
 
@@ -174,6 +186,8 @@ trait Trees extends reflect.internal.Trees { self: Global =>
       new AssignOrNamedArg(lhs, rhs).copyAttrs(tree)
     def SelectFromArray(tree: Tree, qualifier: Tree, selector: Name, erasure: Type) =
       new SelectFromArray(qualifier, selector, erasure).copyAttrs(tree)
+    def ReferenceToBoxed(tree: Tree, idt: Ident) =
+      new ReferenceToBoxed(idt).copyAttrs(tree)
     def TypeTreeWithDeferredRefCheck(tree: Tree) = tree match {
       case dc@TypeTreeWithDeferredRefCheck() => new TypeTreeWithDeferredRefCheck()(dc.check).copyAttrs(tree)
     }
@@ -195,6 +209,11 @@ trait Trees extends reflect.internal.Trees { self: Global =>
       if (qualifier0 == qualifier) && (selector0 == selector) => t
       case _ => this.treeCopy.SelectFromArray(tree, qualifier, selector, erasure)
     }
+    def ReferenceToBoxed(tree: Tree, idt: Ident) = tree match {
+      case t @ ReferenceToBoxed(idt0) 
+      if (idt0 == idt) => t
+      case _ => this.treeCopy.ReferenceToBoxed(tree, idt)
+    }
     def TypeTreeWithDeferredRefCheck(tree: Tree) = tree match {
       case t @ TypeTreeWithDeferredRefCheck() => t
       case _ => this.treeCopy.TypeTreeWithDeferredRefCheck(tree)
@@ -220,6 +239,9 @@ trait Trees extends reflect.internal.Trees { self: Global =>
     case SelectFromArray(qualifier, selector, erasure) =>
       transformer.treeCopy.SelectFromArray(
         tree, transformer.transform(qualifier), selector, erasure)
+    case ReferenceToBoxed(idt) =>
+      transformer.treeCopy.ReferenceToBoxed(
+        tree, transformer.transform(idt) match { case idt1: Ident => idt1 })
     case TypeTreeWithDeferredRefCheck() =>
       transformer.treeCopy.TypeTreeWithDeferredRefCheck(tree)
   }
diff --git a/src/compiler/scala/tools/nsc/transform/LambdaLift.scala b/src/compiler/scala/tools/nsc/transform/LambdaLift.scala
index 2310eae9bb..2180fd4f3a 100644
--- a/src/compiler/scala/tools/nsc/transform/LambdaLift.scala
+++ b/src/compiler/scala/tools/nsc/transform/LambdaLift.scala
@@ -17,6 +17,19 @@ abstract class LambdaLift extends InfoTransform {
 
   /** the following two members override abstract members in Transform */
   val phaseName: String = "lambdalift"
+    
+  /** Converts types of captured variables to *Ref types.
+   */
+  def boxIfCaptured(sym: Symbol, tpe: Type, erasedTypes: Boolean) =
+    if (sym.isCapturedVariable) {
+      val symClass = tpe.typeSymbol
+      def refType(valueRef: Map[Symbol, Symbol], objectRefClass: Symbol) =
+        if (isValueClass(symClass) && symClass != UnitClass) valueRef(symClass).tpe
+        else if (erasedTypes) objectRefClass.tpe
+        else appliedType(objectRefClass.typeConstructor, List(tpe))
+      if (sym.hasAnnotation(VolatileAttr)) refType(volatileRefClass, VolatileObjectRefClass)
+      else refType(refClass, ObjectRefClass)
+    } else tpe
 
   private val lifted = new TypeMap {
     def apply(tp: Type): Type = tp match {
@@ -31,7 +44,8 @@ abstract class LambdaLift extends InfoTransform {
     }
   }
 
-  def transformInfo(sym: Symbol, tp: Type): Type = lifted(tp)
+  def transformInfo(sym: Symbol, tp: Type): Type =
+    boxIfCaptured(sym, lifted(tp), erasedTypes = true)
 
   protected def newTransformer(unit: CompilationUnit): Transformer =
     new LambdaLifter(unit)
@@ -55,7 +69,10 @@ abstract class LambdaLift extends InfoTransform {
 
     /** Buffers for lifted out classes and methods */
     private val liftedDefs = new LinkedHashMap[Symbol, List[Tree]]
-
+    
+    /** True if we are transforming under a ReferenceToBoxed node */
+    private var isBoxedRef = false
+    
     private type SymSet = TreeSet[Symbol]
 
     private def newSymSet = new TreeSet[Symbol](_ isLess _)
@@ -116,22 +133,7 @@ abstract class LambdaLift extends InfoTransform {
           }
           changedFreeVars = true
           debuglog("" + sym + " is free in " + enclosure);
-          if (sym.isVariable && !sym.hasFlag(CAPTURED)) {
-            // todo: We should merge this with the lifting done in liftCode.
-            // We do have to lift twice: in liftCode, because Code[T] needs to see the lifted version
-            // and here again because lazy bitmaps are introduced later and get lifted here.
-            // But we should factor out the code and run it twice.
-            sym setFlag CAPTURED
-            val symClass = sym.tpe.typeSymbol
-            atPhase(phase.next) {
-              sym updateInfo (
-                if (sym.hasAnnotation(VolatileAttr))
-                  if (isValueClass(symClass)) volatileRefClass(symClass).tpe else VolatileObjectRefClass.tpe
-                else
-                  if (isValueClass(symClass)) refClass(symClass).tpe else ObjectRefClass.tpe
-              )
-            }
-          }
+          if (sym.isVariable) sym setFlag CAPTURED
         }
         !enclosure.isClass
       }
@@ -340,7 +342,7 @@ abstract class LambdaLift extends InfoTransform {
       EmptyTree
     }
 
-    private def postTransform(tree: Tree): Tree = {
+    private def postTransform(tree: Tree, isBoxedRef: Boolean = false): Tree = {
       val sym = tree.symbol
       tree match {
         case ClassDef(_, _, _, _) =>
@@ -363,8 +365,19 @@ abstract class LambdaLift extends InfoTransform {
                 }
               case arg => arg
             }
+            /** Wrap expr argument in new *Ref(..) constructor, but make
+             *  sure that Try expressions stay at toplevel.
+             */
+            def refConstr(expr: Tree): Tree = expr match {
+              case Try(block, catches, finalizer) =>
+                Try(refConstr(block), catches map refConstrCase, finalizer)
+              case _ => 
+                Apply(Select(New(TypeTree(sym.tpe)), nme.CONSTRUCTOR), List(expr))
+            }
+            def refConstrCase(cdef: CaseDef): CaseDef = 
+              CaseDef(cdef.pat, cdef.guard, refConstr(cdef.body))
             treeCopy.ValDef(tree, mods, name, tpt1, typer.typedPos(rhs.pos) {
-              Apply(Select(New(TypeTree(sym.tpe)), nme.CONSTRUCTOR), List(constructorArg))
+              refConstr(constructorArg)
             })
           } else tree
         case Return(Block(stats, value)) =>
@@ -388,7 +401,7 @@ abstract class LambdaLift extends InfoTransform {
                 atPos(tree.pos)(proxyRef(sym))
               else tree
             else tree
-          if (sym.isCapturedVariable)
+          if (sym.isCapturedVariable && !isBoxedRef)
             atPos(tree.pos) {
               val tp = tree.tpe
               val elemTree = typer typed Select(tree1 setType sym.tpe, nme.elem)
@@ -406,10 +419,16 @@ abstract class LambdaLift extends InfoTransform {
           tree
       }
     }
+    
+    private def preTransform(tree: Tree) = super.transform(tree) setType lifted(tree.tpe)
 
-    override def transform(tree: Tree): Tree =
-      postTransform(super.transform(tree) setType lifted(tree.tpe))
-
+    override def transform(tree: Tree): Tree = tree match {
+      case ReferenceToBoxed(idt) =>
+        postTransform(preTransform(idt), isBoxedRef = true)
+      case _ =>
+        postTransform(preTransform(tree))
+    }
+      
     /** Transform statements and add lifted definitions to them. */
     override def transformStats(stats: List[Tree], exprOwner: Symbol): List[Tree] = {
       def addLifted(stat: Tree): Tree = stat match {
diff --git a/src/compiler/scala/tools/nsc/transform/LiftCode.scala b/src/compiler/scala/tools/nsc/transform/LiftCode.scala
index 720509644b..171d1df975 100644
--- a/src/compiler/scala/tools/nsc/transform/LiftCode.scala
+++ b/src/compiler/scala/tools/nsc/transform/LiftCode.scala
@@ -110,18 +110,10 @@ abstract class LiftCode extends Transform with TypingTransformers {
       }
     }
 
-    /** Set of mutable local variables that are free in some inner method. */
-    private val freeMutableVars: mutable.Set[Symbol] = new mutable.HashSet
-    private val converted: mutable.Set[Symbol] = new mutable.HashSet // debug
-
     override def transformUnit(unit: CompilationUnit) {
-      freeMutableVars.clear()
-      freeLocalsTraverser(unit.body)
       atPhase(phase.next) {
         super.transformUnit(unit)
       }
-      for (v <- freeMutableVars) //!!! remove
-        assert(converted contains v, "unconverted: " + v + " in " + v.owner + " in unit " + unit)
     }
 
     override def transform(tree: Tree): Tree = {
@@ -137,24 +129,6 @@ abstract class LiftCode extends Transform with TypingTransformers {
               result
             }
           } finally printTypings = saved
-        case ValDef(mods, name, tpt, rhs) if (freeMutableVars(sym)) => // box mutable variables that are accessed from a local closure
-          val tpt1 = TypeTree(sym.tpe) setPos tpt.pos
-          /* Creating a constructor argument if one isn't present. */
-          val constructorArg = rhs match {
-            case EmptyTree => gen.mkZero(atPhase(phase.prev)(sym.tpe))
-            case _ => transform(rhs)
-          }
-          val rhs1 = typer.typedPos(rhs.pos) {
-            Apply(Select(New(TypeTree(sym.tpe)), nme.CONSTRUCTOR), List(constructorArg))
-          }
-          sym resetFlag MUTABLE
-          sym removeAnnotation VolatileAttr
-          converted += sym // dereference boxed variables
-          treeCopy.ValDef(tree, mods &~ MUTABLE, name, tpt1, rhs1)
-        case Ident(name) if freeMutableVars(sym) =>
-          localTyper.typedPos(tree.pos) {
-            Select(tree setType sym.tpe, nme.elem)
-          }
         case _ =>
           super.transform(tree)
       }
@@ -170,74 +144,6 @@ abstract class LiftCode extends Transform with TypingTransformers {
       New(TypeTree(appliedType(definitions.CodeClass.typeConstructor, List(treetpe.widen))),
         List(List(arg)))
     }
-
-    /**
-     * PP: There is apparently some degree of overlap between the CAPTURED
-     *  flag and the role being filled here.  I think this is how this was able
-     *  to go for so long looking only at DefDef and Ident nodes, as bugs
-     *  would only emerge under more complicated conditions such as #3855.
-     *  I'll try to figure it all out, but if someone who already knows the
-     *  whole story wants to fill it in, that too would be great.
-     *
-     *  XXX I found this had been cut and pasted between LiftCode and UnCurry,
-     *  and seems to be running in both.
-     */
-    private val freeLocalsTraverser = new Traverser {
-      var currentMethod: Symbol = NoSymbol
-      var maybeEscaping = false
-
-      def withEscaping(body: => Unit) {
-        val saved = maybeEscaping
-        maybeEscaping = true
-        try body
-        finally maybeEscaping = saved
-      }
-
-      override def traverse(tree: Tree) = tree match {
-        case DefDef(_, _, _, _, _, _) =>
-          val lastMethod = currentMethod
-          currentMethod = tree.symbol
-          try super.traverse(tree)
-          finally currentMethod = lastMethod
-        /** A method call with a by-name parameter represents escape. */
-        case Apply(fn, args) if fn.symbol.paramss.nonEmpty =>
-          traverse(fn)
-          treeInfo.foreachMethodParamAndArg(tree) { (param, arg) =>
-            if (param.tpe != null && isByNameParamType(param.tpe))
-              withEscaping(traverse(arg))
-            else
-              traverse(arg)
-          }
-
-        /** The rhs of a closure represents escape. */
-        case Function(vparams, body) =>
-          vparams foreach traverse
-          withEscaping(traverse(body))
-
-        /**
-         * The appearance of an ident outside the method where it was defined or
-         *  anytime maybeEscaping is true implies escape.
-         */
-        case Ident(_) =>
-          val sym = tree.symbol
-          if (sym.isVariable && sym.owner.isMethod && (maybeEscaping || sym.owner != currentMethod)) {
-            freeMutableVars += sym
-            val symTpe = sym.tpe
-            val symClass = symTpe.typeSymbol
-            atPhase(phase.next) {
-              def refType(valueRef: Map[Symbol, Symbol], objectRefClass: Symbol) =
-                if (isValueClass(symClass) && symClass != UnitClass) valueRef(symClass).tpe
-                else appliedType(objectRefClass.typeConstructor, List(symTpe))
-
-              sym updateInfo (
-                if (sym.hasAnnotation(VolatileAttr)) refType(volatileRefClass, VolatileObjectRefClass)
-                else refType(refClass, ObjectRefClass))
-            }
-          }
-        case _ =>
-          super.traverse(tree)
-      }
-    }
   }
 
   /**
@@ -385,7 +291,10 @@ abstract class LiftCode extends Transform with TypingTransformers {
           else {
             if (sym.isTerm) {
               if (reifyDebug) println("Free: " + sym)
-              mirrorCall("freeVar", reify(sym.name.toString), reify(sym.tpe), Ident(sym))
+              val symtpe = lambdaLift.boxIfCaptured(sym, sym.tpe, erasedTypes = false)
+              def markIfCaptured(arg: Ident): Tree = 
+                if (sym.isCapturedVariable) referenceCapturedVariable(arg) else arg
+              mirrorCall("freeVar", reify(sym.name.toString), reify(symtpe), markIfCaptured(Ident(sym)))
             } else {
               if (reifyDebug) println("Late local: " + sym)
               registerReifiableSymbol(sym)
@@ -471,7 +380,10 @@ abstract class LiftCode extends Transform with TypingTransformers {
       case This(_) if !(boundSyms contains tree.symbol) =>
         reifyFree(tree)
       case Ident(_) if !(boundSyms contains tree.symbol) =>
-        reifyFree(tree)
+        if (tree.symbol.isVariable && tree.symbol.owner.isTerm) {
+          captureVariable(tree.symbol) // Note order dependency: captureVariable needs to come before reifyTree here.
+          mirrorCall("Select", reifyFree(tree), reifyName(nme.elem))
+        } else reifyFree(tree)
       case tt: TypeTree if (tt.tpe != null) =>
         if (!(boundSyms exists (tt.tpe contains _))) mirrorCall("TypeTree", reifyType(tt.tpe))
         else if (tt.original != null) reify(tt.original)
diff --git a/src/compiler/scala/tools/nsc/typechecker/Duplicators.scala b/src/compiler/scala/tools/nsc/typechecker/Duplicators.scala
index 2bed5bffd4..3536608efd 100644
--- a/src/compiler/scala/tools/nsc/typechecker/Duplicators.scala
+++ b/src/compiler/scala/tools/nsc/typechecker/Duplicators.scala
@@ -248,7 +248,7 @@ abstract class Duplicators extends Analyzer {
 
         case vdef @ ValDef(mods, name, tpt, rhs) =>
           // log("vdef fixing tpe: " + tree.tpe + " with sym: " + tree.tpe.typeSymbol + " and " + invalidSyms)
-          if (mods.hasFlag(Flags.LAZY)) vdef.symbol.resetFlag(Flags.MUTABLE)
+          //if (mods.hasFlag(Flags.LAZY)) vdef.symbol.resetFlag(Flags.MUTABLE) // Martin to Iulian: lazy vars can now appear because they are no longer boxed; Please check that deleting this statement is OK.
           vdef.tpt.tpe = fixType(vdef.tpt.tpe)
           vdef.tpe = null
           super.typed(vdef, mode, pt)
diff --git a/src/compiler/scala/tools/nsc/typechecker/Typers.scala b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
index 6b6b905e16..9991836344 100644
--- a/src/compiler/scala/tools/nsc/typechecker/Typers.scala
+++ b/src/compiler/scala/tools/nsc/typechecker/Typers.scala
@@ -4271,6 +4271,10 @@ trait Typers extends Modes with Adaptations with PatMatVirtualiser {
           else
             typedIdent(name)
 
+        case ReferenceToBoxed(idt @ Ident(_)) =>
+          val id1 = typed1(idt, mode, pt) match { case id: Ident => id }
+          treeCopy.ReferenceToBoxed(tree, id1) setType AnyRefClass.tpe         
+
         case Literal(value) =>
           tree setType (
             if (value.tag == UnitTag) UnitClass.tpe
diff --git a/src/library/scala/reflect/api/MacroContext.scala b/src/library/scala/reflect/api/MacroContext.scala
new file mode 100644
index 0000000000..e23357d26e
--- /dev/null
+++ b/src/library/scala/reflect/api/MacroContext.scala
@@ -0,0 +1,15 @@
+package scala.reflect
+package api
+
+trait MacroContext extends Universe {
+  
+  /** Mark a variable as captured; i.e. force boxing in a *Ref type.
+   */
+  def captureVariable(vble: Symbol): Unit
+  
+  /** Mark given identifier as a reference to a captured variable itself
+   *  suppressing dereferencing with the `elem` field.
+   */
+  def referenceCapturedVariable(id: Ident): Tree
+
+}
\ No newline at end of file