8 files changed, 58 insertions, 16 deletions

diff --git a/plugins/extraction/common.ml b/plugins/extraction/common.ml
index 9772ebd641..9aec190d0a 100644
--- a/plugins/extraction/common.ml
+++ b/plugins/extraction/common.ml
@@ -405,7 +405,7 @@ let ref_renaming_fun (k,r) =
     let idg = safe_basename_of_global r in
     match l with
     | [""] -> (* this happens only at toplevel of the monolithic case *)
-      let globs = Id.Set.elements (get_global_ids ()) in
+      let globs = get_global_ids () in
       let id = next_ident_away (kindcase_id k idg) globs in
       Id.to_string id
     | _ -> modular_rename k idg
diff --git a/plugins/extraction/extract_env.ml b/plugins/extraction/extract_env.ml
index f503c572d0..3c46d5c43b 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 5769ff1176..7bbb825b10 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/extraction.ml b/plugins/extraction/extraction.ml
index 7644b49ceb..a227478d0f 100644
--- a/plugins/extraction/extraction.ml
+++ b/plugins/extraction/extraction.ml
@@ -141,6 +141,7 @@ let make_typvar n vl =
     if not (String.contains s '\'') && Unicode.is_basic_ascii s then id
     else id_of_name Anonymous
   in
+  let vl = Id.Set.of_list vl in
   next_ident_away id' vl
 
 let rec type_sign_vl env c =
diff --git a/plugins/extraction/miniml.mli b/plugins/extraction/miniml.mli
index edebba49df..5e967ef379 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 a4c2bcd883..b01b0198d5 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 ||
@@ -769,6 +773,20 @@ let eta_red e =
 	else e
     | _ -> e
 
+(* Performs an eta-reduction when the core is atomic,
+   or otherwise returns None *)
+
+let atomic_eta_red e =
+  let ids,t = collect_lams e in
+  let n = List.length ids in
+  match t with
+  | MLapp (f,a) when test_eta_args_lift 0 n a ->
+     (match f with
+      | MLrel k when k>n -> Some (MLrel (k-n))
+      | MLglob _ | MLexn _ | MLdummy _ -> Some f
+      | _ -> None)
+  | _ -> None
+
 (*s Computes all head linear beta-reductions possible in [(t a)].
   Non-linear head beta-redex become let-in. *)
 
@@ -1053,6 +1071,10 @@ let rec simpl o = function
      simpl o (MLcase(typ,e,br'))
   | MLmagic(MLdummy _ as e) when lang () == Haskell -> e
   | MLmagic(MLexn _ as e) -> e
+  | MLlam _ as e ->
+     (match atomic_eta_red e with
+      | Some e' -> e'
+      | None -> ast_map (simpl o) e)
   | a -> ast_map (simpl o) a
 
 (* invariant : list [a] of arguments is non-empty *)
diff --git a/plugins/extraction/table.ml b/plugins/extraction/table.ml
index ca98f07e8d..995d5fd19d 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 ()].
@@ -750,11 +755,11 @@ let extraction_implicit r l =
 
 let blacklist_table = Summary.ref Id.Set.empty ~name:"ExtrBlacklist"
 
-let modfile_ids = ref []
+let modfile_ids = ref Id.Set.empty
 let modfile_mps = ref MPmap.empty
 
 let reset_modfile () =
-  modfile_ids := Id.Set.elements !blacklist_table;
+  modfile_ids := !blacklist_table;
   modfile_mps := MPmap.empty
 
 let string_of_modfile mp =
@@ -763,7 +768,7 @@ let string_of_modfile mp =
     let id = Id.of_string (raw_string_of_modfile mp) in
     let id' = next_ident_away id !modfile_ids in
     let s' = Id.to_string id' in
-    modfile_ids := id' :: !modfile_ids;
+    modfile_ids := Id.Set.add id' !modfile_ids;
     modfile_mps := MPmap.add mp s' !modfile_mps;
     s'
 
diff --git a/plugins/extraction/table.mli b/plugins/extraction/table.mli
index 7e47d0bc81..cc93f294b3 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