Fix up owner chains to allow `Applyable#apply()` calls to work within lambdas

2 files changed, 24 insertions, 6 deletions

diff --git a/core/src/main/scala/forge/define/Applicative.scala b/core/src/main/scala/forge/define/Applicative.scala
index ac7c6de5..23b5f776 100644
--- a/core/src/main/scala/forge/define/Applicative.scala
+++ b/core/src/main/scala/forge/define/Applicative.scala
@@ -52,7 +52,7 @@ object Applicative {
     import c.universe._
     def rec(t: Tree): Iterator[c.Tree] = Iterator(t) ++ t.children.flatMap(rec(_))
 
-    val bound = collection.mutable.Buffer.empty[(c.Tree, Symbol)]
+    val bound = collection.mutable.Buffer.empty[(c.Tree, ValDef)]
     val targetApplySym = typeOf[Applyable[_]].member(TermName("apply"))
 
     // Derived from @olafurpg's
@@ -73,20 +73,26 @@ object Applicative {
           )
         }
         val tempName = c.freshName(TermName("tmp"))
-        val tempSym = c.internal.newTermSymbol(api.currentOwner, tempName)
+        val tempSym = c.internal.newTermSymbol(c.internal.enclosingOwner, tempName)
         c.internal.setInfo(tempSym, t.tpe)
         val tempIdent = Ident(tempSym)
         c.internal.setType(tempIdent, t.tpe)
-        bound.append((q"${c.prefix}.underlying($fun)", tempSym))
+        c.internal.setFlag(tempSym, (1L << 44).asInstanceOf[c.universe.FlagSet])
+        bound.append((q"${c.prefix}.underlying($fun)", c.internal.valDef(tempSym)))
         tempIdent
       case (t, api) => api.default(t)
     }
 
-    val (exprs, symbols) = bound.unzip
+    val (exprs, bindings) = bound.unzip
 
-    val bindings = symbols.map(c.internal.valDef(_))
 
-    wrapCached(c)(q"${c.prefix}.zipMap(..$exprs){ (..$bindings) => $transformed }")
+    val callback = c.typecheck(q"(..$bindings) => $transformed ")
+
+    val res = q"${c.prefix}.zipMap(..$exprs){ $callback }"
+
+    c.internal.changeOwner(transformed, c.internal.enclosingOwner, callback.symbol)
+
+    wrapCached(c)(res)
   }
   def wrapCached[M[_], T](c: Context)(t: c.Tree) = {
     import c.universe._
diff --git a/core/src/test/scala/forge/ApplicativeTests.scala b/core/src/test/scala/forge/ApplicativeTests.scala
index 000554ef..621dc679 100644
--- a/core/src/test/scala/forge/ApplicativeTests.scala
+++ b/core/src/test/scala/forge/ApplicativeTests.scala
@@ -94,6 +94,18 @@ object ApplicativeTests extends TestSuite {
         counter.value == 1
       )
     }
+    'evaluationsInsideLambdasWork - {
+      // This required some fiddling with owner chains inside the macro to get
+      // working, so ensure it doesn't regress
+      val counter = new Counter()
+      def up = Opt{ "hello" + counter() }
+      val down1 = Opt{ (() => up())() }
+      val down2 = Opt{ Seq(1, 2, 3).map(n => up() * n) }
+      assert(
+        down1 == Some("hello1"),
+        down2 == Some(Seq("hello2", "hello2hello2", "hello2hello2hello2"))
+      )
+    }
   }
 }