add Local[F, Span[F]] instance where F is Kleisli[F, Span[F], *]

build.sbt

@@ -110,7 +110,7 @@ lazy val core = crossProject(JSPlatform, JVMPlatform, NativePlatform)
 
 lazy val coreTests = crossProject(JSPlatform, JVMPlatform, NativePlatform)
   .in(file("modules/core-tests"))
-  .dependsOn(core, testkit)
+  .dependsOn(core, mtl, testkit)
   .enablePlugins(AutomateHeaderPlugin, NoPublishPlugin)
   .settings(commonSettings)
 
@@ -283,7 +283,10 @@ lazy val mtl = crossProject(JSPlatform, JVMPlatform, NativePlatform)
     name := "natchez-mtl",
     description := "cats-mtl bindings for Natchez.",
     libraryDependencies ++= Seq(
-      "org.typelevel" %%% "cats-mtl" % "1.3.0"
+      "org.typelevel" %%% "cats-mtl" % "1.3.0",
+      "org.typelevel" %%% "cats-mtl-laws" % "1.3.0" % Test,
+      "org.typelevel" %%% "discipline-munit" % "2.0.0-M3" % Test,
+      "org.typelevel" %%% "cats-effect-testkit" % "3.4.5" % Test
     )
   )
   .nativeSettings(

modules/core-tests/shared/src/test/scala/InMemorySuite.scala

@@ -4,56 +4,131 @@
 
 package natchez
 
+import cats.arrow.FunctionK
+import cats.{Applicative, ~>}
+import cats.syntax.all._
 import cats.data.Kleisli
-import cats.effect.{IO, MonadCancelThrow}
+import cats.effect.{Concurrent, IO, IOLocal, MonadCancelThrow}
+import cats.mtl.Local
 import munit.CatsEffectSuite
 import natchez.InMemory.Lineage.defaultRootName
+import natchez.mtl._
 
 trait InMemorySuite extends CatsEffectSuite {
   type Lineage = InMemory.Lineage
   val Lineage = InMemory.Lineage
   type NatchezCommand = InMemory.NatchezCommand
   val NatchezCommand = InMemory.NatchezCommand
+}
 
-  trait TraceTest {
-    def program[F[_]: MonadCancelThrow: Trace]: F[Unit]
-    def expectedHistory: List[(Lineage, NatchezCommand)]
-  }
+object InMemorySuite {
+  trait TraceSuite extends InMemorySuite {
+    trait TraceTest {
+      def program[F[_]: MonadCancelThrow: Trace]: F[Unit]
+
+      def expectedHistory: List[(Lineage, NatchezCommand)]
+    }
+
+    def traceTest(name: String, tt: TraceTest): Unit = {
+      test(s"$name - Kleisli")(
+        testTraceKleisli(tt.program[Kleisli[IO, Span[IO], *]](implicitly, _), tt.expectedHistory)
+      )
+      test(s"$name - IOLocal")(testTraceIoLocal(tt.program[IO](implicitly, _), tt.expectedHistory))
+    }
 
-  def traceTest(name: String, tt: TraceTest): Unit = {
-    test(s"$name - Kleisli")(
-      testTraceKleisli(tt.program[Kleisli[IO, Span[IO], *]](implicitly, _), tt.expectedHistory)
+    def testTraceKleisli(
+        traceProgram: Trace[Kleisli[IO, Span[IO], *]] => Kleisli[IO, Span[IO], Unit],
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] = testTrace[Kleisli[IO, Span[IO], *]](
+      traceProgram,
+      root => IO.pure(Trace[Kleisli[IO, Span[IO], *]] -> (k => k.run(root))),
+      expectedHistory
     )
-    test(s"$name - IOLocal")(testTraceIoLocal(tt.program[IO](implicitly, _), tt.expectedHistory))
-  }
 
-  def testTraceKleisli(
-      traceProgram: Trace[Kleisli[IO, Span[IO], *]] => Kleisli[IO, Span[IO], Unit],
-      expectedHistory: List[(Lineage, NatchezCommand)]
-  ): IO[Unit] = testTrace[Kleisli[IO, Span[IO], *]](
-    traceProgram,
-    root => IO.pure(Trace[Kleisli[IO, Span[IO], *]] -> (k => k.run(root))),
-    expectedHistory
-  )
-
-  def testTraceIoLocal(
-      traceProgram: Trace[IO] => IO[Unit],
-      expectedHistory: List[(Lineage, NatchezCommand)]
-  ): IO[Unit] = testTrace[IO](traceProgram, Trace.ioTrace(_).map(_ -> identity), expectedHistory)
-
-  def testTrace[F[_]](
-      traceProgram: Trace[F] => F[Unit],
-      makeTraceAndResolver: Span[IO] => IO[(Trace[F], F[Unit] => IO[Unit])],
-      expectedHistory: List[(Lineage, NatchezCommand)]
-  ): IO[Unit] =
-    InMemory.EntryPoint.create[IO].flatMap { ep =>
-      val traced = ep.root(defaultRootName).use { r =>
-        makeTraceAndResolver(r).flatMap { case (traceInstance, resolve) =>
-          resolve(traceProgram(traceInstance))
+    def testTraceIoLocal(
+        traceProgram: Trace[IO] => IO[Unit],
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] = testTrace[IO](traceProgram, Trace.ioTrace(_).map(_ -> identity), expectedHistory)
+
+    def testTrace[F[_]](
+        traceProgram: Trace[F] => F[Unit],
+        makeTraceAndResolver: Span[IO] => IO[(Trace[F], F[Unit] => IO[Unit])],
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] =
+      InMemory.EntryPoint.create[IO].flatMap { ep =>
+        val traced = ep.root(defaultRootName).use { r =>
+          makeTraceAndResolver(r).flatMap { case (traceInstance, resolve) =>
+            resolve(traceProgram(traceInstance))
+          }
+        }
+        traced *> ep.ref.get.map { history =>
+          assertEquals(history.toList, expectedHistory)
         }
       }
-      traced *> ep.ref.get.map { history =>
-        assertEquals(history.toList, expectedHistory)
-      }
+  }
+
+  trait LocalSuite extends InMemorySuite {
+    type LocalProgram[F[_]] = (EntryPoint[F], Local[F, Span[F]]) => F[Unit]
+
+    trait LocalTest {
+      def program[F[_]: MonadCancelThrow](entryPoint: EntryPoint[F])(implicit
+          L: Local[F, Span[F]]
+      ): F[Unit]
+      def expectedHistory: List[(Lineage, NatchezCommand)]
+    }
+
+    def localTest(name: String, tt: LocalTest): Unit = {
+      test(s"$name - Kleisli")(
+        testLocalKleisli(tt.program[Kleisli[IO, Span[IO], *]](_)(implicitly, _), tt.expectedHistory)
+      )
+      test(s"$name - IOLocal")(
+        testLocalIoLocal(tt.program[IO](_)(implicitly, _), tt.expectedHistory)
+      )
     }
+
+    private def testLocalKleisli(
+        localProgram: LocalProgram[Kleisli[IO, Span[IO], *]],
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] = testProgramGivenEntryPoint[Kleisli[IO, Span[IO], *]](
+      localProgram(_, implicitly),
+      Kleisli.applyK(Span.noop[IO]),
+      expectedHistory
+    )
+
+    private def testLocalIoLocal(
+        localProgram: LocalProgram[IO],
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] =
+      testProgramGivenEntryPoint[IO](
+        ep =>
+          IOLocal(Span.noop[IO])
+            .map { ioLocal =>
+              new Local[IO, Span[IO]] {
+                override def local[A](fa: IO[A])(f: Span[IO] => Span[IO]): IO[A] =
+                  ioLocal.get.flatMap { initial =>
+                    ioLocal.set(f(initial)) >> fa.guarantee(ioLocal.set(initial))
+                  }
+
+                override def applicative: Applicative[IO] = implicitly
+
+                override def ask[E2 >: Span[IO]]: IO[E2] = ioLocal.get
+              }
+            }
+            .flatMap(localProgram(ep, _)),
+        FunctionK.id,
+        expectedHistory
+      )
+
+    private def testProgramGivenEntryPoint[F[_]: Concurrent](
+        localProgram: EntryPoint[F] => F[Unit],
+        fk: F ~> IO,
+        expectedHistory: List[(Lineage, NatchezCommand)]
+    ): IO[Unit] =
+      fk(InMemory.EntryPoint.create[F].flatMap { ep =>
+        localProgram(ep) *> ep.ref.get.map { history =>
+          assertEquals(history.toList, expectedHistory)
+        }
+      })
+  }
+
 }

modules/core-tests/shared/src/test/scala/LocalTraceSpec.scala

@@ -0,0 +1,33 @@
+// Copyright (c) 2019-2020 by Rob Norris and Contributors
+// This software is licensed under the MIT License (MIT).
+// For more information see LICENSE or https://opensource.org/licenses/MIT
+
+package natchez
+package mtl
+
+import cats.effect.{MonadCancelThrow, Trace => _}
+import cats.mtl._
+
+class LocalTraceSpec extends InMemorySuite.LocalSuite {
+
+  private def useTrace[F[_]: Trace]: F[Unit] = Trace[F].log("hello world")
+
+  localTest(
+    "Trace[F] should use root span generated by provided EntryPoint[F] via Local[F, Span[F]]",
+    new LocalTest {
+      override def program[F[_]: MonadCancelThrow](entryPoint: EntryPoint[F])(implicit
+          L: Local[F, Span[F]]
+      ): F[Unit] =
+        entryPoint
+          .root("my root")
+          .use(Local[F, Span[F]].scope(useTrace[F]))
+
+      override def expectedHistory: List[(Lineage, NatchezCommand)] = List(
+        Lineage.Root -> NatchezCommand
+          .CreateRootSpan("my root", Kernel(Map()), Span.Options.Defaults),
+        Lineage.Root("my root") -> NatchezCommand.LogEvent("hello world"),
+        Lineage.Root -> NatchezCommand.ReleaseRootSpan("my root")
+      )
+    }
+  )
+}

modules/core-tests/shared/src/test/scala/SpanCoalesceTest.scala

@@ -6,7 +6,7 @@ package natchez
 
 import cats.effect.MonadCancelThrow
 
-class SpanCoalesceTest extends InMemorySuite {
+class SpanCoalesceTest extends InMemorySuite.TraceSuite {
 
   traceTest(
     "suppress - nominal",

modules/core-tests/shared/src/test/scala/SpanPropagationTest.scala

@@ -7,7 +7,7 @@ package natchez
 import cats.effect.MonadCancelThrow
 import cats.syntax.all._
 
-class SpanPropagationTest extends InMemorySuite {
+class SpanPropagationTest extends InMemorySuite.TraceSuite {
 
   traceTest(
     "propagation",

modules/mtl/shared/src/main/scala/package.scala

@@ -4,13 +4,34 @@
 
 package natchez
 
-import cats.mtl.Local
+import cats._
+import cats.data.Kleisli
 import cats.effect.{Trace => _, _}
+import cats.mtl.{Local, MonadPartialOrder}
 
 package object mtl {
   implicit def natchezMtlTraceForLocal[F[_]](implicit
       ev: Local[F, Span[F]],
       eb: MonadCancel[F, Throwable]
   ): Trace[F] =
     new LocalTrace(ev)
+
+  implicit def localSpanForKleisli[F[_]](implicit
+      F: MonadCancel[F, _]
+  ): Local[Kleisli[F, Span[F], *], Span[Kleisli[F, Span[F], *]]] =
+    new Local[Kleisli[F, Span[F], *], Span[Kleisli[F, Span[F], *]]] {
+      override def local[A](
+          fa: Kleisli[F, Span[F], A]
+      )(f: Span[Kleisli[F, Span[F], *]] => Span[Kleisli[F, Span[F], *]]): Kleisli[F, Span[F], A] =
+        fa.local {
+          f.andThen(_.mapK(Kleisli.applyK(Span.noop[F])))
+            .compose(_.mapK(MonadPartialOrder[F, Kleisli[F, Span[F], *]]))
+        }
+
+      override def applicative: Applicative[Kleisli[F, Span[F], *]] =
+        MonadPartialOrder[F, Kleisli[F, Span[F], *]].monadG
+
+      override def ask[E2 >: Span[Kleisli[F, Span[F], *]]]: Kleisli[F, Span[F], E2] =
+        Kleisli.ask[F, Span[F]].map(_.mapK(MonadPartialOrder[F, Kleisli[F, Span[F], *]]))
+    }
 }

modules/mtl/shared/src/test/scala/LocalInstancesSpec.scala

@@ -0,0 +1,109 @@
+// Copyright (c) 2019-2020 by Rob Norris and Contributors
+// This software is licensed under the MIT License (MIT).
+// For more information see LICENSE or https://opensource.org/licenses/MIT
+
+package natchez
+package mtl
+
+import cats._
+import cats.data._
+import cats.effect.testkit.TestInstances
+import cats.effect.{Trace => _, _}
+import cats.laws.discipline.MiniInt.allValues
+import cats.laws.discipline.{ExhaustiveCheck, MiniInt}
+import cats.laws.discipline.arbitrary._
+import cats.laws.discipline.eq._
+import cats.mtl.laws.discipline.LocalTests
+import cats.syntax.all._
+import munit.DisciplineSuite
+import natchez.mtl.SeededSpan.{seedKey, spanToSeed}
+import org.scalacheck.rng.Seed
+import org.scalacheck.{Arbitrary, Cogen, Gen}
+import org.typelevel.ci.CIStringSyntax
+
+import java.net.URI
+import scala.util.Try
+
+class LocalInstancesSpec extends DisciplineSuite with TestInstances {
+  implicit val ticker: Ticker = Ticker()
+
+  private implicit def exhaustiveCheckSpan[F[_]: Applicative]: ExhaustiveCheck[Span[F]] =
+    ExhaustiveCheck.instance(allValues.map(new SeededSpan[F](_)))
+
+  private def genFromExhaustiveCheck[A: ExhaustiveCheck]: Gen[A] =
+    Gen.oneOf(ExhaustiveCheck[A].allValues)
+
+  private implicit def arbFromExhaustiveCheck[A: ExhaustiveCheck]: Arbitrary[A] =
+    Arbitrary(genFromExhaustiveCheck)
+
+  implicit def comonadIO: Comonad[IO] = new Comonad[IO] {
+    override def extract[A](x: IO[A]): A =
+      unsafeRun(x).fold(
+        throw new RuntimeException("canceled"),
+        throw _,
+        _.get
+      )
+
+    override def coflatMap[A, B](fa: IO[A])(f: IO[A] => B): IO[B] =
+      f(fa).pure[IO]
+
+    override def map[A, B](fa: IO[A])(f: A => B): IO[B] = fa.map(f)
+  }
+
+  private implicit def cogenSpan[F[_]: Comonad]: Cogen[Span[F]] =
+    Cogen(spanToSeed(_))
+
+  private implicit def cogenSpanK[F[_]: Comonad](implicit
+      F: MonadCancel[F, _]
+  ): Cogen[Span[Kleisli[F, Span[F], *]]] =
+    Cogen { (seed: Seed, span: Span[Kleisli[F, Span[F], *]]) =>
+      seed.reseed(
+        spanToSeed(
+          span.mapK(
+            Kleisli.applyK[F, Span[F]](
+              genFromExhaustiveCheck[Span[F]].apply(Gen.Parameters.default, seed).get
+            )
+          )
+        )
+      )
+    }
+
+  private implicit def eqKleisli[F[_], A: ExhaustiveCheck, B](implicit
+      ev: Eq[F[B]]
+  ): Eq[Kleisli[F, A, B]] =
+    Eq.by((x: Kleisli[F, A, B]) => x.run)
+
+  private implicit def kernelEq: Eq[Kernel] = Eq.by(_.toHeaders)
+  private implicit def spanKEq: Eq[Span[Kleisli[IO, Span[IO], *]]] = Eq.by(_.kernel)
+
+  checkAll(
+    "Local[Kleisli[F, Span[F], *], Span[Kleisli[F, Span[F], *]]]",
+    LocalTests[Kleisli[IO, Span[IO], *], Span[Kleisli[IO, Span[IO], *]]](localSpanForKleisli)
+      .local[Int, Int]
+  )
+}
+
+object SeededSpan {
+  val seedKey = ci"seed"
+
+  def spanToSeed[F[_]: Comonad](span: Span[F]): Long =
+    span.kernel.extract.toHeaders
+      .get(seedKey)
+      .flatMap { s =>
+        Try(java.lang.Long.parseUnsignedLong(s, 16)).toOption
+      }
+      .get
+}
+
+private class SeededSpan[F[_]: Applicative](seed: MiniInt) extends Span[F] {
+  override def put(fields: (String, TraceValue)*): F[Unit] = ().pure[F]
+  override def log(fields: (String, TraceValue)*): F[Unit] = ().pure[F]
+  override def log(event: String): F[Unit] = ().pure[F]
+  override def attachError(err: Throwable, fields: (String, TraceValue)*): F[Unit] = ().pure[F]
+  override def kernel: F[Kernel] = Kernel(Map(seedKey -> seed.toInt.toHexString)).pure[F]
+  override def span(name: String, options: Span.Options): Resource[F, Span[F]] =
+    Resource.pure(this)
+  override def traceId: F[Option[String]] = Seed(seed.toInt.toLong).toBase64.some.pure[F]
+  override def spanId: F[Option[String]] = Seed(seed.toInt.toLong).toBase64.some.pure[F]
+  override def traceUri: F[Option[URI]] = none[URI].pure[F]
+}