Handle useStateBy and withPropsChildren

library/coreGeneric/src/main/scala-2/japgolly/scalajs/react/hooks/HookMacros.scala

@@ -27,23 +27,74 @@ class HookMacros(val c: Context) extends MacroUtils {
 
   private implicit def autoTagToType[A](t: c.WeakTypeTag[A]): Type = t.tpe
 
-  private def Box         : Tree = q"_root_.japgolly.scalajs.react.internal.Box"
-  private def Box(t: Type): Type = appliedType(c.typeOf[Box[_]], t)
-  private def Hooks       : Tree = q"_root_.japgolly.scalajs.react.hooks.Hooks"
-  private def JsFn        : Tree = q"_root_.japgolly.scalajs.react.component.JsFn"
-  private def React       : Tree = q"_root_.japgolly.scalajs.react.facade.React"
-  private def ScalaFn     : Tree = q"_root_.japgolly.scalajs.react.component.ScalaFn"
-  private def withHooks          = "withHooks"
+  private def Box          : Tree = q"_root_.japgolly.scalajs.react.internal.Box"
+  private def Box(t: Type) : Type = appliedType(c.typeOf[Box[_]], t)
+  private def HookCtx      : Tree = q"_root_.japgolly.scalajs.react.hooks.HookCtx"
+  private def Hooks        : Tree = q"_root_.japgolly.scalajs.react.hooks.Hooks"
+  private def JsFn         : Tree = q"_root_.japgolly.scalajs.react.component.JsFn"
+  private def PropsChildren: Tree = q"_root_.japgolly.scalajs.react.PropsChildren"
+  private def React        : Tree = q"_root_.japgolly.scalajs.react.facade.React"
+  private def ScalaFn      : Tree = q"_root_.japgolly.scalajs.react.component.ScalaFn"
+  private def withHooks           = "withHooks"
 
-  case class HookDefn(steps: List[HookStep])
-  case class HookStep(name: String, targs: List[Tree], args: List[List[Tree]])
+  private case class HookDefn(steps: List[HookStep])
+
+  private case class HookStep(name: String, targs: List[Tree], args: List[List[Tree]])
+
+  private class HookRewriter(props: Tree, initChildren: Tree, propsChildren: Tree) {
+    private var stmts = Vector.empty[Tree]
+    private var hooks = List.empty[Ident]
+    private var _hookCount = 0
+    private var _usesChildren = false
+
+    def usesChildren() =
+      _usesChildren
+
+    def useChildren(): Unit = {
+      _usesChildren = true
+      this += initChildren
+    }
+
+    def +=(stmt: Tree): Unit =
+      stmts :+= stmt
+
+    def hookCount(): Int =
+      _hookCount
+
+    def nextHookName(suffix: String = ""): TermName =
+      TermName("hook" + (hookCount() + 1) + suffix)
+
+    def registerHook(h: TermName): Unit = {
+      hooks :+= Ident(h)
+      _hookCount += 1
+    }
+
+    def args(): List[Tree] =
+      if (usesChildren())
+        props :: propsChildren :: hooks
+      else
+        props :: hooks
+
+    def ctxArg(): Tree = {
+      val hookCtxObj = if (usesChildren()) q"$HookCtx.withChildren" else HookCtx
+      val create = Apply(hookCtxObj, args())
+      val name = nextHookName("_ctx")
+      this += q"val $name = $create"
+      Ident(name)
+    }
+
+    def wrap(body: Tree): Tree =
+      q"..$stmts; $body"
+  }
+
+  // -------------------------------------------------------------------------------------------------------------------
 
   def render[P, C <: Children, Ctx, CtxFn[_], Step <: SubsequentStep[Ctx, CtxFn]]
             (f: c.Tree)(step: c.Tree, s: c.Tree)
             (implicit P: c.WeakTypeTag[P], C: c.WeakTypeTag[C]): c.Tree = {
 
     implicit val log = MacroLogger()
-    // log.enabled = showCode(c.macroApplication).contains("counter.value")
+    log.enabled = showCode(c.macroApplication).contains("DEBUG") // TODO: DELETE
     log.header()
     log("macroApplication", showRaw(c.macroApplication))
 
@@ -103,72 +154,90 @@ class HookMacros(val c: Context) extends MacroUtils {
           Left(() => "Don't know how to parse " + showRaw(tree))
       }
 
-  private type RenderInliner = (Tree, Init) => Tree
-
-  private def inlineHookDefn(h: HookDefn)(implicit log: MacroLogger): Either[() => String, RenderInliner] = {
-    val init = new Init("hook" + _, lazyVals = false)
+  private def inlineHookDefn(h: HookDefn)(implicit log: MacroLogger): Either[() => String, HookRewriter => Tree] = {
     val it = h.steps.iterator
-    var stepId = 0
     var renderStep: HookStep = null
-    var hooks = List.empty[TermName]
+    var hooks = Vector.empty[HookRewriter => TermName]
+    var withPropsChildren = false
     while (it.hasNext) {
       val step = it.next()
       if (it.hasNext) {
-        stepId += 1
-        inlineHookStep(stepId, step, init) match {
-          case Right(termName) => hooks ::= termName
-          case Left(e) => return Left(e)
-        }
+        if (hooks.isEmpty && step.name == "withPropsChildren")
+          withPropsChildren = true
+        else
+          inlineHookStep(step) match {
+            case Right(h) => hooks :+= h
+            case Left(e)  => return Left(e)
+          }
       } else
         renderStep = step
     }
-    hooks = hooks.reverse
 
-    hookRenderInliner(renderStep, hooks.map(Ident(_))).map { f =>
-      (props, init2) => {
-        init2 ++= init.stmts
-        f(props, init2)
-      }
+    hookRenderInliner(renderStep).map { buildRender => b =>
+      if (withPropsChildren)
+        b.useChildren()
+      for (h <- hooks)
+        b registerHook h(b)
+      buildRender(b)
     }
   }
 
-  private def inlineHookStep(stepId: Int, step: HookStep, init: Init)(implicit log: MacroLogger): Either[() => String, TermName] = {
+  private def inlineHookStep(step: HookStep)(implicit log: MacroLogger): Either[() => String, HookRewriter => TermName] = {
     log("inlineHookStep." + step.name, step)
+
+    def useState(b: HookRewriter, tpe: Tree, body: Tree) = {
+      val rawName = b.nextHookName("_raw")
+      val name    = b.nextHookName()
+      b += q"val $rawName = $React.useStateFn(() => $Box[$tpe]($body))"
+      b += q"val $name = $Hooks.UseState.fromJsBoxed[$tpe]($rawName)"
+      name
+    }
+
     step.name match {
+
       case "useState" =>
-        val stateType = step.targs.head
-        val arg       = step.args.head.head
-        val rawName   = TermName("hook" + stepId + "_raw")
-        val name      = TermName("hook" + stepId)
-        init += q"val $rawName = $React.useStateFn(() => $Box[$stateType]($arg))"
-        init += q"val $name = $Hooks.UseState.fromJsBoxed[$stateType]($rawName)"
-        Right(name)
+        val targ = step.targs.head
+        val arg  = step.args.head.head
+        Right(useState(_, targ, arg))
+
+      case "useStateBy" =>
+        val targ = step.targs.head
+        val arg  = step.args.head.head
+        arg match {
+          case f@ Function(params, _) =>
+            if (params.sizeIs == 1)
+              Right { b =>
+                val ctxArg = b.ctxArg()
+                useState(b, targ, call(f, ctxArg :: Nil))
+              }
+            else
+              Right(b => useState(b, targ, call(f, b.args())))
+
+          case _ =>
+            Left(() => s"Expected a function, found: ${showRaw(arg)}")
+        }
 
       case _ =>
         Left(() => s"Inlining of hook method '${step.name}' not yet supported.")
     }
   }
 
-  private def hookRenderInliner(step: HookStep, hooks: List[Tree])(implicit log: MacroLogger): Either[() => String, RenderInliner] = {
+  private def hookRenderInliner(step: HookStep)(implicit log: MacroLogger): Either[() => String, HookRewriter => Tree] = {
     log("inlineHookRender." + step.name, step)
     step.name match {
       case "render" =>
         @nowarn("msg=exhaustive") val List(List(renderFn), _) = step.args
-        Right { (props, _) =>
-          val args = props :: hooks
-          Apply(Select(renderFn, TermName("apply")), args)
-        }
+        Right(b => call(renderFn, b.args()))
 
       case _ =>
         Left(() => s"Inlining of hook render method '${step.name}' not yet supported.")
     }
   }
 
-  private def inlineHookRawComponent[P](renderInliner: RenderInliner)(implicit P: c.WeakTypeTag[P]): Tree = {
-    val props_unbox = q"props.unbox"
-    val init        = new Init("_i" + _)
-    val render1     = renderInliner(props_unbox, init)
-    val render2     = init.wrap(q"$render1.rawNode")
+  private def inlineHookRawComponent[P](rewrite: HookRewriter => Tree)(implicit P: c.WeakTypeTag[P]): Tree = {
+    val b       = new HookRewriter(q"props.unbox", q"val children = $PropsChildren.fromRawProps(props)", q"children")
+    val render1 = rewrite(b)
+    val render2 = b.wrap(q"$render1.rawNode")
     q"(props => $render2): $JsFn.RawComponent[${Box(P)}]"
   }
 
@@ -178,4 +247,35 @@ class HookMacros(val c: Context) extends MacroUtils {
       $ScalaFn.fromBoxed($JsFn.fromJsFn[${Box(P)}, $C](rawComponent)($summoner))
     """)
   }
+
+  // -------------------------------------------------------------------------------------------------------------------
+
+  private def call(function: Tree, args: List[Tree]): Tree = {
+    import internal._
+
+    function match {
+      case Function(params, body) =>
+
+        // From scala/test/files/run/macro-range/Common_1.scala
+        class TreeSubstituter(from: List[Symbol], to: List[Tree]) extends Transformer {
+          override def transform(tree: Tree): Tree = tree match {
+            case Ident(_) =>
+              def subst(from: List[Symbol], to: List[Tree]): Tree =
+                if (from.isEmpty) tree
+                else if (tree.symbol == from.head) to.head.duplicate
+                else subst(from.tail, to.tail);
+              subst(from, to)
+            case _ =>
+              val tree1 = super.transform(tree)
+              if (tree1 ne tree) setType(tree1, null)
+              tree1
+          }
+        }
+        val t = new TreeSubstituter(params.map(_.symbol), args)
+        t.transform(body)
+
+      case _ =>
+        Apply(Select(function, TermName("apply")), args)
+    }
+  }
 }

library/coreGeneric/src/main/scala/japgolly/scalajs/react/internal/Box.scala

@@ -9,7 +9,7 @@ trait Box[+A] extends js.Object {
 }
 
 object Box {
-  @inline def apply[A](value: A): Box[A] =
+  def apply[A](value: A): Box[A] =
     js.Dynamic.literal(a = value.asInstanceOf[js.Any]).asInstanceOf[Box[A]]
 
   val Unit: Box[Unit] =

library/coreGeneric/src/main/scala/japgolly/scalajs/react/internal/MacroLogger.scala

@@ -69,4 +69,15 @@ class MacroLogger {
 
   def apply(name: => Any, value: => Any)(implicit l: Line): Unit =
     apply(s"$YELLOW$name:$RESET $value")
+
+  def all(name: => Any, values: => Iterable[Any])(implicit l: Line): Unit =
+    if (enabled) {
+      val vs = values
+      val total = vs.size
+      var i = 0
+      for (v <- vs) {
+        i += 1
+        apply(s"$name [$i/$total]", v)
+      }
+    }
 }

library/testEmissions/js/src/main/scala/japgolly/scalajs/react/test/emissions/HooksWithChildren.scala

@@ -0,0 +1,19 @@
+package japgolly.scalajs.react.test.emissions
+
+import japgolly.scalajs.react._
+import japgolly.scalajs.react.vdom.html_<^._
+
+object HooksWithChildren {
+
+  val Component = ScalaFnComponent.withHooks[Int]
+    .withPropsChildren
+    .useState(123)
+    .render { (p, c, s1) =>
+      val sum = p + s1.value + c.count
+      <.button(
+        "Sum = ", sum,
+        ^.onClick --> s1.modState(_ + 1),
+        c
+      )
+    }
+}

library/testEmissions/js/src/main/scala/japgolly/scalajs/react/test/emissions/Main.scala

@@ -14,7 +14,9 @@ object Main {
 
   private val Component = ScalaFnComponent[Unit] { _ =>
     <.div(
-      UseState.Component(),
+      HooksWithChildren.Component(0)(<.div),
+      NoHooksWithChildren.Component(0)(<.div),
+      UseState.Component(0),
     )
   }
 }

library/testEmissions/js/src/main/scala/japgolly/scalajs/react/test/emissions/NoHooksWithChildren.scala

@@ -0,0 +1,14 @@
+package japgolly.scalajs.react.test.emissions
+
+import japgolly.scalajs.react._
+import japgolly.scalajs.react.vdom.html_<^._
+
+object NoHooksWithChildren {
+
+  val Component = ScalaFnComponent.withHooks[Int]
+    .withPropsChildren
+    .render { (p, c) =>
+      val sum = p + c.count
+      <.div("DEBUG = ", sum)
+    }
+}

library/testEmissions/js/src/main/scala/japgolly/scalajs/react/test/emissions/UseState.scala

@@ -5,12 +5,15 @@ import japgolly.scalajs.react.vdom.html_<^._
 
 object UseState {
 
-  val Component = ScalaFnComponent.withHooks[Unit]
+  val Component = ScalaFnComponent.withHooks[Int]
     .useState(123)
-    .render { (_, s) =>
+    .useStateBy((p, s1) => p + s1.value)
+    .useStateBy($ => $.props + $.hook1.value + $.hook2.value)
+    .render { (_, s1, s2, s3) =>
+      val sum = s1.value + s2.value + s3.value
       <.button(
-        "Count is ", s.value,
-        ^.onClick --> s.modState(_ + 1),
+        "Sum = ", sum,
+        ^.onClick --> s1.modState(_ + 1),
       )
     }
 }

library/testEmissions/jvm/src/test/resources/rr-js/fn-in.js

@@ -0,0 +1,7 @@
+import React from 'react'
+
+function App() {
+  return React.createElement("div", {}, "count is: 0");
+}
+
+export default App

library/testEmissions/jvm/src/test/resources/rr-js/fn-out.js

@@ -0,0 +1,12 @@
+import React from 'react';
+
+function App() {
+  return React.createElement("div", {}, "count is: 0");
+}
+
+_c = App;
+export default App;
+
+var _c;
+
+$RefreshReg$(_c, "App");