Coerce the result of a function call to the function's return type

Signed-off-by: Nick Cameron <[email protected]>

Author: nrc

rust/kcl-lib/src/execution/exec_ast.rs

@@ -2164,6 +2164,34 @@ fn assign_args_to_params_kw(
     Ok(())
 }
 
+fn coerce_result_type(
+    result: Result<Option<KclValue>, KclError>,
+    function_expression: NodeRef<'_, FunctionExpression>,
+    exec_state: &mut ExecState,
+) -> Result<Option<KclValue>, KclError> {
+    if let Ok(Some(val)) = result {
+        if let Some(ret_ty) = &function_expression.return_type {
+            let ty = RuntimeType::from_parsed(ret_ty.inner.clone(), exec_state, ret_ty.as_source_range())
+                .map_err(|e| KclError::Semantic(e.into()))?;
+            let val = val.coerce(&ty, exec_state).map_err(|_| {
+                KclError::Semantic(KclErrorDetails {
+                    message: format!(
+                        "This function requires its result to be of type `{}`, but found {}",
+                        ty.human_friendly_type(),
+                        val.human_friendly_type(),
+                    ),
+                    source_ranges: ret_ty.as_source_ranges(),
+                })
+            })?;
+            Ok(Some(val))
+        } else {
+            Ok(Some(val))
+        }
+    } else {
+        result
+    }
+}
+
 async fn call_user_defined_function(
     args: Vec<Arg>,
     memory: EnvironmentRef,
@@ -2184,13 +2212,16 @@ async fn call_user_defined_function(
     let result = ctx
         .exec_block(&function_expression.body, exec_state, BodyType::Block)
         .await;
-    let result = result.map(|_| {
+    let mut result = result.map(|_| {
         exec_state
             .stack()
             .get(memory::RETURN_NAME, function_expression.as_source_range())
             .ok()
             .cloned()
     });
+
+    result = coerce_result_type(result, function_expression, exec_state);
+
     // Restore the previous memory.
     exec_state.mut_stack().pop_env();
 
@@ -2218,13 +2249,16 @@ async fn call_user_defined_function_kw(
     let result = ctx
         .exec_block(&function_expression.body, exec_state, BodyType::Block)
         .await;
-    let result = result.map(|_| {
+    let mut result = result.map(|_| {
         exec_state
             .stack()
             .get(memory::RETURN_NAME, function_expression.as_source_range())
             .ok()
             .cloned()
     });
+
+    result = coerce_result_type(result, function_expression, exec_state);
+
     // Restore the previous memory.
     exec_state.mut_stack().pop_env();