Update translate.ml

src/compiler/gencode/translate.ml

@@ -456,9 +456,9 @@ let rec exp env loop_path code { Zelus.e_desc = desc } =
   | Zelus.Efun { Zelus.f_kind = k; Zelus.f_args = arg_list;
 		 Zelus.f_body = r; Zelus.f_env = f_env } ->
      let ty = Typinfo.get_type r.r_info in
-     let pat_list = List.map (fun a_list -> List.map pattern a_list) arg_list in
+     let pat_list = List.map (fun a_list -> List.map arg a_list) arg_list in
      let env, mem_acc, var_acc = append empty_path f_env Env.empty in
-     let code = expression env e in
+     let code = result env r in
      let code = add_mem_vars_to_code code mem_acc var_acc in
      machine k pat_list code ty
   | Esizeapp _ -> Misc.not_yet_implemented "sizeapp"
@@ -473,6 +473,14 @@ and vardec { Zelus.var_name = id; Zelus.var_info = info } =
   let ty = Typinfo.get_type info in
   Evarpat { id; ty = Interface.type_expression_of_typ ty }
 
+and result env { Zelus.r_desc } =
+  match r_desc with
+  | Exp(e) -> expression env e
+  | Returns { b_vars; b_body; b_env } ->
+     let env, mem_acc, var_acc = append [] b_env env in
+     let eq_code = equation env [] b_body empty_code in
+     add_mem_vars_to_code eq_code mem_acc var_acc
+
 (* Patterns *)
 and pattern { Zelus.pat_desc = desc; Zelus.pat_info = info } =
   let ty = Typinfo.get_type info in