Phase for Slick, that unboxes wrapped tables

Declutter.scala

package scala.slick.compiler

import scala.slick.ast._
import scala.collection.mutable.HashMap

class Declutter extends Phase {
  val name = "declutter"

  def apply(tree: Node, state: CompilationState): Node = {
    val replaced = new HashMap[Symbol, Symbol]
    val unwrapTables = new Transformer {
      def replace = {
        case c @ Comprehension((_, from) :: Nil, Nil, None, Nil, Some(Pure(StructNode(select))), None, None) =>
          val unwrapped = select.forall {
            case (s, Path(col :: _)) => true
            case _ => false
          }
          if (unwrapped) {
            select.foreach {
              case (s, Path(col :: _)) => replaced += s -> col
            }
            from
          } else c
      }
    }
    val replaceRefs = new Transformer {
      def replace = {
        case p @ Path(s :: rest) if replaced.isDefinedAt(s) => Path(replaced(s) :: rest)
      }
    }
    replaceRefs.once(unwrapTables.repeat(tree))
  }
}