Make `SealedTrait.dispatch` allocation-free

2 files changed, 16 insertions, 15 deletions

diff --git a/core/shared/src/main/scala/interface.scala b/core/shared/src/main/scala/interface.scala
index 193a6f9..a4be9ab 100644
--- a/core/shared/src/main/scala/interface.scala
+++ b/core/shared/src/main/scala/interface.scala
@@ -1,6 +1,7 @@
 package magnolia
 
 import language.higherKinds
+import scala.annotation.tailrec
 
 /** represents a subtype of a sealed trait
   *
@@ -157,12 +158,12 @@ final class SealedTrait[Typeclass[_], Type](val typeName: String,
     *                 matches
     *  @return  the result of applying the `handle` lambda to subtype of the sealed trait which
     *           matches the parameter `value` */
-  def dispatch[Return](value: Type)(handle: Subtype[Typeclass, Type] => Return): Return =
-    subtypes
-      .map { sub =>
-        sub.cast.andThen { v =>
-          handle(sub)
-        }
-      }
-      .reduce(_ orElse _)(value)
+  def dispatch[Return](value: Type)(handle: Subtype[Typeclass, Type] => Return): Return = {
+    @tailrec def rec(ix: Int): Return =
+      if (ix < subtypesArray.length) {
+        val sub = subtypesArray(ix)
+        if (sub.cast.isDefinedAt(value)) handle(sub) else rec(ix + 1)
+      } else throw new IllegalArgumentException(s"The given value `$value` is not a sub type of `$typeName`")
+    rec(0)
+  }
 }
diff --git a/core/shared/src/main/scala/magnolia.scala b/core/shared/src/main/scala/magnolia.scala
index f55dfa7..68c0f4f 100644
--- a/core/shared/src/main/scala/magnolia.scala
+++ b/core/shared/src/main/scala/magnolia.scala
@@ -1,6 +1,7 @@
 package magnolia
 
-import scala.reflect._, macros._
+import scala.reflect._
+import macros._
 import scala.collection.immutable.ListMap
 import language.existentials
 import language.higherKinds
@@ -505,15 +506,14 @@ object Magnolia {
     *
     *  This method is intended to be called only from code generated by the Magnolia macro, and
     *  should not be called directly from users' code. */
-  def subtype[Tc[_], T, S <: T](name: String, tc: => Tc[S], isType: T => Boolean, asType: T => S) =
-    new Subtype[Tc, T] {
+  def subtype[Tc[_], T, S <: T](name: String, tc: => Tc[S], isType: T => Boolean, asType: T => S): Subtype[Tc, T] =
+    new Subtype[Tc, T] with PartialFunction[T, S] {
       type SType = S
       def label: String = name
       def typeclass: Tc[SType] = tc
-      def cast: PartialFunction[T, SType] = new PartialFunction[T, S] {
-        def isDefinedAt(t: T) = isType(t)
-        def apply(t: T): SType = asType(t)
-      }
+      def cast: PartialFunction[T, SType] = this
+      def isDefinedAt(t: T) = isType(t)
+      def apply(t: T): SType = asType(t)
     }
 
   /** constructs a new [[Param]] instance