Infer final expression returns in lambdas

README.md

@@ -27,6 +27,7 @@ across package boundaries.
 - `if`, `else`, `while`, and `for (item in items)`
 - function calls
 - lambdas like `{ x: Int -> println(x) }` and `{ println(it) }`
+- Kotlin-style final-expression returns in value lambdas, for example `{ value -> value * 2 }`
 - `class` with primary-constructor fields and methods
 - `interface` with method signatures
 - Rust-style algebraic `enum` declarations with payload variants and exhaustive `match`

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)
+	}
+}

internal/lang/generate_go.go

@@ -613,7 +613,7 @@ func (g *goGenerator) stmt(stmt Stmt, tail []Stmt) error {
 		} else {
 			value, err := g.expr(s.Value, g.currentFunc.ReturnType)
 			if err != nil {
-				return err
+				return fmt.Errorf("returning %s: %w", g.currentFunc.ReturnType, err)
 			}
 			value = g.passthroughValue(s.Value, value)
 			g.line(fmt.Sprintf("return %s", value))
@@ -1210,12 +1210,15 @@ func (g *goGenerator) expr(expr Expr, expectedType string) (string, error) {
 		}
 
 		args := make([]string, 0, len(e.Args))
-		argTypes := g.callArgTypes(e.Callee, len(e.Args))
+		argTypes := g.callArgTypes(e)
 		for i, arg := range e.Args {
 			argType := ""
 			if i < len(argTypes) {
 				argType = argTypes[i]
 			}
+			if _, isNull := arg.(NullExpr); isNull && g.goCallArgumentIsNilable(e, i) {
+				argType = ""
+			}
 			value, err := g.expr(arg, argType)
 			if err != nil {
 				return "", err
@@ -2162,7 +2165,23 @@ func (g *goGenerator) lambda(lambda LambdaExpr, expectedType string) (string, er
 	for _, param := range params {
 		sub.defineType(param.Name, param.Type)
 	}
-	if err := sub.block(lambda.Body); err != nil {
+	if returnType != "" && returnType != "Unit" && !lambdaHasValueReturn(lambda.Body) {
+		if len(lambda.Body) == 0 {
+			return "", fmt.Errorf("value lambda requires a final expression")
+		}
+		last, ok := lambda.Body[len(lambda.Body)-1].(ExprStmt)
+		if !ok {
+			return "", fmt.Errorf("value lambda requires a final expression or explicit return")
+		}
+		if err := sub.block(lambda.Body[:len(lambda.Body)-1]); err != nil {
+			return "", err
+		}
+		value, err := sub.expr(last.Value, returnType)
+		if err != nil {
+			return "", err
+		}
+		sub.line("return " + value)
+	} else if err := sub.block(lambda.Body); err != nil {
 		return "", err
 	}
 	b.Write(sub.buf.Bytes())
@@ -2170,7 +2189,50 @@ func (g *goGenerator) lambda(lambda LambdaExpr, expectedType string) (string, er
 	return b.String(), nil
 }
 
-func (g *goGenerator) callArgTypes(callee Expr, argCount int) []string {
+func (g *goGenerator) callArgTypes(call CallExpr) []string {
+	argCount := len(call.Args)
+	if resolved := exprMeta(call); resolved != nil {
+		switch node := resolved.Node.(type) {
+		case HIRGoCall:
+			if len(node.Params) > 0 {
+				params := node.Params
+				if node.InjectContext {
+					params = params[1:]
+				}
+				return semanticParamStrings(params, argCount, node.Variadic)
+			}
+		case HIRGotlinCall:
+			if node.Target != nil {
+				if function, ok := node.Target.Type.(FunctionType); ok {
+					bindings := map[string]Type{}
+					for index, argument := range call.TypeArgs {
+						if index < len(function.TypeParams) {
+							typ, _ := g.semantic.ResolveType(argument)
+							bindings[function.TypeParams[index]] = typ
+						}
+					}
+					if selector, ok := call.Callee.(SelectorExpr); ok {
+						if _, receiverBindings := classInstance(ResolvedType(selector.Receiver)); receiverBindings != nil {
+							for name, typ := range receiverBindings {
+								bindings[name] = typ
+							}
+						}
+					}
+					for index, argument := range call.Args {
+						if index < len(function.Params) {
+							inferTypeBindings(function.Params[index], ResolvedType(argument), bindings)
+						}
+					}
+					params := make([]Type, len(function.Params))
+					for index, param := range function.Params {
+						params[index] = substituteType(param, bindings)
+					}
+					return semanticParamStrings(params, argCount, function.Variadic)
+				}
+			}
+		}
+	}
+	callee := call.Callee
 	ident, ok := callee.(IdentExpr)
 	if ok {
 		if ident.Name == "runCatching" {
@@ -2192,7 +2254,6 @@ func (g *goGenerator) callArgTypes(callee Expr, argCount int) []string {
 			}
 			return argTypes
 		}
-		return nil
 	}
 
 	selector, ok := selectorPath(callee)
@@ -2209,6 +2270,63 @@ func (g *goGenerator) callArgTypes(callee Expr, argCount int) []string {
 	}
 }
 
+func semanticParamStrings(params []Type, argCount int, variadic bool) []string {
+	result := make([]string, 0, argCount)
+	for index := 0; index < argCount; index++ {
+		paramIndex := index
+		if paramIndex >= len(params) {
+			if !variadic || len(params) == 0 {
+				break
+			}
+			paramIndex = len(params) - 1
+		}
+		param := params[paramIndex]
+		if variadic && paramIndex == len(params)-1 {
+			if list, ok := param.(GenericType); ok && (list.Base.String() == "List" || list.Base.String() == "MutableList") && len(list.Args) == 1 {
+				param = list.Args[0]
+			}
+		}
+		result = append(result, param.String())
+	}
+	return result
+}
+
+func (g *goGenerator) goCallArgumentIsNilable(call CallExpr, index int) bool {
+	resolved := exprMeta(call)
+	if resolved == nil {
+		return false
+	}
+	goCall, ok := resolved.Node.(HIRGoCall)
+	if !ok {
+		return false
+	}
+	params := goCall.Params
+	if goCall.InjectContext && len(params) > 0 {
+		params = params[1:]
+	}
+	if len(params) == 0 {
+		return false
+	}
+	paramIndex := index
+	if paramIndex >= len(params) {
+		if !goCall.Variadic {
+			return false
+		}
+		paramIndex = len(params) - 1
+	}
+	param := params[paramIndex]
+	if goCall.Variadic && paramIndex == len(params)-1 {
+		if list, ok := param.(GenericType); ok && len(list.Args) == 1 {
+			param = list.Args[0]
+		}
+	}
+	switch param.(type) {
+	case GoInterfaceType, GoPointerType, NullableType:
+		return true
+	}
+	return false
+}
+
 func nameUsedInStmts(name string, stmts []Stmt) bool {
 	return nameUsedInStmtsWithShadow(name, stmts, false)
 }

internal/lang/go_types.go

@@ -176,6 +176,7 @@ func semanticTypeFromGoResults(results *gotypes.Tuple) Type {
 }
 
 func semanticTypeFromGo(typ gotypes.Type) Type {
+	typ = gotypes.Unalias(typ)
 	switch value := typ.(type) {
 	case *gotypes.Basic:
 		switch value.Kind() {
@@ -210,6 +211,9 @@ func semanticTypeFromGo(typ gotypes.Type) Type {
 			return NamedType{Name: "Error"}
 		}
 		if value.Obj().Pkg() != nil {
+			if _, ok := value.Underlying().(*gotypes.Interface); ok {
+				return GoInterfaceType{Name: value.Obj().Pkg().Name() + "." + value.Obj().Name()}
+			}
 			return NamedType{Name: value.Obj().Pkg().Name() + "." + value.Obj().Name()}
 		}
 		return NamedType{Name: value.Obj().Name()}

internal/lang/hir.go

@@ -458,6 +458,7 @@ func (resolver *semanticResolver) resolveStmts(stmts []Stmt, scope *Scope, class
 
 func (resolver *semanticResolver) resolveExpr(expr Expr, scope *Scope, class *ClassSymbol, expected Type) (Expr, Type) {
 	environment := TypeEnvironment{Scope: scope, Class: class}
+	var typeOverride Type
 	switch value := expr.(type) {
 	case UnaryExpr:
 		value.Value, _ = resolver.resolveExpr(value.Value, scope, class, UnknownType{})
@@ -468,8 +469,26 @@ func (resolver *semanticResolver) resolveExpr(expr Expr, scope *Scope, class *Cl
 		expr = value
 	case CallExpr:
 		value.Callee, _ = resolver.resolveExpr(value.Callee, scope, class, UnknownType{})
+		signature, hasSignature := resolver.callSignature(value, environment)
+		params := signature.Params
+		if hasSignature && len(params) > 0 && params[0].String() == "context.Context" {
+			explicitContext := len(value.Args) > 0 && resolver.program.TypeOf(value.Args[0], environment).String() == "context.Context"
+			if !explicitContext && contextCanBeOmitted(len(value.Args), len(params), signature.Variadic) {
+				params = params[1:]
+			}
+		}
+		bindings := map[string]Type{}
+		for index, argument := range value.TypeArgs {
+			if index < len(signature.TypeParams) {
+				typ, _ := resolver.program.ResolveType(argument)
+				bindings[signature.TypeParams[index]] = typ
+			}
+		}
 		for index := range value.Args {
 			argumentExpected := Type(UnknownType{})
+			if hasSignature && index < len(params) {
+				argumentExpected = substituteType(params[index], bindings)
+			}
 			if selector, ok := value.Callee.(SelectorExpr); ok {
 				if receiver, ok := selector.Receiver.(IdentExpr); ok && receiver.Name == "Result" {
 					if selector.Name == "Err" {
@@ -482,6 +501,9 @@ func (resolver *semanticResolver) resolveExpr(expr Expr, scope *Scope, class *Cl
 				}
 			}
 			value.Args[index], _ = resolver.resolveExpr(value.Args[index], scope, class, argumentExpected)
+			if hasSignature && index < len(params) {
+				inferTypeBindings(params[index], ResolvedType(value.Args[index]), bindings)
+			}
 		}
 		for index := range value.NamedArgs {
 			value.NamedArgs[index].Value, _ = resolver.resolveExpr(value.NamedArgs[index].Value, scope, class, UnknownType{})
@@ -520,17 +542,40 @@ func (resolver *semanticResolver) resolveExpr(expr Expr, scope *Scope, class *Cl
 		expr = value
 	case LambdaExpr:
 		lambdaScope := NewScope(scope)
+		functionExpected, hasFunctionExpected := expected.(FunctionType)
 		if value.ImplicitIt {
-			_ = lambdaScope.Define(&Symbol{Name: "it", Kind: VariableSymbol, Type: UnknownType{}})
+			itType := Type(UnknownType{})
+			if hasFunctionExpected && len(functionExpected.Params) == 1 {
+				itType = functionExpected.Params[0]
+			}
+			_ = lambdaScope.Define(&Symbol{Name: "it", Kind: VariableSymbol, Type: itType})
 		}
-		for _, param := range value.Params {
+		for index, param := range value.Params {
 			typ, _ := resolver.program.ResolveTypeRef(param.TypeRef)
+			if isUnknownType(typ) && hasFunctionExpected && index < len(functionExpected.Params) {
+				typ = functionExpected.Params[index]
+			}
 			_ = lambdaScope.Define(&Symbol{Name: param.Name, Kind: VariableSymbol, Type: typ})
 		}
 		resolver.resolveStmts(value.Body, lambdaScope, class, functionResult(expected))
+		if hasFunctionExpected {
+			result := functionExpected.Result
+			if !lambdaHasValueReturn(value.Body) && len(value.Body) > 0 {
+				if final, ok := value.Body[len(value.Body)-1].(ExprStmt); ok {
+					actual := ResolvedType(final.Value)
+					if !isUnknownType(actual) {
+						result = actual
+					}
+				}
+			}
+			typeOverride = FunctionType{Params: functionExpected.Params, Result: result, Effects: functionExpected.Effects}
+		}
 		expr = value
 	}
 	typ := resolver.program.TypeOf(expr, environment)
+	if typeOverride != nil {
+		typ = typeOverride
+	}
 	if isUnknownType(typ) && !isUnknownType(expected) {
 		typ = expected
 	}
@@ -562,6 +607,57 @@ func (resolver *semanticResolver) resolveExpr(expr Expr, scope *Scope, class *Cl
 	return withExprMeta(expr, semantic), typ
 }
 
+func (resolver *semanticResolver) callSignature(call CallExpr, environment TypeEnvironment) (FunctionType, bool) {
+	if signature, ok := resolver.program.goCallSignature(call); ok {
+		return signature, true
+	}
+	switch callee := call.Callee.(type) {
+	case IdentExpr:
+		if symbol, ok := resolver.program.Global.Lookup(callee.Name); ok {
+			if signature, ok := symbol.Type.(FunctionType); ok {
+				return signature, true
+			}
+		}
+		if class := resolver.program.ClassInfo[callee.Name]; class != nil {
+			params := make([]Type, len(class.Decl.Fields))
+			for index, field := range class.Decl.Fields {
+				params[index] = class.Fields[field.Name].Type
+			}
+			return FunctionType{TypeParams: class.TypeParams, Params: params, Result: ClassType{Class: class}}, true
+		}
+	case SelectorExpr:
+		if receiver, ok := callee.Receiver.(IdentExpr); ok {
+			if pack := resolver.program.Packages[receiver.Name]; pack != nil {
+				if function := pack.Function(callee.Name); function != nil {
+					if signature, ok := function.Type.(FunctionType); ok {
+						return signature, true
+					}
+				}
+			}
+		}
+		if class, bindings := classInstance(resolver.program.TypeOf(callee.Receiver, environment)); class != nil {
+			if method := class.Method(callee.Name); method != nil {
+				if signature, ok := method.Type.(FunctionType); ok {
+					substituted := substituteType(signature, bindings).(FunctionType)
+					return substituted, true
+				}
+			}
+		}
+	}
+	return FunctionType{}, false
+}
+
+func contextCanBeOmitted(argumentCount, parameterCount int, variadic bool) bool {
+	if variadic {
+		minimum := parameterCount - 2
+		if minimum < 0 {
+			minimum = 0
+		}
+		return argumentCount >= minimum
+	}
+	return argumentCount == parameterCount-1
+}
+
 func (resolver *semanticResolver) hirNode(expr Expr, semantic *HIRExpr) HIRNode {
 	switch semantic.Meaning {
 	case LiteralExpr:

internal/lang/types.go

@@ -51,6 +51,11 @@ type ImportedClassType struct {
 func (ImportedClassType) typeNode()        {}
 func (t ImportedClassType) String() string { return t.Package + "." + t.Class.Name }
 
+type GoInterfaceType struct{ Name string }
+
+func (GoInterfaceType) typeNode()        {}
+func (t GoInterfaceType) String() string { return t.Name }
+
 type NullableType struct{ Element Type }
 
 func (NullableType) typeNode()        {}
@@ -330,6 +335,17 @@ func inferTypeBindings(parameter, argument Type, bindings map[string]Type) {
 				inferTypeBindings(expected.Args[index], actual.Args[index], bindings)
 			}
 		}
+	case FunctionType:
+		actual, ok := argument.(FunctionType)
+		if !ok {
+			return
+		}
+		for index := range expected.Params {
+			if index < len(actual.Params) {
+				inferTypeBindings(expected.Params[index], actual.Params[index], bindings)
+			}
+		}
+		inferTypeBindings(expected.Result, actual.Result, bindings)
 	}
 }
 
@@ -341,6 +357,8 @@ func renderGoType(typ Type) string {
 		return "*" + value.Class.Name
 	case ImportedClassType:
 		return "*" + value.Package + "." + value.Class.Name
+	case GoInterfaceType:
+		return value.Name
 	case NullableType:
 		element := renderGoType(value.Element)
 		switch value.Element.(type) {