lazytrace.scala

package localhost

import scala.tools.nsc
import nsc.Global
import nsc.plugins.Plugin
import nsc.plugins.PluginComponent
import nsc.transform.Transform
import scala.tools.nsc.ast.TreeGen

import scala.tools.nsc.symtab.Flags
import Flags.{PARAM, SYNTHETIC}

class LazyTracePlugin(val global: Global) extends Plugin {
  import global._
  import definitions._

  val name = "lazy-trace"
  val description = "print stack trace in creating lazy value"
  val components = List[PluginComponent](LazyTraceComponent)

  def mkPrintlnTree(tree:Tree) =
    Apply(gen.mkAttributedRef(getMember(PredefModule, "println")), List(tree))

  def mkGetStackTraceTree = {
    val threadClass = definitions.getClass("java.lang.Thread")
    val threadModule = getModule("java.lang.Thread")
    val currentThread = gen.mkAttributedRef(getMember(threadModule, "currentThread"))
    val stackTraceMethod = getMember(threadClass, "getStackTrace")
    
    Apply(Select(Apply(currentThread,Nil) setType currentThread.tpe,
		 stackTraceMethod) setType stackTraceMethod.tpe,
	  List()) setType arrayType(stackTraceMethod.tpe)
  }

  def mkPrintlnFuncTree(parent:Tree) = {
    val anonfun = parent.symbol.newValue(NoPosition, "$anonfun") setInfo(NoType)
    val x = anonfun.newSyntheticValueParam(AnyClass.tpe)

    val param = ValDef(Modifiers(PARAM|SYNTHETIC), x.name, TypeTree(AnyClass.tpe), EmptyTree) setSymbol x
    global.Function(List(param), mkPrintlnTree(Ident(x))) setSymbol anonfun setType functionType(List(AnyClass.tpe), UnitClass.tpe)
  }

  def mkPrintStackTraceTree(parent:Tree) = {
    val stElementClass = definitions.getClass("java.lang.StackTraceElement")
    val refArrayOps = gen.mkAttributedRef(getMember(PredefModule, "refArrayOps"))
    val arrayOpsClass = definitions.getClass("scala.collection.mutable.ArrayOps")
    val foreachMethod = definitions.getMember(arrayOpsClass, "foreach")

    val retType = typeRef(arrayOpsClass.typeConstructor.prefix, arrayOpsClass, List(stElementClass.tpe))
    val param = getMember(PredefModule, "refArrayOps").newSyntheticValueParam(arrayType(stElementClass.tpe))
    val typedArrayOps = TypeApply(refArrayOps, List(TypeTree(stElementClass.tpe))) setType MethodType(List(param), retType)

    Apply(
      TypeApply(
	Select(
	  Apply(typedArrayOps, List(mkGetStackTraceTree)) setType arrayType(stElementClass.tpe),
	  foreachMethod) setType foreachMethod.tpe,
	List(TypeTree(UnitClass.tpe))) setType foreachMethod.tpe,
      List(mkPrintlnFuncTree(parent))) setType UnitClass.tpe
  }

  private object LazyTraceComponent extends PluginComponent with Transform{

    val global: LazyTracePlugin.this.global.type = LazyTracePlugin.this.global
    val runsAfter = List[String]("typer")
    val phaseName = LazyTracePlugin.this.name

    def newTransformer(unit: global.CompilationUnit) = LazyTraceTransformer

    object LazyTraceTransformer extends global.Transformer {
      override def transform(tree: global.Tree) = {
        tree match {
	  case v @ ValDef(mods, name, tpe, rhs) if v.symbol.isLazy =>
	    val treeCopier = new StrictTreeCopier
	    treeCopier.ValDef(v, mods, name, tpe, Block(mkPrintStackTraceTree(tree),rhs))
	  case t => super.transform(t)
	}
      }
    }
  }
}