Infer final expression returns in lambdas

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: