Extraction: slightly better heuristic for Obj.magic simplifications

 On an application (f args) where the head is magic, we first remove Obj.magic
 on arguments before continuing with simplifications (that may push magic down
 inside the arguments).

 For instance, starting with ((Obj.magic f) (Obj.magic (g h))), we now end
 with ((Obj.magic f) (g h)) instead of ((Obj.magic f) ((Obj.magic g) h))) as
 before.

1 files changed, 13 insertions, 11 deletions

diff --git a/plugins/extraction/mlutil.ml b/plugins/extraction/mlutil.ml
index d06583010f..62724f2119 100644
--- a/plugins/extraction/mlutil.ml
+++ b/plugins/extraction/mlutil.ml
@@ -1013,9 +1013,20 @@ let expand_linear_let o id e =
 
 (* Some beta-iota reductions + simplifications. *)
 
+let rec unmagic = function MLmagic e -> unmagic e | e -> e
+let is_magic = function MLmagic _ -> true | _ -> false
+let magic_hd a = match a with
+  | MLmagic _ :: _ -> a
+  | e :: a -> MLmagic e :: a
+  | [] -> assert false
+
 let rec simpl o = function
   | MLapp (f, []) -> simpl o f
-  | MLapp (f, a) -> simpl_app o (List.map (simpl o) a) (simpl o f)
+  | MLapp (MLapp(f,a),a') -> simpl o (MLapp(f,a@a'))
+  | MLapp (f, a) ->
+     (* When the head of the application is magic, no need for magic on args *)
+     let a = if is_magic f then List.map unmagic a else a in
+     simpl_app o (List.map (simpl o) a) (simpl o f)
   | MLcase (typ,e,br) ->
       let br = Array.map (fun (l,p,t) -> (l,p,simpl o t)) br in
       simpl_case o typ br (simpl o e)
@@ -1047,7 +1058,6 @@ let rec simpl o = function
 (* invariant : list [a] of arguments is non-empty *)
 
 and simpl_app o a = function
-  | MLapp (f',a') -> simpl_app o (a'@a) f'
   | MLlam (Dummy,t) ->
       simpl o (MLapp (ast_pop t, List.tl a))
   | MLlam (id,t) -> (* Beta redex *)
@@ -1062,15 +1072,7 @@ and simpl_app o a = function
       (* When we've at least one argument, we permute the magic
          and the lambda, to simplify things a bit (see #2795).
          Alas, the 1st argument must also be magic then. *)
-      let a' = match a with
-        | MLmagic _ :: _ -> a
-        | e::a' -> MLmagic e::a'
-        | [] -> assert false
-      in
-      simpl_app o a' (MLlam (id,MLmagic t))
-  | MLmagic _ as e ->
-      (* When the head is magic, no need for magic on args *)
-      MLapp (e, List.map (function MLmagic e -> e | e -> e) a)
+      simpl_app o (magic_hd a) (MLlam (id,MLmagic t))
   | MLletin (id,e1,e2) when o.opt_let_app ->
       (* Application of a letin: we push arguments inside *)
       MLletin (id, e1, simpl o (MLapp (e2, List.map (ast_lift 1) a)))