Airstream's Split by type macro

SplittableTypeMacros.scala

import com.raquo.airstream.core.{EventStream, Signal, Observable, BaseObservable}
import scala.quoted.{Expr, Quotes, Type}

object SplittableTypeMacros {

  extension [Self[+_] <: Observable[_], I](inline observable: BaseObservable[Self, I]) {
    inline def splitMatch: MatchSplitObservable[Self, I, Nothing] = MatchSplitObservable(observable, Nil, Map.empty[Int, Function[Any, Nothing]])
  }

  extension [Self[+_] <: Observable[_], I, O](inline matchSplitObservable: MatchSplitObservable[Self, I, O]) {
    inline def handleCase[A, B, O1 >: O](inline casePf: PartialFunction[A, B])(inline handleFn: ((B, Signal[B])) => O1) = ${ handleCaseImpl('{ matchSplitObservable }, '{ casePf }, '{ handleFn })}
  }

  extension [I, O](inline matchSplitObservable: MatchSplitObservable[Signal, I, O]) {
    inline def toSignal: Signal[O] = ${ observableImpl('{ matchSplitObservable })}
  }

  extension [I, O](inline matchSplitObservable: MatchSplitObservable[EventStream, I, O]) {
    inline def toStream: EventStream[O] = ${ observableImpl('{ matchSplitObservable })}
  }

  final case class MatchSplitObservable[Self[+_] <: Observable[_] , I, O] (
    observable: BaseObservable[Self, I],
    caseList: List[PartialFunction[Any, Any]],
    handlerMap: Map[Int, Function[Any, O]]
  )

  private def handleCaseImpl[Self[+_] <: Observable[_] : Type, I: Type, O: Type, O1 >: O : Type, A: Type, B: Type](
    matchSplitObservableExpr: Expr[MatchSplitObservable[Self, I, O]],
    casePfExpr: Expr[PartialFunction[A, B]],
    handleFnExpr: Expr[Function[(B, Signal[B]), O1]],
  )(
    using quotes: Quotes
  ): Expr[MatchSplitObservable[Self, I, O1]] = {
    import quotes.reflect.*
    
    matchSplitObservableExpr match {
      case '{ MatchSplitObservable[Self, I, O]($observableExpr, $caseListExpr, $handlerMapExpr)} => innerHandleCaseImpl(observableExpr, caseListExpr, handlerMapExpr, casePfExpr, handleFnExpr)
      case other => report.errorAndAbort("Macro expansion failed, please use `splitMatch` instead of creating new MatchSplitObservable explicitly")
    }
  }

  private def innerHandleCaseImpl[Self[+_] <: Observable[_] : Type, I: Type, O: Type, O1 >: O : Type, A: Type, B: Type](
    observableExpr: Expr[BaseObservable[Self, I]],
    caseListExpr: Expr[List[PartialFunction[Any, Any]]],
    handlerMapExpr: Expr[Map[Int, Function[Any, O]]],
    casePfExpr: Expr[PartialFunction[A, B]],
    handleFnExpr: Expr[Function[(B, Signal[B]), O1]]
  )(
    using quotes: Quotes
  ): Expr[MatchSplitObservable[Self, I, O1]] = {
    val caseExprList = exprOfListToListOfExpr(caseListExpr)

    val nextCaseExprList = casePfExpr.asExprOf[PartialFunction[Any, Any]] :: caseExprList

    val nextCaseListExpr = listOfExprToExprOfList(nextCaseExprList)

    '{ MatchSplitObservable[Self, I, O1]($observableExpr, $nextCaseListExpr, ($handlerMapExpr + ($handlerMapExpr.size -> $handleFnExpr.asInstanceOf[Function[Any, O1]]))) }
  }

  private def exprOfListToListOfExpr(
    pfListExpr: Expr[List[PartialFunction[Any, Any]]]
  )(
    using quotes: Quotes
  ): List[Expr[PartialFunction[Any, Any]]] = {
    import quotes.reflect.*
    
    pfListExpr match {
      case '{ $headExpr :: (${tailExpr}: List[PartialFunction[Any, Any]]) } =>
        headExpr :: exprOfListToListOfExpr(tailExpr)
      case '{ Nil } => Nil
      case _ => report.errorAndAbort("Macro expansion failed, please use `handleCase` instead of modify MatchSplitObservable explicitly")
    }
    
  }

  private def listOfExprToExprOfList(
    pfExprList: List[Expr[PartialFunction[Any, Any]]]
  )(
    using quotes: Quotes
  ): Expr[List[PartialFunction[Any, Any]]] = {
    import quotes.reflect.*
    
    pfExprList match
      case head :: tail => '{ $head :: ${listOfExprToExprOfList(tail)} }
      case Nil => '{ Nil }
  }

  private def observableImpl[Self[+_] <: Observable[_] : Type, I: Type, O: Type](
    matchSplitObservableExpr: Expr[MatchSplitObservable[Self, I, O]]
  )(
    using quotes: Quotes
  ): Expr[Self[O]] = {
    import quotes.reflect.*

    matchSplitObservableExpr match {
      case '{ MatchSplitObservable[Self, I, O]($_, Nil, $_)} =>
        report.errorAndAbort("Macro expansion failed, need at least one handleCase")
      case '{ MatchSplitObservable[Self, I, O]($observableExpr, $caseListExpr, $handlerMapExpr)} => 
        '{
          $observableExpr
            .map(i => ${ innerObservableImpl('i, caseListExpr) })
            .asInstanceOf[Observable[(Int, Any)]]
            .matchStreamOrSignal(
              ifSignal = _.splitOne(_._1) { case (idx, (_, b), dataSignal) =>
                val bSignal = dataSignal.map(_._2)
                $handlerMapExpr.apply(idx).apply(b -> bSignal)
              },
              ifStream = _.splitOne(_._1) { case (idx, (_, b), dataSignal) =>
                val bSignal = dataSignal.map(_._2)
                $handlerMapExpr.apply(idx).apply(b -> bSignal)
              }
            ).asInstanceOf[Self[O]]
        }
      case _ => report.errorAndAbort("Macro expansion failed, please use `splitMatch` instead of creating new MatchSplitObservable explicitly")
    }
  }

  private def innerObservableImpl[I: Type](
    iExpr: Expr[I],
    caseListExpr: Expr[List[PartialFunction[Any, Any]]]
  )(
    using quotes: Quotes
  ): Expr[(Int, Any)] = {
    import quotes.reflect.*

    val caseExprList = exprOfListToListOfExpr(caseListExpr)

    val allCaseDefLists = caseExprList.reverse.zipWithIndex.flatMap { case (caseExpr, idx) =>
      caseExpr.asTerm match {
        case Block(List(DefDef(_, _, _, Some(Match(_, caseDefList)))), _) => {
          caseDefList.map { caseDef =>
            val idxExpr = Expr.apply(idx)
            val newRhsExpr = '{ val res = ${caseDef.rhs.asExprOf[Any]}; ($idxExpr, res)}
            CaseDef.copy(caseDef)(caseDef.pattern, caseDef.guard, newRhsExpr.asTerm)
          }
        }
        case _ => report.errorAndAbort("Macro expansion failed, please use `handleCase` with annonymous partial function")
      }
    }

    Match(iExpr.asTerm, allCaseDefLists.map(_.changeOwner(Symbol.spliceOwner))).asExprOf[(Int, Any)]
  }

}