Infer final expression returns in lambdas

internal/lang/expression_function_test.go

@@ -84,3 +84,59 @@ fun handle(request: *http.Request) = withContext(request.context()) {
 		}
 	}
 }
+
+func TestLambdaReturnsFinalExpression(t *testing.T) {
+	program, err := Parse(`package demo
+fun transform<T, R>(value: T, block: (T) -> R): R = block(value)
+fun answer(): Int = transform(21) { value -> value * 2 }`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	output, err := GenerateGo(program)
+	if err != nil {
+		t.Fatal(err)
+	}
+	code := string(output)
+	for _, expected := range []string{"func(value int) int", "return value * 2", "func answer() int"} {
+		if !strings.Contains(code, expected) {
+			t.Fatalf("missing %q:\n%s", expected, code)
+		}
+	}
+}
+
+func TestGoCallbackReturnsFinalExpression(t *testing.T) {
+	program, err := Parse(`package demo
+import sort
+fun order(values: MutableList<Int>) {
+    sort.slice(values) { left, right -> values[left] < values[right] }
+}`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	output, err := GenerateGo(program)
+	if err != nil {
+		t.Fatal(err)
+	}
+	code := string(output)
+	for _, expected := range []string{"func(left int, right int) bool", "return values[left] < values[right]"} {
+		if !strings.Contains(code, expected) {
+			t.Fatalf("missing %q:\n%s", expected, code)
+		}
+	}
+}
+
+func TestUnitLambdaKeepsFinalExpressionAsStatement(t *testing.T) {
+	program, err := Parse(`package demo
+fun execute(block: () -> Unit) = block()
+fun start() = execute { println("started") }`)
+	if err != nil {
+		t.Fatal(err)
+	}
+	output, err := GenerateGo(program)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if strings.Contains(string(output), `return fmt.Println("started")`) {
+		t.Fatalf("Unit lambda returned final expression:\n%s", output)
+	}
+}