Initial rewrite to move complex nexps in fn sigs into constraints

(for monomorphisation, off for now because the analysis needs extended).
Also tighten up orig_nexp, make Lem backend replace # in type variables.

7 files changed, 199 insertions, 9 deletions

diff --git a/src/monomorphise.ml b/src/monomorphise.ml
index be8ff5d4..8ff5dc20 100644
--- a/src/monomorphise.ml
+++ b/src/monomorphise.ml
@@ -3817,10 +3817,172 @@ let add_bitvector_casts (Defs defs) =
   in Defs (List.map rewrite_def defs)
 end
 
+let replace_nexp_in_typ env typ orig new_nexp =
+  let rec aux (Typ_aux (t,l) as typ) =
+    match t with
+    | Typ_id _
+    | Typ_var _
+        -> false, typ
+    | Typ_fn (arg,res,eff) ->
+       let f1, arg = aux arg in
+       let f2, res = aux res in
+       f1 || f2, Typ_aux (Typ_fn (arg, res, eff),l)
+    | Typ_tup typs ->
+       let fs, typs = List.split (List.map aux typs) in
+       List.exists (fun x -> x) fs, Typ_aux (Typ_tup typs,l)
+    | Typ_exist (kids,nc,typ') -> (* TODO avoid capture *)
+       let f, typ' = aux typ' in
+       f, Typ_aux (Typ_exist (kids,nc,typ'),l)
+    | Typ_app (id, targs) ->
+       let fs, targs = List.split (List.map aux_targ targs) in
+       List.exists (fun x -> x) fs, Typ_aux (Typ_app (id, targs),l)
+  and aux_targ (Typ_arg_aux (ta,l) as typ_arg) =
+    match ta with
+    | Typ_arg_nexp nexp ->
+       if prove env (nc_eq nexp orig)
+       then true, Typ_arg_aux (Typ_arg_nexp new_nexp,l)
+       else false, typ_arg
+    | Typ_arg_typ typ ->
+       let f, typ = aux typ in
+       f, Typ_arg_aux (Typ_arg_typ typ,l)
+    | Typ_arg_order _ -> false, typ_arg
+  in aux typ
+
+(* TODO: replace with more informative mangled name *)
+let fresh_top_kid =
+  let counter = ref 0 in
+  fun () ->
+    let n = !counter in
+    let () = counter := n+1 in
+    mk_kid ("mono#nexp#" ^ string_of_int n)
+
+let rewrite_toplevel_nexps (Defs defs) =
+  let find_nexp env nexp_map nexp =
+    let is_equal (kid,nexp') = prove env (nc_eq nexp nexp') in
+    List.find is_equal nexp_map
+  in
+  let rec rewrite_typ_in_spec env nexp_map (Typ_aux (t,ann) as typ_full) =
+    match t with
+    | Typ_fn (args,res,eff) ->
+       let nexp_map, args = rewrite_typ_in_spec env nexp_map args in
+       let nexp_map, res = rewrite_typ_in_spec env nexp_map res in
+       nexp_map, Typ_aux (Typ_fn (args,res,eff),ann)
+    | Typ_tup typs ->
+       let nexp_map, typs =
+         List.fold_right (fun typ (nexp_map,t) ->
+           let nexp_map, typ = rewrite_typ_in_spec env nexp_map typ in
+           (nexp_map, typ::t)) typs (nexp_map,[])
+       in nexp_map, Typ_aux (Typ_tup typs,ann)
+    | _ ->
+       let typ' = Env.base_typ_of env typ_full in
+       let nexp_opt =
+         match destruct_atom_nexp env typ' with
+         | Some nexp -> Some nexp
+         | None ->
+            if is_bitvector_typ typ' then
+              let (_,size,_,_) = vector_typ_args_of typ' in
+              Some size
+            else None
+       in match nexp_opt with
+       | None -> nexp_map, typ_full
+       | Some (Nexp_aux (Nexp_constant _,_))
+       | Some (Nexp_aux (Nexp_var _,_)) -> nexp_map, typ_full
+       | Some nexp ->
+          let nexp_map, kid =
+            match find_nexp env nexp_map nexp with
+            | (kid,_) -> nexp_map, kid
+            | exception Not_found ->
+               let kid = fresh_top_kid () in
+               (kid, nexp)::nexp_map, kid
+          in
+          let new_nexp = nvar kid in
+          (* Try to avoid expanding the original type *)
+          let changed, typ = replace_nexp_in_typ env typ_full nexp new_nexp in
+          if changed then nexp_map, typ
+          else nexp_map, snd (replace_nexp_in_typ env typ' nexp new_nexp)
+  in
+  let rewrite_valspec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (tqs,typ),ts_l),id,ext_opt,is_cast),ann)) =
+    match tqs, ext_opt "lem" with (* does it really matter if this is wrong? *)
+    | _, Some _
+    | TypQ_aux (TypQ_no_forall,_), _ -> None
+    | TypQ_aux (TypQ_tq qs, tq_l), None ->
+       let env = env_of_annot ann in
+       let env = add_typquant tqs env in
+       let nexp_map, typ = rewrite_typ_in_spec env [] typ in
+       match nexp_map with
+       | [] -> None
+       | _ ->
+          let new_vars = List.map (fun (kid,nexp) -> QI_aux (QI_id (KOpt_aux (KOpt_none kid,Generated Unknown)), Generated tq_l)) nexp_map in
+          let new_constraints = List.map (fun (kid,nexp) -> QI_aux (QI_const (nc_eq (nvar kid) nexp), Generated tq_l)) nexp_map in
+          let tqs = TypQ_aux (TypQ_tq (qs @ new_vars @ new_constraints),tq_l) in
+          let vs =
+            VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (tqs,typ),ts_l),id,ext_opt,is_cast),ann) in
+          Some (id, nexp_map, vs)
+  in
+  let rewrite_typ_in_body env nexp_map typ =
+    let rec aux (Typ_aux (t,l) as typ_full) =
+    match t with
+    | Typ_tup typs -> Typ_aux (Typ_tup (List.map aux typs),l)
+    | Typ_exist (kids,nc,typ') -> (* TODO: avoid shadowing *)
+       Typ_aux (Typ_exist (kids,(* TODO? *) nc, aux typ'),l)
+    | Typ_app (id,targs) -> Typ_aux (Typ_app (id,List.map aux_targ targs),l)
+    | _ -> typ_full
+    and aux_targ (Typ_arg_aux (ta,l) as ta_full) =
+      match ta with
+      | Typ_arg_typ typ -> Typ_arg_aux (Typ_arg_typ (aux typ),l)
+      | Typ_arg_order _ -> ta_full
+      | Typ_arg_nexp nexp ->
+         match find_nexp env nexp_map nexp with
+         | (kid,_) -> Typ_arg_aux (Typ_arg_nexp (nvar kid),l)
+         | exception Not_found -> ta_full
+    in aux typ
+  in
+  let rewrite_one_exp nexp_map (e,ann) =
+    match e with
+    | E_cast (typ,e') -> E_aux (E_cast (rewrite_typ_in_body (env_of_annot ann) nexp_map typ,e'),ann)
+    | E_sizeof nexp ->
+       (match find_nexp (env_of_annot ann) nexp_map nexp with
+       | (kid,_) -> E_aux (E_sizeof (nvar kid),ann)
+       | exception Not_found -> E_aux (e,ann))
+    | _ -> E_aux (e,ann)
+  in
+  let rewrite_one_pat nexp_map (p,ann) =
+    match p with
+    | P_typ (typ,p') -> P_aux (P_typ (rewrite_typ_in_body (env_of_annot ann) nexp_map typ,p'),ann)
+    | _ -> P_aux (p,ann)
+  in
+  let rewrite_body nexp_map pexp =
+    let open Rewriter in
+    fold_pexp { id_exp_alg with
+      e_aux = rewrite_one_exp nexp_map;
+      pat_alg = { id_pat_alg with p_aux = rewrite_one_pat nexp_map }
+    } pexp
+  in
+  let rewrite_funcl spec_map (FCL_aux (FCL_Funcl (id,pexp),ann) as funcl) =
+    match Bindings.find id spec_map with
+    | nexp_map -> FCL_aux (FCL_Funcl (id,rewrite_body nexp_map pexp),ann)
+    | exception Not_found -> funcl
+  in
+  let rewrite_def spec_map def =
+    match def with
+    | DEF_spec vs -> (match rewrite_valspec vs with
+      | None -> spec_map, def
+      | Some (id, nexp_map, vs) -> Bindings.add id nexp_map spec_map, DEF_spec vs)
+    | DEF_fundef (FD_aux (FD_function (recopt,tann,eff,funcls),ann)) ->
+       spec_map,
+       DEF_fundef (FD_aux (FD_function (recopt,tann,eff,List.map (rewrite_funcl spec_map) funcls),ann))
+    | _ -> spec_map, def
+  in
+  let _, defs = List.fold_left (fun (spec_map,t) def ->
+    let spec_map, def = rewrite_def spec_map def in
+    (spec_map, def::t)) (Bindings.empty, []) defs
+  in Defs (List.rev defs)
+
 type options = {
   auto : bool;
   debug_analysis : int;
   rewrites : bool;
+  rewrite_toplevel_nexps : bool;
   rewrite_size_parameters : bool;
   all_split_errors : bool;
   continue_anyway : bool;
@@ -3835,14 +3997,17 @@ let recheck defs =
   r
 
 let monomorphise opts splits env defs =
+  let defs =
+    if opts.rewrite_toplevel_nexps
+      then rewrite_toplevel_nexps defs
+      else defs
+  in
   let (defs,env) =
     if opts.rewrites then
       let defs = MonoRewrites.mono_rewrite defs in
-    (* TODO: is this necessary? *)
       recheck defs
-    else (defs,env)
+    else recheck defs
   in
-(*let _ = Pretty_print.pp_defs stdout defs in*)
   let ok_analysis, new_splits, extra_splits =
     if opts.auto
     then
diff --git a/src/monomorphise.mli b/src/monomorphise.mli
index 3e561e32..3baa7d12 100644
--- a/src/monomorphise.mli
+++ b/src/monomorphise.mli
@@ -52,6 +52,7 @@ type options = {
   auto : bool;                    (* Analyse ast to find splits for monomorphisation *)
   debug_analysis : int;           (* Debug output level for the automatic analysis *)
   rewrites : bool;                (* Experimental rewrites for variable-sized operations *)
+  rewrite_toplevel_nexps : bool;  (* Move complex nexps in function signatures into constraints *)
   rewrite_size_parameters : bool; (* Make implicit type parameters explicit for (e.g.) lem *)
   all_split_errors : bool;
   continue_anyway : bool;
diff --git a/src/pretty_print_lem.ml b/src/pretty_print_lem.ml
index 79eada3b..97b618e8 100644
--- a/src/pretty_print_lem.ml
+++ b/src/pretty_print_lem.ml
@@ -170,15 +170,19 @@ let is_regtyp (Typ_aux (typ, _)) env = match typ with
   | _ -> false
 
 let doc_nexp_lem nexp =
+  let nice_kid kid =
+    let (Kid_aux (Var kid,l)) = orig_kid kid in
+    Kid_aux (Var (String.map (function '#' -> '_' | c -> c) kid),l)
+  in
   let (Nexp_aux (nexp, l) as full_nexp) = nexp_simp nexp in
   match nexp with
   | Nexp_constant i -> string ("ty" ^ Big_int.to_string i)
-  | Nexp_var v -> string (string_of_kid (orig_kid v))
+  | Nexp_var v -> string (string_of_kid (nice_kid v))
   | _ ->
      let rec mangle_nexp (Nexp_aux (nexp, _)) = begin
        match nexp with
        | Nexp_id id -> string_of_id id
-       | Nexp_var kid -> string_of_id (id_of_kid (orig_kid kid))
+       | Nexp_var kid -> string_of_id (id_of_kid (nice_kid kid))
        | Nexp_constant i -> Pretty_print_lem_ast.lemnum Big_int.to_string i
        | Nexp_times (n1, n2) -> mangle_nexp n1 ^ "_times_" ^ mangle_nexp n2
        | Nexp_sum (n1, n2) -> mangle_nexp n1 ^ "_plus_" ^ mangle_nexp n2
diff --git a/src/process_file.ml b/src/process_file.ml
index ffd61fae..e9d04c5b 100644
--- a/src/process_file.ml
+++ b/src/process_file.ml
@@ -190,6 +190,7 @@ let opt_ddump_raw_mono_ast = ref false
 let opt_dmono_analysis = ref 0
 let opt_auto_mono = ref false
 let opt_mono_rewrites = ref false
+let opt_mono_complex_nexps = ref false
 let opt_dall_split_errors = ref false
 let opt_dmono_continue = ref false
 
@@ -199,6 +200,7 @@ let monomorphise_ast locs type_env ast =
     auto = !opt_auto_mono;
     debug_analysis = !opt_dmono_analysis;
     rewrites = !opt_mono_rewrites;
+    rewrite_toplevel_nexps = !opt_mono_complex_nexps;
     rewrite_size_parameters = !Pretty_print_lem.opt_mwords;
     all_split_errors = !opt_dall_split_errors;
     continue_anyway = !opt_dmono_continue;
diff --git a/src/process_file.mli b/src/process_file.mli
index 54415621..2f418ff9 100644
--- a/src/process_file.mli
+++ b/src/process_file.mli
@@ -75,6 +75,7 @@ val opt_dmono_analysis : int ref
 val opt_dmono_continue : bool ref
 val opt_auto_mono : bool ref
 val opt_mono_rewrites : bool ref
+val opt_mono_complex_nexps : bool ref
 val opt_dall_split_errors : bool ref
 
 type out_type =
diff --git a/src/sail.ml b/src/sail.ml
index 61e38f8a..2632f9a4 100644
--- a/src/sail.ml
+++ b/src/sail.ml
@@ -164,6 +164,9 @@ let options = Arg.align ([
   ( "-mono_rewrites",
     Arg.Set Process_file.opt_mono_rewrites,
     " turn on rewrites for combining bitvector operations");
+  ( "-mono_complex_nexps",
+    Arg.Set Process_file.opt_mono_complex_nexps,
+    " move complex size expressions in function signatures into constraints");
   ( "-dall_split_errors",
     Arg.Set Process_file.opt_dall_split_errors,
     " display all case split errors from monomorphisation, rather than one");
diff --git a/src/type_check.ml b/src/type_check.ml
index 06fbed0a..760345a6 100644
--- a/src/type_check.ml
+++ b/src/type_check.ml
@@ -140,7 +140,9 @@ let string_of_bind (typquant, typ) = string_of_typquant typquant ^ ". " ^ string
 let orig_kid (Kid_aux (Var v, l) as kid) =
   try
     let i = String.rindex v '#' in
-    Kid_aux (Var ("'" ^ String.sub v (i + 1) (String.length v - i - 1)), l)
+    if i >= 3 && String.sub v 0 3 = "'fv" then
+      Kid_aux (Var ("'" ^ String.sub v (i + 1) (String.length v - i - 1)), l)
+    else kid
   with
   | Not_found -> kid
 
@@ -3553,7 +3555,13 @@ let infer_funtyp l env tannotopt funcls =
      end
   | Typ_annot_opt_aux (Typ_annot_opt_none, _) -> typ_error l "Cannot infer function type for unannotated function"
 
-let mk_val_spec typq typ id = DEF_spec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (typq, typ), Parse_ast.Unknown), id, (fun _ -> None), false), (Parse_ast.Unknown, None)))
+let mk_val_spec env typq typ id =
+  let eff =
+    match typ with
+    | Typ_aux (Typ_fn (_,_,eff),_) -> eff
+    | _ -> no_effect
+  in
+  DEF_spec (VS_aux (VS_val_spec (TypSchm_aux (TypSchm_ts (typq, typ), Parse_ast.Unknown), id, (fun _ -> None), false), (Parse_ast.Unknown, Some (env,typ,eff))))
 
 let check_tannotopt env typq ret_typ = function
   | Typ_annot_opt_aux (Typ_annot_opt_none, _) -> ()
@@ -3595,7 +3603,7 @@ let check_fundef env (FD_aux (FD_function (recopt, tannotopt, effectopt, funcls)
     if not have_val_spec
     then
       let typ = Typ_aux (Typ_fn (vtyp_arg, vtyp_ret, eff), vl) in
-      [mk_val_spec quant typ id], Env.add_val_spec id (quant, typ) env, eff
+      [mk_val_spec env quant typ id], Env.add_val_spec id (quant, typ) env, eff
     else [], env, declared_eff
   in
   let env = Env.define_val_spec id env in
@@ -3608,6 +3616,7 @@ let check_fundef env (FD_aux (FD_function (recopt, tannotopt, effectopt, funcls)
    context. We have to destructure the various kinds of val specs, but
    the difference is irrelevant for the typechecker. *)
 let check_val_spec env (VS_aux (vs, (l, _))) =
+  let annotate vs typ eff = DEF_spec (VS_aux (vs, (l, Some (env,typ,eff)))) in
   let (id, quants, typ, env) = match vs with
     | VS_val_spec (TypSchm_aux (TypSchm_ts (quants, typ), _) as typschm, id, ext_opt, is_cast) ->
        typ_debug ("VS typschm: " ^ string_of_id id ^ ", " ^ string_of_typschm typschm);
@@ -3623,7 +3632,12 @@ let check_val_spec env (VS_aux (vs, (l, _))) =
        let env = if is_cast then Env.add_cast id env else env in
        (id, quants, typ, env)
   in
-  [DEF_spec (VS_aux (vs, (l, None)))], Env.add_val_spec id (quants, Env.expand_synonyms (add_typquant quants env) typ) env
+  let eff =
+    match typ with
+    | Typ_aux (Typ_fn (_,_,eff),_) -> eff
+    | _ -> no_effect
+  in
+  [annotate vs typ eff], Env.add_val_spec id (quants, Env.expand_synonyms (add_typquant quants env) typ) env
 
 let check_default env (DT_aux (ds, l)) =
   match ds with