Allow mutually recursive functions

5 files changed, 164 insertions, 81 deletions

diff --git a/src/ast_util.ml b/src/ast_util.ml
index d839b5a4..b0f4c052 100644
--- a/src/ast_util.ml
+++ b/src/ast_util.ml
@@ -306,6 +306,22 @@ let id_loc = function
 let kid_loc = function
   | Kid_aux (_, l) -> l
 
+let def_loc = function
+  | DEF_kind (KD_aux (_, (l, _)))
+  | DEF_type (TD_aux (_, (l, _)))
+  | DEF_fundef (FD_aux (_, (l, _)))
+  | DEF_val (LB_aux (_, (l, _)))
+  | DEF_spec (VS_aux (_, (l, _)))
+  | DEF_default (DT_aux (_, l))
+  | DEF_scattered (SD_aux (_, (l, _)))
+  | DEF_reg_dec (DEC_aux (_, (l, _))) ->
+    l
+  | DEF_internal_mutrec _
+  | DEF_comm _
+  | DEF_overload _
+  | DEF_fixity _ ->
+    Parse_ast.Unknown
+
 let string_of_id = function
   | Id_aux (Id v, _) -> v
   | Id_aux (DeIid v, _) -> "(deinfix " ^ v ^ ")"
diff --git a/src/ast_util.mli b/src/ast_util.mli
index abadd759..ef367e4e 100644
--- a/src/ast_util.mli
+++ b/src/ast_util.mli
@@ -140,6 +140,8 @@ val map_letbind_annot : ('a annot -> 'b annot) -> 'a letbind -> 'b letbind
 val id_loc : id -> Parse_ast.l
 val kid_loc : kid -> Parse_ast.l
 
+val def_loc : 'a def -> Parse_ast.l
+
 (* For debugging and error messages only: Not guaranteed to produce
    parseable SAIL, or even print all language constructs! *)
 (* TODO: replace with existing pretty-printer *)
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 2025b127..23fc8287 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -1523,81 +1523,90 @@ let get_id = function
 
 module StringSet = Set.Make(String)
 
-let rec doc_fundef_lem sequential mwords (FD_aux(FD_function(r, typa, efa, fcls),fannot)) =
+let doc_fundef_rhs_lem sequential mwords (FD_aux(FD_function(r, typa, efa, funcls),fannot) as fd) =
+  match funcls with
+  | [] -> failwith "FD_function with empty function list"
+  | [FCL_aux(FCL_Funcl(id,pat,exp),_)] ->
+    (prefix 2 1)
+      ((separate space)
+         [doc_id_lem id;
+          (doc_pat_lem sequential mwords true pat);
+          equals])
+      (doc_fun_body_lem sequential mwords exp)
+  | _ ->
+    let clauses =
+      (separate_map (break 1))
+        (fun fcl -> separate space [pipe;doc_funcl_lem sequential mwords fcl]) funcls in
+    (prefix 2 1)
+      ((separate space) [doc_id_lem (id_of_fundef fd);equals;string "function"])
+      (clauses ^/^ string "end")
+
+let doc_mutrec_lem sequential mwords = function
+  | [] -> failwith "DEF_internal_mutrec with empty function list"
+  | fundefs ->
+    string "let rec " ^^
+    separate_map (hardline ^^ string "and ")
+      (doc_fundef_rhs_lem sequential mwords) fundefs
+
+let rec doc_fundef_lem sequential mwords (FD_aux(FD_function(r, typa, efa, fcls),fannot) as fd) =
   match fcls with
   | [] -> failwith "FD_function with empty function list"
