Extraction: some minor fixes regarding recursive extraction

src/mlw/compile.ml

@@ -206,7 +206,6 @@ module ML = struct
 
 end
 
-
 (** Translation from Mlw to ML *)
 
 module Translate = struct
@@ -346,8 +345,9 @@ module Translate = struct
         List.map def cty_app.cty_args in
       args @ extra_args in
     let eapp = ML.mk_expr (Mltree.Eapp (rsc, args)) (Mltree.C cty_app)
-                                                    cty_app.cty_effect in
-    ML.mk_expr (Mltree.Efun (args_f, eapp)) (Mltree.C cty_app) cty_app.cty_effect
+        cty_app.cty_effect in
+    ML.mk_expr (Mltree.Efun (args_f, eapp)) (Mltree.C cty_app)
+      cty_app.cty_effect
 
   (* function arguments *)
   let filter_params args =

src/tools/why3extract.ml

@@ -87,8 +87,8 @@ let () =
     Whyconf.Args.exit_with_usage option_list usage_msg
   end
 
-let opt_rec_single = !opt_rec_single
 let opt_modu_flat  = !opt_modu_flat
+let opt_rec_single = !opt_rec_single
 
 let opt_output = match opt_modu_flat, !opt_output with
   | Modular, None ->
@@ -273,19 +273,19 @@ let find_decl mm id =
   let m = translate_module m in
   Ident.Mid.find id m.Mltree.mod_known
 
-let rec visit mm id =
+let rec visit_mem mm id =
   if not (Ident.Hid.mem visited id) then begin
     try
       let d = find_decl mm id in
       Ident.Hid.add visited id ();
-      ML.iter_deps (visit mm) d;
+      ML.iter_deps (visit_mem mm) d;
       toextract := id :: !toextract
     with Not_found -> ()
   end
 
-(* let visit mm id = *)
-(*   if opt_rec_single = Recursive then visit mm id *)
-(*   else toextract := id :: !toextract *)
+let visit mm id =
+  if opt_rec_single = Recursive then visit_mem mm id
+  else toextract := id :: !toextract
 
 let flat_extraction mm = function
   | File fname ->
@@ -298,17 +298,17 @@ let flat_extraction mm = function
       end;
       Mstr.add s m mm in
     Mstr.fold do_m mmf mm
-  | Module (path, m) ->
-    let m = find_module_path mm path m in (* FIXME: catch Not_found here *)
-    let m = translate_module m in
-    Ident.Mid.iter (fun id _ -> visit mm id) m.Mltree.mod_known;
-    mm
-  | Symbol (path, m, s) ->
-    let m = find_module_path mm path m in
+  | Module (path, ms) ->
+    let m = find_module_path mm path ms in (* FIXME: catch Not_found here *)
+    let m_t = translate_module m in
+    Ident.Mid.iter (fun id _ -> visit mm id) m_t.Mltree.mod_known;
+    Mstr.add ms m mm
+  | Symbol (path, ms, s) ->
+    let m = find_module_path mm path ms in
     let ns = m.mod_export in
     let id = find_symbol_id ns s in
     visit mm id;
-    mm
+    Mstr.add ms m mm
 
 let () =
   try
@@ -318,7 +318,7 @@ let () =
     | Flat ->
       let mm = Queue.fold flat_extraction Mstr.empty opt_queue in
       let visit_m _ m = let tm = translate_module m in
-        Ident.Mid.iter (fun id _ -> visit mm id) tm.Mltree.mod_known in
+        Ident.Mid.iter (fun id _ -> visit_mem mm id) tm.Mltree.mod_known in
       Mstr.iter visit_m mm;
       let (_fg, pargs, pr) = Pdriver.lookup_printer opt_driver in
       let cout = match opt_output with
@@ -334,7 +334,14 @@ let () =
         let m = translate_module pm in
         let d = Ident.Mid.find id m.Mltree.mod_known in
         pr pargs ~flat:true pm fmt d in
-      List.iter (extract fmt) (List.rev !toextract);
+      let idl = List.rev !toextract in
+      let is_local id =
+        let (path, m, _) = Pmodule.restore_path id in
+        path = [] || Mstr.mem m mm in
+      let idl = match opt_rec_single with
+        | Single -> List.filter is_local idl
+        | Recursive -> idl in
+      List.iter (extract fmt) idl;
       if cout <> stdout then close_out cout
   with e when not (Debug.test_flag Debug.stack_trace) ->
     eprintf "%a@." Exn_printer.exn_printer e;