From 72606e78188ccb57fc6d2466051261ebc113960b Mon Sep 17 00:00:00 2001 From: pavel Date: Thu, 27 Aug 2026 23:31:46 +0200 Subject: [PATCH] Infer final expression returns in lambdas --- README.md | 1 + internal/lang/expression_function_test.go | 56 ++++++++++ internal/lang/generate_go.go | 128 +++++++++++++++++++++- internal/lang/go_types.go | 4 + internal/lang/hir.go | 100 ++++++++++++++++- internal/lang/types.go | 18 +++ 6 files changed, 300 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 8e86c32..1af4c00 100644 --- a/README.md +++ b/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` diff --git a/internal/lang/expression_function_test.go b/internal/lang/expression_function_test.go index 0ed4fa3..54d912c 100644 --- a/internal/lang/expression_function_test.go +++ b/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(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) { + 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) + } +} diff --git a/internal/lang/generate_go.go b/internal/lang/generate_go.go index 066ac22..bc87b1c 100644 --- a/internal/lang/generate_go.go +++ b/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) } diff --git a/internal/lang/go_types.go b/internal/lang/go_types.go index 1e66431..51284fe 100644 --- a/internal/lang/go_types.go +++ b/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()} diff --git a/internal/lang/hir.go b/internal/lang/hir.go index b52b8ea..fe9e64b 100644 --- a/internal/lang/hir.go +++ b/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: diff --git a/internal/lang/types.go b/internal/lang/types.go index 6d8642f..847d373 100644 --- a/internal/lang/types.go +++ b/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) {