Slight code cleanup

2 files changed, 13 insertions, 9 deletions

diff --git a/core/shared/src/main/scala/magnolia.scala b/core/shared/src/main/scala/magnolia.scala
index 32a9aa6..e0ebdbb 100644
--- a/core/shared/src/main/scala/magnolia.scala
+++ b/core/shared/src/main/scala/magnolia.scala
@@ -96,7 +96,7 @@ object Magnolia {
 
     def checkMethod(termName: String, category: String, expected: String): Unit = {
       val term = TermName(termName)
-      val combineClass = c.prefix.tree.tpe.baseClasses
+      val combineClass = prefixType.baseClasses
         .find { cls =>
           cls.asType.toType.decl(term) != NoSymbol
         }
@@ -268,6 +268,7 @@ object Magnolia {
             ..$preAssignments
             val $paramsVal: $scalaPkg.Array[$magnoliaPkg.Param[$typeConstructor, $genericType]] =
               new $scalaPkg.Array(${assignments.length})
+            
             ..$assignments
 
             $typeNameDef
@@ -282,11 +283,14 @@ object Magnolia {
                   val msg = "`" + $typeName.full + "` has " + $paramsVal.length + " fields, not " + $fieldValues.size
                   throw new java.lang.IllegalArgumentException(msg)
                 }
+                
                 new $genericType(..${caseParams.zipWithIndex.map {
-          case (typeclass, idx) =>
-            val arg = q"$fieldValues($idx).asInstanceOf[${typeclass.paramType}]"
-            if (typeclass.repeated) q"$arg: _*" else arg
-        }})}))
+                  case (typeclass, idx) =>
+                    val arg = q"$fieldValues($idx).asInstanceOf[${typeclass.paramType}]"
+                    if (typeclass.repeated) q"$arg: _*" else arg
+                 }})
+              }
+            ))
           }""")
       } else if (isSealedTrait) {
         val genericSubtypes = classType.get.knownDirectSubclasses.to[List]
@@ -308,11 +312,11 @@ object Magnolia {
 
         val subtypesVal: TermName = TermName(c.freshName("subtypes"))
 
-        val typeclasses = for (subType <- subtypes) yield {
+        val typeclasses = subtypes.map { subtype =>
           val path = CoproductType(genericType.toString)
           val frame = stack.Frame(path, resultType, assignedName)
-          subType -> stack.recurse(frame, appliedType(typeConstructor, subType)) {
-            typeclassTree(subType, typeConstructor)
+          subtype -> stack.recurse(frame, appliedType(typeConstructor, subtype)) {
+            typeclassTree(subtype, typeConstructor)
           }
         }
 
diff --git a/tests/src/main/scala/tests.scala b/tests/src/main/scala/tests.scala
index a311291..f56afcb 100644
--- a/tests/src/main/scala/tests.scala
+++ b/tests/src/main/scala/tests.scala
@@ -243,7 +243,7 @@ object Tests extends TestApp {
       } catch {
         case NonFatal(e) => e.getMessage
       }
-    }.assert(_ == "Cannot patch value `Person(Bob,42)`, expected 2 fields but got 3")
+    }.assert(_ == "Cannot patch value `Person(Bob,42)`, expected 2 fields but got 3", _ => "test failed")
 
     test("throw on an illegal patch attempt with field type mismatch") {
       // these two implicits can be removed once https://github.com/propensive/magnolia/issues/58 is closed