-  | [FCL_aux (FCL_Funcl(id,pat,exp),_)]
-    when not (Env.is_extern id (env_of exp)) ->
-     (prefix 2 1)
-       ((separate space)
-          [(string "let") ^^ (doc_rec_lem r) ^^ (doc_id_lem id);
-           (doc_pat_lem sequential mwords true pat);
-           equals])
-       (doc_fun_body_lem sequential mwords exp)
+  | FCL_aux (FCL_Funcl(id,_,exp),_) :: _
+    when string_of_id id = "execute" || string_of_id id = "initial_analysis" ->
+    let (_,auxiliary_functions,clauses) =
+      List.fold_left
+        (fun (already_used_fnames,auxiliary_functions,clauses) funcl ->
+          match funcl with
+          | FCL_aux (FCL_Funcl (Id_aux (Id _,l),pat,exp),annot) ->
+             let ctor, l, argspat, pannot = (match pat with
+               | P_aux (P_app (Id_aux (Id ctor,l),argspat),pannot) ->
+                 (ctor, l, argspat, pannot)
+               | P_aux (P_id (Id_aux (Id ctor,l)), pannot) ->
+                 (ctor, l, [], pannot)
+               | _ ->
+                 raise (Reporting_basic.err_unreachable l
+                   "unsupported parameter pattern in function clause")) in
+             let rec pick_name_not_clashing_with already_used candidate =
+               if StringSet.mem candidate already_used then
+                 pick_name_not_clashing_with already_used (candidate ^ "'")
+               else candidate in
+             let aux_fname = pick_name_not_clashing_with already_used_fnames (string_of_id id ^ "_" ^ ctor) in
+             let already_used_fnames = StringSet.add aux_fname already_used_fnames in
+             let fcl = FCL_aux (FCL_Funcl (Id_aux (Id aux_fname,l),
+                                           P_aux (P_tup argspat,pannot),exp),annot) in
+             let auxiliary_functions =
+                auxiliary_functions ^^ hardline ^^ hardline ^^
+                  doc_fundef_lem sequential mwords (FD_aux (FD_function(r,typa,efa,[fcl]),fannot)) in
+             (* Bind complex patterns to names so that we can pass them to the
+                auxiliary function *)
+             let name_pat idx (P_aux (p,a)) = match p with
+               | P_as (pat,_) -> P_aux (p,a) (* already named *)
+               | P_lit _ -> P_aux (p,a) (* no need to name a literal *)
+               | P_id _ -> P_aux (p,a) (* no need to name an identifier *)
+               | _ -> P_aux (P_as (P_aux (p,a), Id_aux (Id ("arg" ^ string_of_int idx),l)),a) in
+             let named_argspat = List.mapi name_pat argspat in
+             let named_pat = P_aux (P_app (Id_aux (Id ctor,l),named_argspat),pannot) in
+             let doc_arg idx (P_aux (p,(l,a))) = match p with
+               | P_as (pat,id) -> doc_id_lem id
+               | P_lit lit -> doc_lit_lem sequential mwords false lit a
+               | P_id id -> doc_id_lem id
+               | _ -> string ("arg" ^ string_of_int idx) in
+             let clauses =
+               clauses ^^ (break 1) ^^
+                 (separate space
+                    [pipe;doc_pat_lem sequential mwords false named_pat;arrow;
+                     string aux_fname;
+                     parens (separate comma (List.mapi doc_arg named_argspat))]) in
+             (already_used_fnames,auxiliary_functions,clauses)
+        ) (StringSet.empty,empty,empty) fcls in
+
+    auxiliary_functions ^^ hardline ^^ hardline ^^
+    (prefix 2 1)
+      ((separate space) [string "let" ^^ doc_rec_lem r ^^ doc_id_lem id;equals;string "function"])
+      (clauses ^/^ string "end")
   | FCL_aux (FCL_Funcl(id,_,exp),_) :: _
     when not (Env.is_extern id (env_of exp)) ->
