Merge pull request #63 from sirthias/master

Make `SealedTrait.dispatch` allocation-free

2 files changed, 33 insertions, 31 deletions

diff --git a/core/shared/src/main/scala/interface.scala b/core/shared/src/main/scala/interface.scala
index 193a6f9..a82a54c 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
   *
@@ -128,7 +129,7 @@ abstract class CaseClass[Typeclass[_], Type] private[magnolia] (
     *
     *  For efficiency, this sequence is implemented by an `Array`, but upcast to a
     *  [[scala.collection.Seq]] to hide the mutable collection API. */
-  def parameters: Seq[Param[Typeclass, Type]] = parametersArray
+  final def parameters: Seq[Param[Typeclass, Type]] = parametersArray
 }
 
 /** represents a sealed trait and the context required to construct a new typeclass instance
@@ -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..2e7dae8 100644
--- a/core/shared/src/main/scala/magnolia.scala
+++ b/core/shared/src/main/scala/magnolia.scala
@@ -1,6 +1,8 @@
 package magnolia
 
-import scala.reflect._, macros._
+import scala.util.control.NonFatal
+import scala.reflect._
+import macros._
 import scala.collection.immutable.ListMap
 import language.existentials
 import language.higherKinds
@@ -111,7 +113,7 @@ object Magnolia {
           )
         }
       val firstParamBlock = combineClass.asType.toType.decl(term).asTerm.asMethod.paramLists.head
-      if (firstParamBlock.length != 1)
+      if (firstParamBlock.lengthCompare(1) != 0)
         c.abort(c.enclosingPosition,
                 s"magnolia: the method `combine` should take a single parameter of type $expected")
     }
@@ -123,9 +125,9 @@ object Magnolia {
     def findType(key: Type): Option[TermName] =
       recursionStack(c.enclosingPosition).frames.find(_.genericType == key).map(_.termName(c))
 
-    case class Typeclass(typ: c.Type, tree: c.Tree)
+    final case class Typeclass(typ: c.Type, tree: c.Tree)
 
-    def recurse[T](path: TypePath, key: Type, value: TermName)(fn: => T): Option[T] = {
+    def recurse[A](path: TypePath, key: Type, value: TermName)(fn: => A): Option[A] = {
       val oldRecursionStack = recursionStack.get(c.enclosingPosition)
       recursionStack = recursionStack.updated(
         c.enclosingPosition,
@@ -135,7 +137,7 @@ object Magnolia {
       )
 
       try Some(fn)
-      catch { case e: Exception => None } finally {
+      catch { case NonFatal(_) => None } finally {
         val currentStack = recursionStack(c.enclosingPosition)
         recursionStack = recursionStack.updated(c.enclosingPosition, currentStack.pop())
       }
@@ -292,7 +294,7 @@ object Magnolia {
             m.asMethod
         }
 
-        case class CaseParam(sym: c.universe.MethodSymbol,
+        final case class CaseParam(sym: c.universe.MethodSymbol,
                              repeated: Boolean,
                              typeclass: c.Tree,
                              paramType: c.Type,
@@ -505,15 +507,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
@@ -524,7 +525,7 @@ object Magnolia {
                          isRepeated: Boolean,
                          typeclassParam: Tc[P],
                          defaultVal: => Option[P],
-                         deref: T => P) = new Param[Tc, T] {
+                         deref: T => P): Param[Tc, T] = new Param[Tc, T] {
     type PType = P
     def label: String = name
     def repeated: Boolean = isRepeated
@@ -541,31 +542,31 @@ object Magnolia {
                           obj: Boolean,
                           valClass: Boolean,
                           params: Array[Param[Tc, T]],
-                          constructor: (Param[Tc, T] => Any) => T) =
+                          constructor: (Param[Tc, T] => Any) => T): CaseClass[Tc, T] =
     new CaseClass[Tc, T](name, obj, valClass, params) {
       def construct[R](param: Param[Tc, T] => R): T = constructor(param)
     }
 }
 
-private[magnolia] case class DirectlyReentrantException()
+private[magnolia] final case class DirectlyReentrantException()
     extends Exception("attempt to recurse directly")
 
 private[magnolia] object Deferred { def apply[T](method: String): T = ??? }
 
 private[magnolia] object CompileTimeState {
 
-  sealed class TypePath(path: String) { override def toString = path }
-  case class CoproductType(typeName: String) extends TypePath(s"coproduct type $typeName")
+  sealed abstract class TypePath(path: String) { override def toString = path }
+  final case class CoproductType(typeName: String) extends TypePath(s"coproduct type $typeName")
 
-  case class ProductType(paramName: String, typeName: String)
+  final case class ProductType(paramName: String, typeName: String)
       extends TypePath(s"parameter '$paramName' of product type $typeName")
 
-  case class ChainedImplicit(typeName: String)
+  final case class ChainedImplicit(typeName: String)
       extends TypePath(s"chained implicit of type $typeName")
 
-  case class ImplicitNotFound(genericType: String, path: List[TypePath])
+  final case class ImplicitNotFound(genericType: String, path: List[TypePath])
 
-  case class Stack(cache: Map[whitebox.Context#Type, Option[whitebox.Context#Tree]],
+  final case class Stack(cache: Map[whitebox.Context#Type, Option[whitebox.Context#Tree]],
                    frames: List[Frame],
                    errors: List[ImplicitNotFound]) {
 
@@ -583,7 +584,7 @@ private[magnolia] object CompileTimeState {
     def pop(): Stack = Stack(cache, frames.tail, errors)
   }
 
-  case class Frame(path: TypePath,
+  final case class Frame(path: TypePath,
                    genericType: whitebox.Context#Type,
                    term: whitebox.Context#TermName) {
     def termName(c: whitebox.Context): c.TermName = term.asInstanceOf[c.TermName]