refactoring of expression rewriting

Ignore-this: 3713a9bc8818619cb8fa0afdc14d6a41

- factorize the code between lwt and async stuff
- handle type aliases

darcs-hash:20121015094544-c41ad-167ffe2a9f5ecb8973914e9e57b9879c09f561af

src/lib/uTop_main.ml

@@ -214,7 +214,7 @@ let print_out_phrase term string =
   Lwt_main.run (LTerm.fprints term styled)
 
 (* +-----------------------------------------------------------------+
-   | Lwt_main.run auto-insertion                                     |
+   | Toplevel expression rewriting                                   |
    +-----------------------------------------------------------------+ *)
 
 #if ocaml_version >= (4, 0, 0)
@@ -226,31 +226,125 @@ let with_loc loc str = {
 let with_loc loc str = str
 #endif
 
+(* A rule for rewriting a toplevel expression. *)
+type rewrite_rule = {
+  required_values : Longident.t list;
+  (* Values that must exist and be persistent for the rule to apply. *)
+  rewrite : Location.t -> Parsetree.expression -> Parsetree.expression;
+  (* The rewrite function. *)
+  enabled : bool React.signal;
+  (* Whether the rule is enabled or not. *)
+}
+
+(* Rewrite rules, indexed by the identifier of the type
+   constructor. *)
+let rewrite_rules : (Longident.t, rewrite_rule) Hashtbl.t = Hashtbl.create 42
+
 let longident_lwt_main_run = Longident.Ldot (Longident.Lident "Lwt_main", "run")
-let longident_async_main_run = Longident.parse "Async.Std.Thread_safe.block_on_async_exn"
+let longident_async_core_thread_safe_block_on_async_exn =
+  Longident.parse "Async.Std.Thread_safe.block_on_async_exn"
 let longident_unit = Longident.Lident "()"
 
+(* Wrap <expr> into: fun () -> <expr> *)
+let wrap_unit loc e =
+  let i = with_loc loc longident_unit in
+  let p = {
+    Parsetree.ppat_desc = Parsetree.Ppat_construct (i, None, false);
+    Parsetree.ppat_loc = loc;
+  } in
+  {
+    Parsetree.pexp_desc = Parsetree.Pexp_function ("", None, [(p, e)]);
+    Parsetree.pexp_loc = loc;
+  }
+
+let () =
+  (* Rewrite Lwt.t expressions to Lwt_main.run <expr> *)
+  Hashtbl.add rewrite_rules (Longident.Ldot (Longident.Lident "Lwt", "t")) {
+    required_values = [longident_lwt_main_run];
+    rewrite = (fun loc e -> {
+      Parsetree.pexp_desc =
+        Parsetree.Pexp_apply
+          ({ Parsetree.pexp_desc = Parsetree.Pexp_ident (with_loc loc longident_lwt_main_run);
+             Parsetree.pexp_loc = loc },
+           [("", e)]);
+      Parsetree.pexp_loc = loc;
+    });
+    enabled = UTop.auto_run_lwt;
+  };
+
+  (* Rewrite Async.Std.Defered.t expressions to
+     Async_core.Thread_safe.block_on_async_exn (fun () -> <expr>). *)
+  let async_rewrite = {
+    required_values = [longident_async_core_thread_safe_block_on_async_exn];
+    rewrite = (fun loc e -> {
+      Parsetree.pexp_desc =
+        Parsetree.Pexp_apply
+          ({ Parsetree.pexp_desc = Parsetree.Pexp_ident
+              (with_loc loc longident_async_core_thread_safe_block_on_async_exn);
+             Parsetree.pexp_loc = loc },
+           [("", wrap_unit loc e)]);
+      Parsetree.pexp_loc = loc;
+    });
+    enabled = UTop.auto_run_async;
+  } in
+
+  Hashtbl.add rewrite_rules (Longident.parse "Async_core.Deferred.t") async_rewrite;
+  Hashtbl.add rewrite_rules (Longident.parse "Async.Std.Deferred.t") async_rewrite
+
+(* Returns whether the argument is a toplevel expression. *)
 let is_eval = function
   | { Parsetree.pstr_desc = Parsetree.Pstr_eval _ } -> true
   | _ -> false
 
-let rec is_lwt_t typ =
+(* Returns whether the given path is persistent. *)
-  match typ.Types.desc with
+let rec is_persistent_path = function
-    | Types.Tlink typ ->
+  | Path.Pident id -> Ident.persistent id
-        is_lwt_t typ
+  | Path.Pdot (p, _, _) -> is_persistent_path p
-    | Types.Tconstr (Path.Pdot (Path.Pident id, "t", -1), _, _) ->
+  | Path.Papply (_, p) -> is_persistent_path p
-        Ident.persistent id && Ident.name id = "Lwt"
-    | _ ->
-        false
 
-let rec is_async_t typ =
+(* Convert a path to a long identifier. *)
+let rec longident_of_path path =
+  match path with
+    | Path.Pident id ->
+      Longident.Lident (Ident.name id)
+    | Path.Pdot (path, s, _) ->
+      Longident.Ldot (longident_of_path path, s)
+    | Path.Papply (p1, p2) ->
+      Longident.Lapply (longident_of_path p1, longident_of_path p2)
+
+(* Returns the rewrite rule associated to a type, if any. *)
+let rec rule_of_type typ =
   match typ.Types.desc with
     | Types.Tlink typ ->
-        is_async_t typ
+      rule_of_type typ
     | Types.Tconstr (path, _, _) ->
-        let n = Path.name path in
+      if is_persistent_path path then
-        (n = "Async_core.Deferred.t") || (n = "Async.Std.Deferred.t")
+        try
+          Some (Hashtbl.find rewrite_rules (longident_of_path path))
+        with Not_found ->
+          rule_of_alias path
+      else
+        rule_of_alias path
     | _ ->
+      None
+
+and rule_of_alias path =
+  match try Some (Env.find_type path !Toploop.toplevel_env) with Not_found -> None with
+    | Some {
+        Types.type_kind = Types.Type_abstract;
+        Types.type_private = Asttypes.Public;
+        Types.type_manifest = Some typ;
+      } ->
+      rule_of_type typ
+    | _ ->
+      None
+
+(* Check that the given long identifier is present in the environment
+   and is persistent. *)
+let is_persistent_in_env longident =
+  try
+    is_persistent_path (fst (Env.lookup_value longident !Toploop.toplevel_env))
+  with Not_found ->
     false
 
 #if ocaml_version >= (4, 0, 0)
@@ -261,94 +355,31 @@ let str_items_of_typed_structure tstr = tstr
 let str_desc_of_typed_str_item tstr = tstr
 #endif
 
-let insert_lwt_main_run phrase =
+let rewrite_str_item pstr_item tstr_item =
-  match phrase with
-    | Parsetree.Ptop_def pstr ->
-        let env = !Toploop.toplevel_env in
-        let lwt_main_run_is_the_real_one =
-          try
-            match Env.lookup_value longident_lwt_main_run env with
-              | Path.Pdot (Path.Pident id, "run", 0), _ ->
-                  Ident.persistent id
-              | _ ->
-                  false
-          with Not_found ->
-            false
-        in
-        if lwt_main_run_is_the_real_one && List.exists is_eval pstr then
-          let tstr, _, _ = Typemod.type_structure env pstr Location.none in
-          let tstr = str_items_of_typed_structure tstr in
-          Parsetree.Ptop_def
-            (List.map2
-               (fun pstr_item tstr_item ->
   match pstr_item, str_desc_of_typed_str_item tstr_item with
-                    | { Parsetree.pstr_desc = Parsetree.Pstr_eval e; Parsetree.pstr_loc = loc },
+    | ({ Parsetree.pstr_desc = Parsetree.Pstr_eval e;
-                      Typedtree.Tstr_eval { Typedtree.exp_type = typ } when is_lwt_t typ ->
+         Parsetree.pstr_loc = loc },
-                        {
+       Typedtree.Tstr_eval { Typedtree.exp_type = typ }) -> begin
-                          Parsetree.pstr_desc =
+      match rule_of_type typ with
-                            Parsetree.Pstr_eval {
+        | Some rule ->
-                              Parsetree.pexp_desc =
+          if React.S.value rule.enabled && List.for_all is_persistent_in_env rule.required_values then
-                                Parsetree.Pexp_apply
+            { Parsetree.pstr_desc = Parsetree.Pstr_eval (rule.rewrite loc e);
-                                  ({ Parsetree.pexp_desc = Parsetree.Pexp_ident (with_loc loc longident_lwt_main_run);
+              Parsetree.pstr_loc = loc }
-                                     Parsetree.pexp_loc = loc },
-                                   [("", e)]);
-                              Parsetree.pexp_loc = loc;
-                            };
-                          Parsetree.pstr_loc = loc;
-                        }
-                    | _ ->
-                        pstr_item)
-               pstr tstr)
           else
-          phrase
+            pstr_item
-    | Parsetree.Ptop_dir _ ->
+        | None ->
-        phrase
+          pstr_item
+    end
+    | _ ->
+      pstr_item
 
-let insert_async_main_run phrase =
+let rewrite phrase =
-  let wrap_unit loc e =
-    let open Parsetree in
-    let i = with_loc loc longident_unit in
-    let p = { ppat_desc = Ppat_construct (i, None, false); ppat_loc=loc } in
-    { pexp_desc=Pexp_function ("", None, [p,e]); pexp_loc=loc } in
   match phrase with
     | Parsetree.Ptop_def pstr ->
-        let env = !Toploop.toplevel_env in
+      if (UTop.get_auto_run_lwt () || UTop.get_auto_run_async ()) && List.exists is_eval pstr then
-        let async_main_run_is_the_real_one =
+        let tstr, _, _ = Typemod.type_structure !Toploop.toplevel_env pstr Location.none in
-          try
-            let path, _ = Env.lookup_value longident_async_main_run env in
-            let rec is_persistent = function
-             | Path.Pident id -> Ident.persistent id
-             | Path.Pdot(t, _, _) -> is_persistent t
-             | Path.Papply(_, p) -> is_persistent p
-            in
-            is_persistent path
-          with Not_found ->
-            false
-        in
-        if async_main_run_is_the_real_one && List.exists is_eval pstr then
-          let tstr, _, _ = Typemod.type_structure env pstr Location.none in
         let tstr = str_items_of_typed_structure tstr in
-          Parsetree.Ptop_def
+        Parsetree.Ptop_def (List.map2 rewrite_str_item pstr tstr)
-            (List.map2
-               (fun pstr_item tstr_item ->
-                  match pstr_item, str_desc_of_typed_str_item tstr_item with
-                    | { Parsetree.pstr_desc = Parsetree.Pstr_eval e; Parsetree.pstr_loc = loc },
-                      Typedtree.Tstr_eval { Typedtree.exp_type = typ } when is_async_t typ ->
-                        {
-                          Parsetree.pstr_desc =
-                            Parsetree.Pstr_eval {
-                              Parsetree.pexp_desc =
-                                Parsetree.Pexp_apply
-                                  ({ Parsetree.pexp_desc = Parsetree.Pexp_ident (with_loc loc longident_async_main_run);
-                                     Parsetree.pexp_loc = loc },
-                                   [("", (wrap_unit loc e))]);
-                              Parsetree.pexp_loc = loc;
-                            };
-                          Parsetree.pstr_loc = loc;
-                        }
-                    | _ ->
-                        pstr_item)
-               pstr tstr)
       else
         phrase
     | Parsetree.Ptop_dir _ ->
@@ -405,10 +436,8 @@ let rec loop term =
 
   match phrase_opt with
     | Some phrase ->
-        (* Add Lwt_main.run to toplevel evals. *)
+        (* Rewrite toplevel expressions. *)
-        let phrase = if UTop.get_auto_run_lwt () then insert_lwt_main_run phrase else phrase in
+        let phrase = rewrite phrase in
-        (* Add Async execution to toplevel evals. *)
-        let phrase = if UTop.get_auto_run_async () then insert_async_main_run phrase else phrase in
         (* Set the margin of standard formatters. *)
         let cols = (LTerm.size term).cols in
         update_margin Format.std_formatter cols;
@@ -626,9 +655,8 @@ module Emacs(M : sig end) = struct
           send "accept" "";
           List.iter (send "stderr") (split_at ~trim:true '\n' warnings);
           if add_to_history then LTerm_history.add UTop.history input;
-          (* Add Lwt_main.run to toplevel evals. *)
+          (* Rewrite toplevel expressions. *)
-          let phrase = if UTop.get_auto_run_lwt () then insert_lwt_main_run phrase else phrase in
+          let phrase = rewrite phrase in
-          let phrase = if UTop.get_auto_run_async () then insert_async_main_run phrase else phrase in
           try
             ignore (Toploop.execute_phrase true Format.std_formatter phrase)
           with exn ->