6 files changed, 35 insertions, 12 deletions

diff --git a/plugins/extraction/extract_env.ml b/plugins/extraction/extract_env.ml
index f503c572d..3c46d5c43 100644
--- a/plugins/extraction/extract_env.ml
+++ b/plugins/extraction/extract_env.ml
@@ -342,7 +342,7 @@ let rec extract_structure env mp reso ~all = function
 
 and extract_mexpr env mp = function
   | MEwith _ -> assert false (* no 'with' syntax for modules *)
-  | me when lang () != Ocaml ->
+  | me when lang () != Ocaml || Table.is_extrcompute () ->
       (* In Haskell/Scheme, we expand everything.
          For now, we also extract everything, dead code will be removed later
          (see [Modutil.optimize_struct]. *)
@@ -570,11 +570,12 @@ let print_structure_to_file (fn,si,mo) dry struc =
 let reset () =
   Visit.reset (); reset_tables (); reset_renaming_tables Everything
 
-let init modular library =
+let init ?(compute=false) modular library =
   check_inside_section (); check_inside_module ();
   set_keywords (descr ()).keywords;
   set_modular modular;
   set_library library;
+  set_extrcompute compute;
   reset ();
   if modular && lang () == Scheme then error_scheme ()
 
@@ -684,8 +685,22 @@ let extraction_library is_rec m =
   List.iter print struc;
   reset ()
 
+(** For extraction compute, we flatten all the module structure,
+    getting rid of module types or unapplied functors *)
+
+let flatten_structure struc =
+  let rec flatten_elem (lab,elem) = match elem with
+    |SEdecl d -> [d]
+    |SEmodtype _ -> []
+    |SEmodule m -> match m.ml_mod_expr with
+      |MEfunctor _ -> []
+      |MEident _ | MEapply _ -> assert false (* should be expanded *)
+      |MEstruct (_,elems) -> flatten_elems elems
+  and flatten_elems l = List.flatten (List.map flatten_elem l)
+  in flatten_elems (List.flatten (List.map snd struc))
+
 let structure_for_compute c =
-  init false false;
+  init false false ~compute:true;
   let env = Global.env () in
   let ast, mlt = Extraction.extract_constr env c in
   let ast = Mlutil.normalize ast in
@@ -694,8 +709,7 @@ let structure_for_compute c =
   let () = ast_iter_references add_ref add_ref add_ref ast in
   let refs = Refset.elements !refs in
   let struc = optimize_struct (refs,[]) (mono_environment refs []) in
-  let flatstruc = List.map snd (List.flatten (List.map snd struc)) in
-  flatstruc, ast, mlt
+  (flatten_structure struc), ast, mlt
 
 (* For the test-suite :
    extraction to a temporary file + run ocamlc on it *)
diff --git a/plugins/extraction/extract_env.mli b/plugins/extraction/extract_env.mli
index 5769ff117..7bbb825b1 100644
--- a/plugins/extraction/extract_env.mli
+++ b/plugins/extraction/extract_env.mli
@@ -34,5 +34,4 @@ val print_one_decl :
 (* Used by Extraction Compute *)
 
 val structure_for_compute :
-  Term.constr ->
-    Miniml.ml_flat_structure * Miniml.ml_ast * Miniml.ml_type
+  Term.constr -> (Miniml.ml_decl list) * Miniml.ml_ast * Miniml.ml_type
diff --git a/plugins/extraction/miniml.mli b/plugins/extraction/miniml.mli
index edebba49d..5e967ef37 100644
--- a/plugins/extraction/miniml.mli
+++ b/plugins/extraction/miniml.mli
@@ -187,8 +187,6 @@ type ml_structure = (ModPath.t * ml_module_structure) list
 
 type ml_signature = (ModPath.t * ml_module_sig) list
 
-type ml_flat_structure = ml_structure_elem list
-
 type unsafe_needs = {
   mldummy : bool;
   tdummy : bool;
diff --git a/plugins/extraction/mlutil.ml b/plugins/extraction/mlutil.ml
index b16b430a8..b01b0198d 100644
--- a/plugins/extraction/mlutil.ml
+++ b/plugins/extraction/mlutil.ml
@@ -127,11 +127,15 @@ let rec mgu = function
   | Taxiom, Taxiom -> ()
   | _ -> raise Impossible
 
-let needs_magic p = try mgu p; false with Impossible -> true
+let skip_typing () = lang () == Scheme || is_extrcompute ()
 
-let put_magic_if b a = if b && lang () != Scheme then MLmagic a else a
+let needs_magic p =
+  if skip_typing () then false
+  else try mgu p; false with Impossible -> true
 
-let put_magic p a = if needs_magic p && lang () != Scheme then MLmagic a else a
+let put_magic_if b a = if b then MLmagic a else a
+
+let put_magic p a = if needs_magic p then MLmagic a else a
 
 let generalizable a =
   lang () != Ocaml ||
diff --git a/plugins/extraction/table.ml b/plugins/extraction/table.ml
index 30e3b520f..995d5fd19 100644
--- a/plugins/extraction/table.ml
+++ b/plugins/extraction/table.ml
@@ -250,6 +250,11 @@ let modular () = !modular_ref
 let set_library b = library_ref := b
 let library () = !library_ref
 
+let extrcompute = ref false
+
+let set_extrcompute b = extrcompute := b
+let is_extrcompute () = !extrcompute
+
 (*s Printing. *)
 
 (* The following functions work even on objects not in [Global.env ()].
diff --git a/plugins/extraction/table.mli b/plugins/extraction/table.mli
index 7e47d0bc8..cc93f294b 100644
--- a/plugins/extraction/table.mli
+++ b/plugins/extraction/table.mli
@@ -165,6 +165,9 @@ val modular : unit -> bool
 val set_library : bool -> unit
 val library : unit -> bool
 
+val set_extrcompute : bool -> unit
+val is_extrcompute : unit -> bool
+
 (*s Table for custom inlining *)
 
 val to_inline : global_reference -> bool