Scala Typesafe Context

ctx-scala-2.scala

final case class Ctx[T <: Ctx.Key[_]] private (private val map: Map[Ctx.Key[_], Any])

object Ctx {
  trait Key[T] { type Value = T; override def toString: String = getClass.getTypeName.replace("$", "") }

  implicit final class CtxExtension[Self <: Ctx[_]](private val self: Self) extends AnyVal {
    def put[K <: Key[_]](key: K)(value: key.Value): Self with Ctx[K] =
      Ctx(self.map + (key -> value)).asInstanceOf[Self with Ctx[K]]

    def get[K <: Key[_]](key: K)(implicit ev: Self <:< Ctx[K]): key.Value = self.map.get(key) match {
      case Some(value) => value.asInstanceOf[key.Value]
      case None => throw new IllegalArgumentException
    }

    def in[R <: Ctx[_]](fun: Self => Self with R): Self with R = fun(self)
  }

  val empty: Ctx[Nothing] = Ctx(Map())
}

object Application extends App {
  object Span    extends Ctx.Key[String]
  object FlowId  extends Ctx.Key[Int]
  object Effect  extends Ctx.Key[Unit]
  object Unknown extends Ctx.Key[Nothing]

  def logic1
  [T <: Ctx[Span.type] with Ctx[FlowId.type]]
  (ctx: T)
  : T = {
    val span: String = ctx.get(Span)
    val flowId: Int = ctx.get(FlowId)
    // Will not compile since T doesn't define Unknown requirement
    // println(ctx.get(Unknown))
    ctx.put(Span)(s"$span[$flowId]")
  }

  def logic2
  [T <: Ctx[Span.type]]
  (ctx: T)
  : T with Ctx[Effect.type] = {
    println(s"Performing action for span = ${ctx.get(Span)}")
    ctx.put(Effect)(())
  }

  def logic3
  [T <: Ctx[FlowId.type] with Ctx[Span.type]
                         with Ctx[Effect.type]]
  (ctx: T)
  : T = {
    println(s"After effect[flowId=${ctx.get(FlowId)}]: $ctx")
    ctx
  }

  val ctx = Ctx.empty
    .put(Span)("span")
    .put(FlowId)(1)

  val result = ctx
    .in(logic1)
    .in(logic2)
    .in(logic3)

  println(result)
}

ctx-scala-3.scala

final case class Ctx[T <: Ctx.Key[?]] private(private val map: Map[Ctx.Key[?], Any] = Map())

object Ctx:
  trait Key[T]:
    type Value = T
    override val toString: String = getClass.getTypeName.replace("$", "")

  val empty: Ctx[Nothing] = Ctx()

  extension [Self <: Ctx[?]] (self: Self)
    def put[B <: Key[?]](key: B, value: key.Value): Self & Ctx[B] =
      Ctx(self.map + (key -> value)).asInstanceOf[Self & Ctx[B]]

    def get[B <: Key[?]](key: B)(using Self <:< Ctx[B]): key.Value = self.map.get(key) match {
      case Some(value) => value.asInstanceOf[key.Value]
      case None        => throw new IllegalArgumentException
    }

// Usage

object Span    extends Ctx.Key[String]
object FlowId  extends Ctx.Key[Int]
object Effect  extends Ctx.Key[Unit]
object Unknown extends Ctx.Key[Nothing]

type Logic1Req = Ctx[Span.type] & Ctx[FlowId.type]

def logic1[T <: Logic1Req](ctx: T): T =
  val span: String = ctx.get(Span)
  val flowId: Int = ctx.get(FlowId)
  // Will not compile since Logic1Req doesn't define Unknown requirement
  // println(ctx.get(Unknown))
  ctx.put(Span, s"$span[$flowId]")

type Logic2Req = Ctx[Span.type]
type Logic2Out = Ctx[Effect.type]

def logic2[T <: Logic2Req](ctx: T): T & Logic2Out =
  println(s"Performing action for span = ${ctx.get(Span)}")
  ctx.put(Effect, ())

type Logic3Req = Ctx[FlowId.type] & Ctx[Effect.type]

def logic3[T <: Logic3Req](ctx: T): T =
  println(s"After effect[flowId=${ctx.get(FlowId)}]: $ctx")
  ctx

@main def hello: Unit =
  val ctx = Ctx.empty
    .put(Span, "span")
    .put(FlowId, 1)
  // Will not compile since braking input constraints defined by logic3
  // val pipeline = logic1.andThen(logic3).andThen(logic2)
  val pipeline = logic1.andThen(logic2).andThen(logic3)
  pipeline(ctx)