package macros.utils

import scala.language.experimental.macros
import scala.reflect.macros.blackbox._
import chisel3._
import chisel3.experimental.SourceInfo
import chisel3.util.BitPat

object bis {

    /** Executes `block` if the bswitch condition matchs `v`.
      */
    def apply(v: BitPat)(block: => Any): Unit = {
        locally((v, block))
        require(
            false,
            "The 'bis' keyword may not be used outside of a bswitch."
        )
    }
}

class BSwitchContext(
    cond: UInt,
    whenContext: Option[WhenContext],
    lits: Set[BitPat]
) {
    def bis(
        pat: BitPat
    )(block: => Any)(implicit
        sourceInfo: SourceInfo
    ): BSwitchContext = {
        require(
            !lits.exists(w => {
                w.overlap(pat)
            }),
            "all bis conditions must be mutually exclusive!"
        )

        // def instead of val so that logic ends up in legal place
        def p = (cond === pat)
        whenContext match {
            case Some(w) =>
                new BSwitchContext(
                    cond,
                    Some(w.elsewhen(p)(block)),
                    lits + pat
                )
            case None =>
                new BSwitchContext(
                    cond,
                    Some(when(p)(block)),
                    lits + pat
                )
        }
    }
}

object bswitch {
    def apply[T <: Element](cond: T)(x: => Any): Unit = macro impl
    def impl(c: Context)(cond: c.Tree)(x: c.Tree): c.Tree = {
        import c.universe._
        val q"..$body" = x
        val res = body.foldLeft(
            q"""new macros.utils.BSwitchContext($cond, None, Set.empty)"""
        ) { case (acc, tree) =>
            tree match {
                case q"macros.utils.bis.apply( ..$params )( ..$body )" =>
                    q"$acc.bis( ..$params )( ..$body )"
                case _ =>
                    throw new Exception(
                        s"Cannot include blocks that do not begin with bis() in bswitch."
                    )
            }
        }
        q"""{ $res }"""
    }
}
