Fix issue with context receive in lambdas

Summary:
Fixes issue reported by at #471

Formatting might be a bit odd, but the goal of this diff is solely to stop of the issue with running ktfmt.

Reviewed By: strulovich

Differential Revision: D58194198

fbshipit-source-id: 51123c1b3fef65c2c0803689c9173cb7c50cb41f

Author: Nivaldo Bondança

core/src/main/java/com/facebook/ktfmt/format/KotlinInputAstVisitor.kt

@@ -2290,6 +2290,9 @@ class KotlinInputAstVisitor(
 
   override fun visitFunctionType(type: KtFunctionType) {
     builder.sync(type)
+
+    type.contextReceiverList?.let { visitContextReceiverList(it) }
+
     val receiver = type.receiver
     if (receiver != null) {
       visit(receiver)

core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt

@@ -7199,6 +7199,13 @@ class FormatterTest {
       |  context(SomethingElse)
       |
       |  private class NestedClass {}
+      |
+      |  fun <T> testSuspend(
+      |    mock: T,
+      |    block: suspend context(SomeContext) T.() -> Unit,
+      |  ) = startCoroutine {
+      |    T.block()
+      |  }
       |}
       |"""
             .trimMargin()
@@ -7216,6 +7223,13 @@ class FormatterTest {
       |
       |  context(SomethingElse)
       |  private class NestedClass {}
+      |
+      |  fun <T> testSuspend(
+      |      mock: T,
+      |      block:
+      |          suspend context(SomeContext)
+      |          T.() -> Unit,
+      |  ) = startCoroutine { T.block() }
       |}
       |"""
             .trimMargin()