Fix layouts in flambda1 partial application simplification (#1549)

* Fix layouts in flambda1 partial application simplification

* add back assert

Author: ccasin

middle_end/flambda/inline_and_simplify.ml

@@ -777,10 +777,13 @@ and simplify_apply env r ~(apply : Flambda.apply) : Flambda.t * R.t =
                 ~function_decls ~lhs_of_application ~closure_id_being_applied
                 ~function_decl ~value_set_of_closures ~dbg ~reg_close ~mode
                 ~inlined_requested ~specialise_requested ~result_layout
-            else if nargs > 0 && nargs < arity then
+            else if nargs > 0 && nargs < arity then begin
+              assert(Lambda.compatible_layout Lambda.layout_function
+                       result_layout);
               simplify_partial_application env r ~lhs_of_application
                 ~closure_id_being_applied ~function_decl ~args ~mode ~dbg
-                ~inlined_requested ~specialise_requested ~result_layout
+                ~inlined_requested ~specialise_requested
+            end
             else
               Misc.fatal_errorf "Function with arity %d when simplifying \
                   application expression: %a"
@@ -809,7 +812,7 @@ and simplify_full_application env r ~function_decls ~lhs_of_application
 
 and simplify_partial_application env r ~lhs_of_application
       ~closure_id_being_applied ~function_decl ~args ~mode ~dbg
-      ~inlined_requested ~specialise_requested ~result_layout
+      ~inlined_requested ~specialise_requested
   =
   let arity = A.function_arity function_decl in
   assert (arity > List.length args);
@@ -867,10 +870,9 @@ and simplify_partial_application env r ~lhs_of_application
         inlined = Default_inlined;
         specialise = Default_specialise;
         probe = None;
-        result_layout;
+        result_layout = function_decl.A.return_layout;
       }
     in
-    assert(Lambda.compatible_layout function_decl.A.return_layout result_layout);
     let closure_variable =
       Variable.rename ~debug_info:(Closure_id.debug_info closure_id_being_applied)
         (Closure_id.unwrap closure_id_being_applied)