-    (*    let sep = hardline ^^ pipe ^^ space in *)
-    (match id with
-    | Id_aux (Id fname,idl)
-      when fname = "execute" || fname = "initial_analysis" ->
-      let (_,auxiliary_functions,clauses) = 
-        List.fold_left 
-          (fun (already_used_fnames,auxiliary_functions,clauses) funcl ->
-            match funcl with
-            | FCL_aux (FCL_Funcl (Id_aux (Id _,l),pat,exp),annot) ->
-               let ctor, l, argspat, pannot = (match pat with
-                 | P_aux (P_app (Id_aux (Id ctor,l),argspat),pannot) ->
-                   (ctor, l, argspat, pannot)
-                 | P_aux (P_id (Id_aux (Id ctor,l)), pannot) ->
-                   (ctor, l, [], pannot)
-                 | _ ->
-                   raise (Reporting_basic.err_unreachable l
-                     "unsupported parameter pattern in function clause")) in
-               let rec pick_name_not_clashing_with already_used candidate =
-                 if StringSet.mem candidate already_used then
-                   pick_name_not_clashing_with already_used (candidate ^ "'")
-                 else candidate in
-               let aux_fname = pick_name_not_clashing_with already_used_fnames (fname ^ "_" ^ ctor) in
-               let already_used_fnames = StringSet.add aux_fname already_used_fnames in
-               let fcl = FCL_aux (FCL_Funcl (Id_aux (Id aux_fname,l),
-                                             P_aux (P_tup argspat,pannot),exp),annot) in
-               let auxiliary_functions = 
-                  auxiliary_functions ^^ hardline ^^ hardline ^^
-                    doc_fundef_lem sequential mwords (FD_aux (FD_function(r,typa,efa,[fcl]),fannot)) in
-               (* Bind complex patterns to names so that we can pass them to the
-                  auxiliary function *)
-               let name_pat idx (P_aux (p,a)) = match p with
-                 | P_as (pat,_) -> P_aux (p,a) (* already named *)
-                 | P_lit _ -> P_aux (p,a) (* no need to name a literal *)
-                 | P_id _ -> P_aux (p,a) (* no need to name an identifier *)
-                 | _ -> P_aux (P_as (P_aux (p,a), Id_aux (Id ("arg" ^ string_of_int idx),l)),a) in
-               let named_argspat = List.mapi name_pat argspat in
-               let named_pat = P_aux (P_app (Id_aux (Id ctor,l),named_argspat),pannot) in
-               let doc_arg idx (P_aux (p,(l,a))) = match p with
-                 | P_as (pat,id) -> doc_id_lem id
-                 | P_lit lit -> doc_lit_lem sequential mwords false lit a
-                 | P_id id -> doc_id_lem id
-                 | _ -> string ("arg" ^ string_of_int idx) in
-               let clauses =
-                 clauses ^^ (break 1) ^^
-                   (separate space
-                      [pipe;doc_pat_lem sequential mwords false named_pat;arrow;
-                       string aux_fname;
-                       parens (separate comma (List.mapi doc_arg named_argspat))]) in
-               (already_used_fnames,auxiliary_functions,clauses)
-          ) (StringSet.empty,empty,empty) fcls in
-
-      auxiliary_functions ^^ hardline ^^ hardline ^^
-      (prefix 2 1)
-        ((separate space) [string "let" ^^ doc_rec_lem r ^^ doc_id_lem id;equals;string "function"])
-        (clauses ^/^ string "end")
-    | _ ->
-      let clauses =
-        (separate_map (break 1))
-          (fun fcl -> separate space [pipe;doc_funcl_lem sequential mwords fcl]) fcls in
-      (prefix 2 1)
-        ((separate space) [string "let" ^^ doc_rec_lem r ^^ doc_id_lem id;equals;string "function"])
-        (clauses ^/^ string "end"))
+    string "let" ^^ (doc_rec_lem r) ^^ (doc_fundef_rhs_lem sequential mwords fd)
   | _ -> empty
 
 
@@ -1651,6 +1660,8 @@ let rec doc_def_lem sequential mwords def =
 
   | DEF_default df -> (empty,empty)
   | DEF_fundef f_def -> (empty,group (doc_fundef_lem sequential mwords f_def) ^/^ hardline)
+  | DEF_internal_mutrec fundefs ->
+     (empty, doc_mutrec_lem sequential mwords fundefs ^/^ hardline)
   | DEF_val lbind -> (empty,group (doc_let_lem sequential mwords false lbind) ^/^ hardline)
   | DEF_scattered sdef -> failwith "doc_def_lem: shoulnd't have DEF_scattered at this point"
 
