メソッドの呼び出し結果をキャッシュするMacro

CallCacheMacro.scala

import java.util.concurrent.ConcurrentHashMap
import scala.collection.JavaConverters._
import scala.concurrent.Future
import scala.reflect.macros.blackbox.Context
import scala.language.experimental.macros

/**
  *
  * メッソドの呼び出しをキャッシュし、一回しか実行されないようにする機能を提供するためのマクロ
  *
  * def hoge(i: Int) = println(i)
  *
  * object MyCache extends CallCache
  *
  * MyCache.withCache(hoge(2))// 2 is printed
  * MyCache.withCache(hoge(2))// nothing is printed. So call result is cached
  * MyCache.withCache(hoge(3))// 3 is printed. Because args for 3 is not cached.
  *
  *
  * キャッシュ可能なのは、単メソッド呼び出しのみ
  *
  * Created by takezoux2 on 2016/08/01.
  */
object CallCacheMacro {


  def autoCacheImpl[T: c.WeakTypeTag](c: Context)(anyCall: c.Expr[T]) : c.Expr[T] = {
    import c.universe._
    val t = implicitly[WeakTypeTag[T]]

    def extractMethodNameAndArgs(call : Tree) : (String,List[c.universe.Tree]) = {
      call match{
          // hoge(2)(3)(4)などのコードを分解
        case Apply(inner @ Apply(_,args1),args2) => {
          val (mn,args) = extractMethodNameAndArgs(inner)
          (mn,args ++ args2)
        }
        case Apply(methodName,args) => {
          (methodName.toString(),args)
        }
        case _ => {
          c.abort(c.enclosingPosition,s"Not simple method call.${anyCall}")
        }
      }
    }

    val (methodName,args) = extractMethodNameAndArgs(anyCall.tree)

    val mn = methodName.toString() + ":"
    val makeKey = q"""{${mn} + Seq(${ args:_*}).mkString(":")}"""

    val tree = if(t.tpe <:< typeOf[Future[_]]){
      q"""${c.prefix}.withFutureCache(${makeKey},{ ${anyCall} })"""
    }else{
      q"""${c.prefix}.withCache(${makeKey},{ ${anyCall} })"""
    }
    //println(tree)
    c.Expr[T](tree)
  }


}

trait CallCache{

  private var cache = new ConcurrentHashMap[String,Any]().asScala

  implicit def ec = scala.concurrent.ExecutionContext.global


  def withCache[T](key: String,func: => T) : T = {
    cache.getOrElseUpdate(key,{
      func
    }).asInstanceOf[T]
  }

  def withFutureCache[T](key: String, func: => Future[T]) : Future[T] = {
    cache.getOrElseUpdate(key,{
      val f = func
      f.map(v => {
        cache += (key -> v)
      })
    }) match{
      case f : Future[_] => f.asInstanceOf[Future[T]]
      case v => Future.successful(v.asInstanceOf[T])
    }
  }


  /**
    * 同じ関数、引数が呼ばれた場合、キャッシュを返す。
    * Future型を返す場合も適切にキャッシュされる
    * @param anyCall
    * @tparam T
    * @return
    */
  def withCache[T](anyCall : T) : T = macro CallCacheMacro.autoCacheImpl[T]


}