ResourceExtractorDbIndex.scala

package p
import scalafix.util.SemanticdbIndex
import scala.meta._


class ResourceExtractorDbIndex(index: SemanticdbIndex) extends SemanticdbIndex {

  private[this] var cache: Map[Symbol, Option[Defn]] = Map.empty

  def document(s: Symbol): Option[Document] =
    index.documents.find(d =>
      d.symbols.exists(rs => rs.symbol == s)
    )

  def definition(s: Symbol): Option[Defn] = {

    def collectDefnsInDocument(): List[Defn] =
      (for {
        d <- document(s)
        src <- d.input.parse[Source].toOption
      } yield src.collect {
        case defn: Defn.Class => defn
        case defn: Defn.Trait => defn
      }).getOrElse(Nil)

    if (cache.contains(s)) cache(s)
    else {
      val symDefs =
        collectDefnsInDocument().map {
          case defn: Defn.Class =>
            index.symbol(defn.name).map(_ -> Some(defn))
          case defn: Defn.Trait =>
            index.symbol(defn.name).map(_ -> Some(defn))
        }

      val newCache = symDefs.foldLeft(cache) {
        case (c, sd) =>
          if (sd.isEmpty) c
          else c + sd.get
      }

      cache = newCache
      cache.getOrElse(s, None)
    }

  }

  def findDefinition(s: Symbol) : Either[Error, Defn] =
    definition(s).toRight(new Error(s"No definition found for $s"))

  def namedDefinitions : Seq[Either[Error, (Symbol, Defn)]] =
    documents.flatMap{ d =>
      d.input.parse[Source] match {
        case e: Parsed.Error =>
          Left(new Error(e.toString(), e.details)) :: Nil
        case Parsed.Success(src) =>
          val res = src.collect {
            case defn: Defn.Class =>
              findSymbol(defn.name).map{s =>
                cache += s -> Some(defn)
                s -> defn
              }
            case defn: Defn.Trait =>
              findSymbol(defn.name).map{s =>
                cache += s -> Some(defn)
                s -> defn
              }
          }
          res
      }
    }

  def sourcepath: Sourcepath = index.sourcepath

  def classpath: Classpath = index.classpath

  def database: Database = index.database

  def names: Seq[ResolvedName] = index.names

  def symbol(position: Position): Option[Symbol] = index.symbol(position)

  def symbol(tree: Tree): Option[Symbol] = index.symbol(tree)

  def findSymbol(tree: Tree): Either[Error, Symbol] =
    index.symbol(tree).toRight(new Error(s"No symbol found for $tree"))


  def denotation(symbol: Symbol): Option[Denotation] = index.denotation(symbol)

  def denotation(tree: Tree): Option[Denotation] = index.denotation(tree)

  def withDocuments(documents: Seq[Document]): SemanticdbIndex = index.withDocuments(documents)
}