diff --git a/src/rewriter.ml b/src/rewriter.ml
index 69beac4d..bcfb731a 100644
--- a/src/rewriter.ml
+++ b/src/rewriter.ml
@@ -579,6 +579,7 @@ let rewrite_fun rewriters (FD_aux (FD_function(recopt,tannotopt,effectopt,funcls
 let rewrite_def rewriters d = match d with
   | DEF_type _ | DEF_kind _ | DEF_spec _ | DEF_default _ | DEF_reg_dec _ | DEF_comm _ | DEF_overload _ | DEF_fixity _ -> d
   | DEF_fundef fdef -> DEF_fundef (rewriters.rewrite_fun rewriters fdef)
+  | DEF_internal_mutrec fdefs -> DEF_internal_mutrec (List.map (rewriters.rewrite_fun rewriters) fdefs)
   | DEF_val letbind -> DEF_val (rewriters.rewrite_let rewriters letbind)
   | DEF_scattered _ -> raise (Reporting_basic.err_unreachable Parse_ast.Unknown "DEF_scattered survived to rewritter")
 
@@ -2432,20 +2433,51 @@ let rewrite_fix_fun_effs (Defs defs) =
     (Bindings.add id fun_eff fun_effs,
      funcls @ [FCL_aux (FCL_Funcl (id, pat, exp), (l, annot))]) in
 
-  let rewrite_def (fun_effs, defs) = function
-  | DEF_fundef (FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) ->
+  let rewrite_fundef (fun_effs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) =
     let (fun_effs, funcls) = List.fold_left rewrite_funcl (fun_effs, []) funcls in
+    let is_funcl_rec (FCL_aux (FCL_Funcl (id, _, exp), _)) =
+      fst (fold_exp
+        { (compute_exp_alg false (||) ) with
+          e_app = (fun (f, es) ->
+            let (rs, es) = List.split es in
+            (List.fold_left (||) (string_of_id f = string_of_id id) rs,
+             E_app (f, es)));
+          e_app_infix = (fun ((r1,e1), f, (r2,e2)) ->
+            (r1 || r2 || (string_of_id f = string_of_id id),
+             E_app_infix (e1, f, e2))) } exp) in
+    let is_rec = List.exists is_funcl_rec funcls in
     (* Repeat once for recursive functions:
        propagates union of effects to all clauses *)
-    let (fun_effs, funcls) = List.fold_left rewrite_funcl (fun_effs, []) funcls in
-    (fun_effs, defs @ [DEF_fundef (FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a))])
+    let recopt, (fun_effs, funcls) =
+      if is_rec then
+        Rec_aux (Rec_rec, Parse_ast.Unknown),
+        List.fold_left rewrite_funcl (fun_effs, []) funcls
+      else recopt, (fun_effs, funcls) in
+    (fun_effs, FD_aux (FD_function (recopt, tannotopt, effopt, funcls), a)) in
+
+  let rec rewrite_fundefs (fun_effs, fundefs) =
+    match fundefs with
+    | fundef :: fundefs ->
+      let (fun_effs, fundef) = rewrite_fundef (fun_effs, fundef) in
+      let (fun_effs, fundefs) = rewrite_fundefs (fun_effs, fundefs) in
+      (fun_effs, fundef :: fundefs)
+    | [] -> (fun_effs, []) in
+
+  let rewrite_def (fun_effs, defs) = function
+  | DEF_fundef fundef ->
+    let (fun_effs, fundef) = rewrite_fundef (fun_effs, fundef) in
+    (fun_effs, defs @ [DEF_fundef fundef])
+  | DEF_internal_mutrec fundefs ->
+    let (fun_effs, fundefs) = rewrite_fundefs (fun_effs, fundefs) in
+    (fun_effs, defs @ [DEF_internal_mutrec fundefs])
   | DEF_val (LB_aux (LB_val (pat, exp), a)) ->
     (fun_effs, defs @ [DEF_val (LB_aux (LB_val (pat, rewrite_exp fun_effs exp), a))])
   | def -> (fun_effs, defs @ [def]) in
 
-  if !Type_check.opt_no_effects
-  then Defs (snd (List.fold_left rewrite_def (Bindings.empty, []) defs))
-  else Defs defs
+  (* if !Type_check.opt_no_effects
+  then *)
+  Defs (snd (List.fold_left rewrite_def (Bindings.empty, []) defs))
+  (* else Defs defs *)
 
 (* Turn constraints into numeric expressions with sizeof *)
 let rewrite_constraint =
@@ -3018,6 +3050,8 @@ let rewrite_defs_letbind_effects  =
           rewrap (LB_val (pat, n_exp_term (effectful exp) exp))
       end
     | DEF_fundef fdef -> DEF_fundef (rewrite_fun rewriters fdef)
+    | DEF_internal_mutrec fdefs ->
+      DEF_internal_mutrec (List.map (rewrite_fun rewriters) fdefs)
     | d -> d in
   rewrite_defs_base
     {rewrite_exp = rewrite_exp
diff --git a/src/spec_analysis.ml b/src/spec_analysis.ml
index c9e4c6e7..0d299988 100644
--- a/src/spec_analysis.ml
+++ b/src/spec_analysis.ml
@@ -315,6 +315,7 @@ let rec fv_of_exp consider_var bound used set (E_aux (e,(_,tannot))) : (Nameset.
   | E_for(id,from,to_,by,_,body) ->
     let _,used,set = list_fv bound used set [from;to_;by] in
     fv_of_exp consider_var (Nameset.add (string_of_id id) bound) used set body
+  | E_loop(_, cond, body) -> list_fv bound used set [cond; body]
   | E_vector_access(v,i) -> list_fv bound used set [v;i]
   | E_vector_subrange(v,i1,i2) -> list_fv bound used set [v;i1;i2]
   | E_vector_update(v,i,e) -> list_fv bound used set [v;i;e]
@@ -457,6 +458,7 @@ let fv_of_fun consider_var (FD_aux (FD_function(rec_opt,tannot_opt,_,funcls),_))
       let pat_bs,pat_ns = pat_bindings consider_var base_bounds ns pat in
       let _, exp_ns,_ = fv_of_exp consider_var pat_bs pat_ns Nameset.empty exp in
       exp_ns) funcls mt in
+  (* let _ = Printf.eprintf "Function %s uses %s\n" fun_name (set_to_string (Nameset.union ns ns_r)) in *)
   init_env fun_name,Nameset.union ns ns_r
 
 let fv_of_vspec consider_var (VS_aux(vspec,_)) = match vspec with
@@ -575,6 +577,20 @@ let rec print_dependencies orig_queue work_queue names =
            print_dependencies orig_queue orig_queue uses));
     print_dependencies orig_queue wq names
 
+let merge_mutrecs defs =
+  let merge_aux ((binds', uses'), def) ((binds, uses), fundefs) =
+    let fundefs = match def with
+      | DEF_fundef fundef -> fundef :: fundefs
+      | DEF_internal_mutrec fundefs' -> fundefs' @ fundefs
+      | _ ->
+        (* let _ = Pretty_print_sail2.pp_defs stderr (Defs [def]) in *)
+        raise (Reporting_basic.err_unreachable (def_loc def)
+          "Trying to merge non-function definition with mutually recursive functions") in
+    (* let _ = Printf.eprintf " - Merging %s (using %s)\n" (set_to_string binds') (set_to_string uses') in *)
+    ((Nameset.union binds' binds, Nameset.union uses' uses), fundefs) in
+  let ((binds, uses), fundefs) = List.fold_right merge_aux defs ((mt, mt), []) in
+  ((binds, uses), DEF_internal_mutrec fundefs)
+
 let rec topological_sort work_queue defs =
   match work_queue with
   | [] -> List.rev defs
@@ -588,11 +604,15 @@ let rec topological_sort work_queue defs =
     else
       match List.stable_sort compare_dbts work_queue with (*We wait to sort until there are no 0 dependency nodes on top*)
       | [] -> failwith "sort shrunk the list???"
-      | (((n,uses),_)::_) as wq ->
+      | (((n,uses),def)::rest) as wq ->
         if (Nameset.cardinal uses = 0)
         then topological_sort wq defs
-        else let _ = Printf.eprintf "Uses on failure are %s, binds are %s\n" (set_to_string uses) (set_to_string n)
-          in let _ = print_dependencies wq wq uses in failwith "A dependency was unmet"
+        else
+          let _ = Printf.eprintf "Merging (potentially) mutually recursive definitions %s and %s\n" (set_to_string n) (set_to_string uses) in
+          let is_used ((binds', uses'), def') = not(Nameset.is_empty(Nameset.inter binds' uses)) in
+          let (used, rest) = List.partition is_used rest in
+          let wq = merge_mutrecs (((n,uses),def)::used) :: rest in
+          topological_sort wq defs
 
 let rec add_to_partial_order ((binds,uses),def) = function
   | [] ->