com.github.plokhotnyuk.jsoniter_scala.macros.LowLevelQuoteUtil.scala Maven / Gradle / Ivy

Go to download
package com.github.plokhotnyuk.jsoniter_scala.macros

import scala.quoted._

private[macros] object LowLevelQuoteUtil {

   /**
    * When we mix generation via trees and via quotes, we can have incorrect owners inside tree,
    * when tree with owner Symbol.currentOwner was inserted inside quote in other scope.
    *
    * We call deepChangeOwner when all val-s which defined in term is not used outside term
    * and over
    **/
  def deepChangeOwner(using Quotes)(tree: quotes.reflect.Tree, owner: quotes.reflect.Symbol, traceFlag: Boolean): quotes.reflect.Tree = {
    import quotes.reflect._

    val mapper = new TreeMap {
      override def transformTree(tree: Tree)(owner: Symbol): Tree =
        try super.transformTree(tree)(owner) catch {
          case ex: IllegalStateException =>
            if (traceFlag) println(s"${ex.getMessage()}\ntree=$tree")
            throw ex
        }

      override def transformStatement(tree: Statement)(owner: Symbol): Statement = tree match {
        case d: Definition =>
          if (d.symbol.owner != owner) throw new IllegalStateException("Invalid owner in definition.")
          d
        case _ => super.transformStatement(tree)(owner)
      }

      override def transformTerm(tree: Term)(owner: Symbol): Term = tree match {
        case Inlined(orgin, bindings, body) =>
          if (bindings.isEmpty) transformTerm(body)(owner)
          else super.transformTerm(tree)(owner)
        case bl@Block(statements, expr) =>
          var needTopLevelChange = false
          var incorrectSymbols: List[Symbol] = Nil
          statements.foreach {
            case d: Definition =>
              if (d.symbol.owner != owner) {
                incorrectSymbols = d.symbol :: incorrectSymbols
                needTopLevelChange = true
              }
            case other =>
          }
          val r =
            if (needTopLevelChange) {
               if (traceFlag) print(s"fixing owners for ${incorrectSymbols.mkString(",")}")
               bl.changeOwner(owner)
            } else bl
          val nStatements = r.statements.map {
            case t: Term => transformTerm(t)(owner)
            case other => other
          }
          val nExpr = transformTerm(r.expr)(owner)
          Block.copy(r)(nStatements, nExpr)
        case tif@If(cond, ifTrue, ifFalse) =>
          If.copy(tif)(transformTerm(cond)(owner), transformTerm(ifTrue)(owner), transformTerm(ifFalse)(owner))
        case _ =>
          super.transformTerm(tree)(owner)
      }

      override def transformTypeTree(tree: TypeTree)(owner: Symbol): TypeTree = tree // don't navigate over types.

      def checkInvalidOwner(term: Term, owner: Symbol): Boolean = checkOwner(term, owner, traceFlag, false, true)
    }
    mapper.transformTree(tree)(owner)
   }

  def checkOwner(using Quotes)(term: quotes.reflect.Term, ownerToCheck: quotes.reflect.Symbol, traceFlag: Boolean = true, throwFlag: Boolean = true, onlyFirst: Boolean = false): Boolean = {
    import quotes.reflect._

    var topLevelFound = false
    val treeTraverser = new TreeTraverser() {
      var wasException = false

      override def traverseTree(tree: Tree)(owner: Symbol): Unit = {
        if (topLevelFound && onlyFirst) return
        var foundInvalidOwner = false
        tree match {
          case d: Definition =>
            if (tree.symbol.owner != owner) {
              foundInvalidOwner = true
              topLevelFound = true
              if (traceFlag) println(s"checkOwner: owner mismatch for ${tree.show}, expectd owner: ${owner}, have ${tree.symbol.maybeOwner}")
              if (!wasException) {
                wasException = true
                throw new IllegalStateException(s"invlid owner, expected: ${owner}, have ${tree.symbol.owner}")
              }
            }
          case _ =>
        }
        try {
          if (!foundInvalidOwner) traverseTreeChildren(tree)(owner)
        } catch {
          case ex: IllegalStateException =>
            if (traceFlag) println(s"in tree:  $tree\n")
            if (throwFlag) throw ex
        }
     }
    }
    treeTraverser.foldTree((), term)(ownerToCheck)
    topLevelFound
  }
}