Source file sct_checker_forward.ml

open Utils
open Annotations
open Prog
open Constraints

module CT = Ct_checker_forward

module S = Syntax

(* ----------------------------------------------------------- *)
let pp_var fmt x = Printer.pp_var ~debug:false fmt x

let pp_var_i fmt x = pp_var fmt (L.unloc x)

let pp_expr fmt e = Printer.pp_expr ~debug:false fmt e

let pp_lval fmt x = Printer.pp_lval ~debug:false fmt x

let pp_vset fmt xs =
  Format.fprintf fmt "{@[ %a @]}"  (pp_list ",@ " pp_var) (Sv.elements xs)


(* ----------------------------------------------------------- *)
let ssecret = "secret"
let spublic = "public"
let stransient = "transient"
let smsf = "msf"

(* ----------------------------------------------------------- *)
(* Info provided by the user                                   *)
type ulevel =
  | Secret
  | Transient
  | Public
  | Msf

(* -------------------------------------------------------------- *)
(* Special operators to deal with msf                             *)
type special_op =
  | Init_msf
  | Update_msf
  | Mov_msf
  | Protect
  | Other

let is_special o =
  match o with
  | Sopn.Opseudo_op (Pseudo_operator.Ospill _) -> assert false
  | Sopn.Opseudo_op _ -> Other
  | Oasm _ -> Other
  | Oslh o ->
    match o with
    | SLHinit   -> Init_msf
    | SLHupdate -> Update_msf
    | SLHmove   -> Mov_msf
    | SLHprotect _ | SLHprotect_ptr _ -> Protect
    | SLHprotect_ptr_fail _ -> assert false

(* -----------------------------------------------------------*)
(* Types generated by the system                              *)

(* Attempt of convention: n stands for normal (non-speculative),
   s for speculative *)

type vty =
  | Direct of VlPairs.t
  | Indirect of VlPairs.t * VlPairs.t

let pp_vlpair ~top fmt (n, s) =
  if Lvl.is_public n && Lvl.is_public s then
    Format.fprintf fmt "%s%s" (if top then "#" else "") spublic
  else if Lvl.is_public n && Lvl.is_secret s then Format.fprintf fmt "%s%s" (if top then "#" else "") stransient
  else if Lvl.is_secret n && Lvl.is_secret s then Format.fprintf fmt "%s%s" (if top then "#" else "")ssecret
  else
    Format.fprintf fmt "%s{ n = %a, s = %a}" (if top then "#poly = " else "") Lvl.pp n Lvl.pp s

let pp_vty fmt = function
  | Direct vlp -> pp_vlpair ~top:true fmt vlp
  | Indirect(p_vlp, v_vlp) ->
      Format.fprintf fmt "#[ptr = %a, val = %a]" (pp_vlpair ~top:false) p_vlp (pp_vlpair ~top:false) v_vlp

type vfty =
  | IsMsf
  | IsNormal of vty

let pp_vfty fmt = function
  | IsMsf -> Format.fprintf fmt "#%s" smsf
  | IsNormal ty -> pp_vty fmt ty

(* Either a function does not modify the MSF, or it does at a certain location,
   and if that location is a function call we add the trace of calls until the
   offending instruction. *)
type modmsf =
  | Modified of L.i_loc * (L.i_loc * funname) list
  | NotModified

let is_Modified m =
  match m with
  | Modified _ -> true
  | NotModified -> false

type ty_fun = {
    modmsf               : modmsf;
    tyin                 : vfty list;
    tyout                : vfty list;
    constraints          : C.constraints;
    resulting_corruption : VlPairs.t (* resulting memory corruption after function call *)
  }

type ('info,'asm) fenv = {
    env_ty     : ty_fun Hf.t;
    env_def    : ('info,'asm) func list;
  }

module FEnv = struct
  let get_fun_def fenv fn = List.find (fun f -> F.equal f.f_name fn) fenv.env_def

  let get_fty fenv fn =
    try Hf.find fenv.env_ty fn with Not_found -> assert false

end

let pp_modmsf fmt modmsf =
  let s =
    match modmsf with
    | Modified _ -> "modmsf"
    | NotModified -> "nomodmsf"
  in
  Format.fprintf fmt "%s" s

let pp_funty fmt (fname, tyfun) =
  Format.fprintf fmt
    "@[<v>%a %s : @[%a@] ->@ @[%a@]@ \
      output corruption: %a@.\
      @ constraints:@ @[%a@]@]@."
    pp_modmsf tyfun.modmsf
    fname
    (pp_list " *@ " pp_vfty) tyfun.tyin
    (pp_list " *@ " pp_vfty) tyfun.tyout
    pp_vty (Direct (tyfun.resulting_corruption))
    C.pp tyfun.constraints


(* --------------------------------------------------------------- *)
(* Checks if the status of msf is modified by the instructions     *)

let is_inline i =
  match Annot.ensure_uniq1 "inline" Annot.none i.i_annot with
  | Some _ -> true
  | None   -> false

let rec modmsf_i fenv i =
  let modified_here = Modified(i.i_loc, []) in
  match i.i_desc with
  | Csyscall _ | Cwhile _ -> modified_here
  | Cif(_, c0, c1) ->
    if is_inline i
    then
      let r = modmsf_c fenv c0 in
      if is_Modified r then r else modmsf_c fenv c1
    else modified_here
  | Cassgn _ -> NotModified
  | Copn (_, _, o, _) ->
    begin match is_special o with
    | Init_msf -> modified_here (* LFENCE modifies msf *)
    | Update_msf -> modified_here (* not sure it is needed *)
    | Mov_msf | Protect | Other -> NotModified
    end
  | Cfor(_, _, c) -> modmsf_c fenv c
  | Ccall (_, f, _) ->
    match (FEnv.get_fty fenv f).modmsf with
    | Modified (l, tr) -> Modified(i.i_loc, (l, f) :: tr)
    | NotModified -> NotModified

and modmsf_c fenv c =
  List.map (modmsf_i fenv) c
  |> List.find_opt is_Modified
  |> Option.default NotModified

let error ~loc =
  hierror ~loc:(Lone loc) ~kind:"speculative constant type checker"

let warn ~loc = warning SCTchecker (L.i_loc0 loc)

(* --------------------------------------------------------- *)
(* Inference of the variables that need to contain msf       *)
(* This code does not need to be trusted it is simply        *)
(* used as an oracle                                         *)

let is_register ~direct x =
  is_reg_kind x.v_kind && (not direct || is_reg_direct_kind x.v_kind)

let ensure_register ~direct x =
  if not (is_register ~direct (L.unloc x)) then
    error ~loc:(L.loc x)
      "variable %a must be a reg%s" pp_var_i x
      (if direct then "" else " (ptr)")

let reg_lval ~direct loc x =
  match x with
  | Laset (_, _, _, x, _)
  | Lvar x -> ensure_register ~direct x; x
  | _ ->
      error ~loc "L-value %a must be a reg%s" pp_lval x
        (if direct then "" else " (ptr)")

let reg_lval_opt ~direct loc =
  function
  | Lnone _ -> None
  | x -> Some (reg_lval ~direct loc x)

let reg_expr_opt ~direct = function
  | Pget (_, _, _, x, _) | Pvar x ->
      if is_gkvar x && is_register ~direct (L.unloc x.gv)
      then Some x.gv
      else None
  | _ -> None

let reg_expr ~direct loc e =
  match reg_expr_opt ~direct e with
  | Some x -> x
  | None -> error ~loc "expression %a must be a reg%s" pp_expr e
        (if direct then "" else " (ptr)")


let rec infer_msf_i ~withcheck fenv (tbl:(L.i_loc, Sv.t) Hashtbl.t) i ms =
  let loc = i.i_loc.L.base_loc in

  let check_x ms x =
    if Sv.mem (L.unloc x) ms then
      error ~loc:(L.loc x)
        "instruction assigns %a, which is required to be a msf"
        pp_var_i x in

  let check ms x =
    match x with
    | Lvar x -> check_x ms x
    | _ -> () in

  let checks ms xs = List.iter (check ms) xs in

  let pp_modmsf_trace fmt tr =
    let pp_item fmt (l, fn) =
      Format.fprintf fmt
        "@[<h>the function %s destroys MSFs at %a@]"
        fn.fn_name
        L.pp_iloc l
    in
    Format.fprintf fmt "Trace:@;<0 2>@[<v>%a@]" (pp_list "" pp_item) tr
  in

  let check_call ~loc fn modmsf ms =
    match modmsf with
    | Modified(l, tr) ->
      if not (Sv.is_empty ms) && withcheck then
        error ~loc
          "@[<h>this function call destroys MSFs and %a are required.@]@;%a"
          pp_vset ms
          pp_modmsf_trace ((l, fn) :: tr)
    | NotModified -> ()
  in

  match i.i_desc with
  | Csyscall _ ->
      if not (Sv.is_empty ms) && withcheck then
        error ~loc "syscalls destroy msf variables, %a are required" pp_vset ms;
      (* withcheck => is_empty ms *)
      ms

  | Cif (_, c1, c2) ->
    let ms1 = infer_msf_c ~withcheck fenv tbl c1 ms in
    let ms2 = infer_msf_c ~withcheck fenv tbl c2 ms in
    Sv.union ms1 ms2

  | Cfor(x, _, c) ->
    check_x ms x;
    let rec loop ms =
      let ms' = infer_msf_c ~withcheck fenv tbl c ms in
      if Sv.subset ms' ms then (Hashtbl.add tbl i.i_loc ms; ms)
      else loop (Sv.union ms' ms) in
    loop ms

  | Cwhile (_, c1, _, _, c2) ->
    (* c1; while e do c2; c1 *)
    let rec loop ms =
      let ms1 = infer_msf_c ~withcheck fenv tbl c1 ms in
      let ms2 = infer_msf_c ~withcheck fenv tbl c2 ms1 in
      if Sv.subset ms2 ms then (Hashtbl.add tbl i.i_loc ms1; ms1)
      else loop (Sv.union ms2 ms) in
    loop ms

  | Cassgn(Lvar x, tag, _, e) when Sv.mem (L.unloc x) ms ->
      (* We need to allow assignments to MSF if the compiler introduces them,
         to be able to use the checker after inlining. *)
      let gets_removed =
        match tag with
        | AT_none | AT_keep -> false
        | AT_rename | AT_inline | AT_phinode -> true
      in
      begin match reg_expr_opt ~direct:true e with
      | Some x' when gets_removed ->
          Sv.add (L.unloc x') (Sv.remove (L.unloc x) ms)
      | _ -> error ~loc "assignment to MSF variable %a not allowed" pp_var_i x
      end

  | Cassgn _ ->
    ms

  | Ccall(xs, f, es) ->
    let fty = FEnv.get_fty fenv f in
    let ms =
      let doout ms vfty x =
        match vfty with
        | IsMsf -> let x = reg_lval ~direct:true loc x in Sv.remove (L.unloc x) ms
        | _ -> ms in
      let ms = List.fold_left2 doout ms fty.tyout xs in
      check_call ~loc f fty.modmsf ms;
      ms
    in
    let doin ms vfty e =
      match vfty with
      | IsMsf -> let x = reg_expr ~direct:true loc e in Sv.add (L.unloc x) ms
      | _ -> ms in
    List.fold_left2 doin ms fty.tyin es

  | Copn (xs, _, o, es) ->
    match is_special o, xs, es with
    | Init_msf, [x], _ ->
      let x = Option.map_default (fun x -> Sv.singleton (L.unloc x)) Sv.empty (reg_lval_opt ~direct:true loc x) in
      if not (Sv.subset ms x) then
        error ~loc "only %a is ensured to be msf after init_msf, %a are required" pp_vset x pp_vset ms;
      Sv.empty

    | Init_msf, _, _ -> assert false

    | Update_msf, [xms], [_e; ms'] ->
      let xms = reg_lval ~direct:true loc xms and ms' = reg_expr ~direct:true loc ms' in
      if not (Sv.subset ms (Sv.singleton (L.unloc xms))) then
        error ~loc "only %a is ensured to be msf after update_msf, %a are required" pp_var_i xms pp_vset ms;
      Sv.singleton (L.unloc ms')

    | Update_msf, _, _ -> assert false

    | Mov_msf, [x], [x'] ->
      let x = reg_lval ~direct:true loc x and x' = reg_expr ~direct:true loc x' in
      Sv.add (L.unloc x') (Sv.remove (L.unloc x) ms)

    | Mov_msf, _, _ -> assert false

    | Protect, xs, [_; ms'] ->
      checks ms xs;
      let ms' = reg_expr ~direct:true loc ms' in
      Sv.add (L.unloc ms') ms

    | Protect, _, _ -> assert false

    | Other, xs, _ -> checks ms xs; ms

and infer_msf_c ~withcheck fenv tbl c ms =
  List.fold_right (infer_msf_i ~withcheck fenv tbl) c ms


(* --------------------------------------------------------- *)
(* Typing environment                                        *)

module Env : sig
  type env  (* constraints holder *)
  type venv (* type variables association *)

  val init : unit -> env

  val empty : env -> venv
  val constraints : env -> C.constraints
  val add_var : env -> venv -> var -> vty -> venv
  val public    : env -> Lvl.t
  val secret    : env -> Lvl.t

  val public2    : env -> VlPairs.t
  val transient  : env -> VlPairs.t
  val secret2    : env -> VlPairs.t

  val dpublic    : env -> vty
  val dsecret    : env -> vty

  val get   : venv -> var -> vty
  val get_i : venv -> var_i -> vty
  val gget  : venv -> int ggvar -> vty

  val fresh  : ?name:string -> env -> Lvl.t
  val fresh2 : ?name:string -> env -> VlPairs.t

  val init_ty : env -> venv -> var -> vty -> venv
  val set_ty : env -> venv -> var_i -> vty -> venv
  val set_init_msf : env -> venv -> var_i option -> venv

  val max : env -> venv -> venv -> venv

  val get_msf_oracle : env -> (L.i_loc, Sv.t) Hashtbl.t
  val msf_oracle : env -> L.i_loc -> Sv.t

  val freshen : ?min:Constraints.VlPairs.t -> env -> venv -> venv
  val ensure_le : L.t -> venv -> venv -> unit
  val clone_for_call : env -> ty_fun -> vfty list * vfty list * VlPairs.t
          (* output type, input type, output corruption *)

  val corruption : env -> venv -> VlPairs.t -> venv
  val corruption_speculative : env -> venv -> VlPairs.t -> venv

  val get_resulting_corruption : venv -> VlPairs.t

  val venv_forget : Sv.t -> venv -> venv
  (** Restrict knowledge to the given set of variables. *)

end = struct

  type env = {
      constraints : C.constraints;
      msf_oracle  : (L.i_loc, Sv.t) Hashtbl.t;
    }

  type venv = {
      vtype : vty Mv.t;
      vars : Sv.t;
      resulting_corruption : VlPairs.t;
      public2 : VlPairs.t (* needed in 'get'. This is dirty.
        Ideally, env and venv should be merged, but this requires
        some work and restructuring of the module signature. *)
    }

  let init () =
    { constraints = C.init ();
      msf_oracle  = Hashtbl.create 97; }

  let constraints env = env.constraints

  let public    env = C.public    env.constraints
  let secret    env = C.secret    env.constraints

  let public2   env = (public env, public env)
  let transient env = (public env, secret env)
  let secret2   env = (secret env, secret env)

  let dpublic   env = Direct (public2 env)
  let dsecret   env = Direct (secret2 env)

  let fresh ?name env = C.fresh ?name env.constraints

  let fresh2 ?name env = (C.fresh ?name env.constraints,
                          C.fresh ?name env.constraints)

  let empty env =
    { vtype = Mv.empty;
      vars = Sv.empty;
      resulting_corruption = fresh2 env;
      public2 = public2 env;
    }

  let get venv x =
    try match x.v_kind with
    | Global -> Direct venv.public2
    | _ -> Mv.find x venv.vtype
    with Not_found -> assert false

  let get_i venv x = get venv (L.unloc x)
  let gget venv x = get_i venv x.gv

  let add_le_var ty1 ty2 =
    match ty1, ty2 with
    | Direct ty1, Direct ty2 -> VlPairs.add_le ty1 ty2
    | Indirect (l1, x1), Indirect (l2, x2) -> VlPairs.add_le l1 l2; VlPairs.add_le x1 x2
    | _ -> assert false

  let init_ty env venv x xty =
      (* warning Always (L.i_loc0 x.v_dloc) "initialising variable %a" pp_var x; *)
    let nxty =
      let public2 = public2 env in
      match x.v_kind, xty with
      | Const, Direct le -> VlPairs.add_le le public2; xty
      | Inline, Direct le -> VlPairs.add_le le public2; xty

      | Global, Direct _
      | Stack Direct, Direct _
      | Stack (Pointer _), Indirect _
      | Reg (_, Direct), Direct _
      | Reg (_, Pointer _), Indirect _ -> xty

      | _, _ -> assert false
    in
    { venv with
        vtype = Mv.add x nxty venv.vtype;
        vars = Sv.add x venv.vars
    }

  let add_var env venv x vty =
    assert (not (Mv.mem x venv.vtype));
    try init_ty env venv x vty
    with Lvl.Unsat _unsat ->
      error ~loc:(x.v_dloc)
        "invalid security annotations for %a, this leads to invalid constraints"
           pp_var x

  let set_ty env venv x nxty =
    init_ty env venv (L.unloc x) nxty

  (* Initialises type for InitMsf operations. Acts as a Fence operator *)
  let set_init_msf env venv ms =
    let venv =
      match ms with
      | Some ms -> set_ty env venv ms (dpublic env)
      | None -> venv in
    let operate_fence _ = function
        | Direct le -> Direct (VlPairs.normalise le)
        | Indirect(lp, le) -> Indirect(VlPairs.normalise lp, VlPairs.normalise le)
    in
    { venv with vtype = Mv.mapi operate_fence venv.vtype }


  let max env venv1 venv2 =
    let merge1 l1 l2 =
      if Lvl.equal l1 l2 then l1
      else
        let l = fresh env in
        Lvl.add_le l1 l; Lvl.add_le l2 l;
        l in

    let merge (n1, s1) (n2, s2) =
      (merge1 n1 n2, merge1 s1 s2) in

    let merge_var _ oty1 oty2 =
      match oget oty1, oget oty2 with
      | Direct le1, Direct le2 -> Some (Direct (merge le1 le2))
      | Indirect(lp1, le1), Indirect(lp2, le2) -> Some (Indirect (merge lp1 lp2, merge le1 le2))
      | _ -> assert false
    in
    { venv1 with
      vtype = Mv.merge merge_var venv1.vtype venv2.vtype;
      resulting_corruption = merge venv1.resulting_corruption venv2.resulting_corruption
    }

  let get_msf_oracle env = env.msf_oracle

  let msf_oracle env loc =
    try Hashtbl.find env.msf_oracle loc with Not_found -> assert false

  (* freshen all variables in environment env, venv, with
     possibly a minimum (typically memory corruption) *)
  let freshen ?min env venv =
    let fresh ~in_memory le =
      let l = fresh2 env in
      if in_memory && min != None then VlPairs.add_le (oget min) l;
      VlPairs.add_le le l;
      l
    in
    { venv with vtype = Sv.fold (fun x vtype ->
        let in_memory = match x.v_kind with
          | Wsize.Global (* likely unused as global variables are not in venv.vars *)
          | Stack _ -> true
          | Const | Inline | Reg _ -> false
        in
        let ty =
          match Mv.find x vtype with
          | Direct le -> Direct (fresh ~in_memory le)
          | Indirect(lp, le) ->
             Indirect(fresh ~in_memory lp, fresh ~in_memory:true le) (* the pointed values are in memory *)
        in
        Mv.add x ty vtype) venv.vars venv.vtype }

  let ensure_le loc venv1 venv2 =
    let add_le_silent _ oty1 oty2 = add_le_var (oget oty1) (oget oty2); None in
    try ignore (Mv.merge add_le_silent venv1.vtype venv2.vtype)
    with Lvl.Unsat _unsat ->
      error ~loc "constraints caused by the loop cannot be satisfied"

  let clone_for_call (env:env) (tyfun:ty_fun) =
    let subst1 = C.clone tyfun.constraints env.constraints in
    let subst (n, s) = (subst1 n, subst1 s) in
    let subst_ty = function
      | IsMsf -> IsMsf
      | IsNormal ty -> let ty =
            match ty with
            | Direct le -> Direct (subst le)
            | Indirect(lp, le) -> Indirect(subst lp, subst le) in
          IsNormal ty in
    List.map subst_ty tyfun.tyout, List.map subst_ty tyfun.tyin,
    subst tyfun.resulting_corruption

  let corruption env venv ty =
    VlPairs.add_le ty venv.resulting_corruption; (* update corruption level *)
    freshen ~min:ty env venv

  let corruption_speculative env venv (_, s) = corruption env venv (public env, s)

  let get_resulting_corruption venv = venv.resulting_corruption

  let venv_forget lv x =
    { x with
      vtype = Mv.filter (fun k _ -> Sv.mem k lv) x.vtype;
      vars = Sv.inter x.vars lv }

end


(* --------------------------------------------------------- *)
let error_unsat loc (_ : Lvl.t list * Lvl.t * Lvl.t) pp e ety ety' =
  error ~loc
    "%a has type %a but should be at most %a"
    pp e pp_vty ety pp_vty ety'

let ssafe_test x aa ws i =
  let x = L.unloc x in
  match x.v_kind, x.v_ty, i with
  | Reg (_, Direct), _, _ -> true
  | _, Arr (ws1, len), Pconst v ->
      let len = Z.of_int (arr_size ws1 len) in
      let v = Z.of_int (access_offset aa ws (Z.to_int v)) in
      let v_max = Z.add v (Z.of_int (size_of_ws ws - 1)) in
      Z.(leq zero v && lt v_max len)
  | _ -> false

let content_ty = function
  | Direct le | Indirect (_, le) -> le

(* --------------------------------------------------------- *)
(* Type checking of expressions                              *)

let rec ty_expr env venv loc (e:expr) : vty =
  match e with
  | Pconst _ | Pbool _ | Parr_init _ -> Env.dpublic env

  | Pvar x -> Env.gget venv x

  | Pget (_, aa, ws, x, i) ->
      ensure_public_address env venv loc x.gv;
      ensure_public env venv loc i;
      let ty = Env.fresh2 env
      and xty = Env.gget venv x in
      if not (ssafe_test x.gv aa ws i) then VlPairs.add_le_speculative (Env.secret env) ty;
      VlPairs.add_le (content_ty xty) ty;
      Direct ty

    (* in the case of sub-arrays, no operation is performed, and there is now
       an alias on the values. Thus the type must be *equal* *)
  | Psub (_, _, _, x, i) ->
      ensure_public env venv loc i;
      Env.gget venv x

  | Pload (_, _, i) ->
      ensure_public env venv loc i;
      Env.dsecret env

  | Papp1(o, e)      ->
    let public = not (CT.is_ct_op1 o) in
    ty_exprs_max ~public env venv loc [e]
  | Papp2(o, e1, e2) ->
    let public = not (CT.is_ct_op2 o) in
    ty_exprs_max ~public env venv loc [e1; e2]
  | PappN(o, es)     ->
    let public = not (CT.is_ct_opN o) in
    ty_exprs_max ~public env venv loc es

  | Pif(_, e1, e2, e3) ->
      let ty1 = ty_expr env venv loc e1 in
      let ty2 = ty_expr env venv loc e2 in
      let ty3 = ty_expr env venv loc e3 in
      match ty1 with
      | Indirect _ -> assert false
      | Direct l1 ->
        let do_indirect lp2 le2 lp3 le3 =
          let lp = Env.fresh2 env in
          let le = Env.fresh2 env in
          (* The condition expression is also added to the constraints because it can be deduced from the result value*)
          VlPairs.add_le l1 lp; VlPairs.add_le lp2 lp; VlPairs.add_le lp3 lp;
          VlPairs.add_le l1 le; VlPairs.add_le le2 le; VlPairs.add_le le3 le;
          Indirect(lp, le) in
        match ty2, ty3 with
        | Direct l2, Direct l3 ->
          let le = Env.fresh2 env in
          VlPairs.add_le l1 le; VlPairs.add_le l2 le; VlPairs.add_le l3 le;
          Direct le
        | Indirect(lp2, le2), Indirect(lp3, le3) ->
          do_indirect lp2 le2 lp3 le3
        | Indirect(lp2, le2), Direct le3 ->
          do_indirect lp2 le2 (Env.public2 env) le3
        | Direct le2, Indirect (lp3, le3) ->
          do_indirect (Env.public2 env) le2 lp3 le3

and ensure_smaller env venv loc e l =
  let ety = ty_expr env venv loc e in
  match ety with
  | Direct le | Indirect (le, _) ->
    try VlPairs.add_le le l
    with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety (Direct l)

and ensure_public env venv loc e = ensure_smaller env venv loc e (Env.public2 env)

and ensure_public_address env venv loc x =
  let ety = Env.get_i venv x in
  match ety with
  | Direct _ -> () (* stack or reg arrays have public addresses by definition *)
  | Indirect (le, _) -> try VlPairs.add_le le (Env.public2 env)
      with Lvl.Unsat unsat -> error_unsat loc unsat pp_var_i x ety (Direct (Env.public2 env))

and ty_exprs_max ~(public:bool) env venv loc es : vty =
  let l = if public then Env.public2 env else Env.fresh2 env in
  List.iter (fun e -> ensure_smaller env venv loc e l) es;
  Direct l


(* ------------------------------------------------------------- *)
(* Compare expressions up to constant folding *)
let expr_equal a b =
  let fcp =
    let open Glob_options in
    match !target_arch with
    | X86_64 -> X86_decl.x86_fcp
    | ARM_M4 -> Arm_decl.arm_fcp 
    | RISCV  -> Riscv_decl.riscv_fcp
  in 
  let normalize e =
    e |> Conv.cexpr_of_expr |> Constant_prop.(const_prop_e fcp None empty_cpm) in
  Expr.eq_expr (normalize a) (normalize b)

(* ------------------------------------------------------------- *)
(* MSF status                                                    *)

module MSF : sig

  type t
  val toinit : t
  val exact1 : var_i -> t
  val add : var_i -> t -> t
  val update   : t -> var  -> t

  val enter_if : t -> expr -> t
  val max : t -> t -> t

  val check_msf : t -> var_i -> unit
  val check_msf_trans : t -> var_i -> expr -> unit
  val is_msf    : t -> var_i -> bool
  val is_msf_exact    : t -> var_i -> bool
  val check_msf_exact : t -> var_i -> unit
  val loop : Env.env -> L.i_loc -> t -> t
  val end_loop : L.t -> t -> t -> t

  val pp : Format.formatter -> t -> unit

  end = struct

  type t = Sv.t * expr option

  let toinit = (Sv.empty, None)
  let exact xs = (xs, None)
  let trans xs e = (xs, Some e)

  let exact1 x =
    ensure_register ~direct:true x;
    exact (Sv.singleton (L.unloc x))

  let add x (xs, oe) =
    ensure_register ~direct:true x;
    let loc = L.loc x in
    let x = L.unloc x in
    Stdlib.Option.iter (fun e ->
        if Sv.mem x (vars_e e) then
          error ~loc "%a cannot become an MSF as the current status depends on it (%a)" pp_var x pp_expr e
      ) oe;
    (Sv.add x xs, oe)

  let update (xs, oe) x =
    match oe with
    | Some e when Sv.mem x (vars_e e) -> toinit
    | _ -> (Sv.remove x xs, oe)

  let enter_if msf e =
    match msf  with
    | (_, Some _) -> toinit
    | (xs, None) -> (xs, Some e)

  let max (xs1, oe1) (xs2, oe2) =
    match oe1, oe2 with
    | None, None -> Sv.inter xs1 xs2, None
    | Some e1, Some e2 when expr_equal e1 e2 -> Sv.inter xs1 xs2, Some e1
    | _, _ -> toinit

  let check_msf (xs, _) ms =
    if not (Sv.mem (L.unloc ms) xs) then
      error ~loc:(L.loc ms)
        "the variable %a is not known to be a msf, only %a are"
        pp_var_i ms pp_vset xs

  let check_msf_trans ((_, ob) as msf) ms b =
    match ob with
    | None -> error ~loc:(L.loc ms) "MSF is not Trans"
    | Some b' ->
        if not (expr_equal b b') then
          error ~loc:(L.loc ms)
          "the expression %a need to be equal to@ %a"
            pp_expr b pp_expr  b';
        check_msf msf ms

  let is_msf (xs, _) ms = Sv.mem (L.unloc ms) xs

  let is_msf_exact (xs, ob) ms =
    match ob with
    | Some _ -> false
    | None -> Sv.mem (L.unloc ms) xs

  let check_msf_exact ((_, ob) as msf) ms =
    match ob with
    | Some b ->
      error ~loc:(L.loc ms) "MSF is Trans@ %a"  pp_expr b
    | None -> check_msf msf ms

  let pp fmt (xs, oe) =
    match oe with
    | Some e -> Format.fprintf fmt "Trans %a %a" pp_vset xs pp_expr e
    | None   -> Format.fprintf fmt "Exact %a" pp_vset xs

  let loop env loc ((xs, oe) as msf) =
    let xs' = Env.msf_oracle env loc in
    if Sv.subset xs' xs then (xs', oe)
    else
      error ~loc:(loc.L.base_loc)
        "current msf = %a, it should contain %a"
        pp msf pp_vset (Sv.diff xs' xs)

  let end_loop loc ((xsi, oei) as msfi) ((xso, oeo) as msfo)=
    if Sv.subset xsi xso then
      begin match oei, oeo with
      | None, None -> msfi
      | Some ei, Some eo when expr_equal ei eo -> msfi
      | _, _ ->
          if not (Sv.is_empty xsi) then
            error ~loc
            "msf is %a it should be be at least %a"
            pp msfo pp msfi;
          toinit
    end
    else
      error ~loc
        "msf is %a it should be be at least %a"
        pp msfo pp msfi

end


(* --------------------------------------------------------- *)
(* Type checking of lvalue                                   *)
type msf_e = MSF.t * Env.venv

let ty_lval env ((msf, venv) as msf_e : msf_e) x ety : msf_e =
  (* First path the type ety to make it consistant with the variable info *)
  match x with
  | Lnone _ -> msf_e
  | Lvar x ->
      (* TODO assumption: p = e when p is a pointer and e a direct value means p
         points to a new position, where the expression is *)
      (* as opposed to assigning the pointer directly to the given value *)
      (* likewise, assuming x = p means storing the value pointed by p in x *)
      (* what with p = [ x ]? I believe it does not compile *)
      let lp, le =
        match ety with
        | Direct le -> Env.public2 env, le
        | Indirect(lp, le) -> lp, le in
      let xty = if is_ptr (kind_i x)
          then Indirect(lp, le)
          else Direct le in
      let msf = MSF.update msf (L.unloc x) in
      let venv = Env.set_ty env venv x xty in
      msf, venv

  | Lmem(_, _, vi, i) ->
      ensure_public env venv vi i;
        (* programmes are assumed to be safe, thus corruption from memory store
           with [x + i] is speculative only *)
      msf, Env.corruption_speculative env venv (content_ty ety)

  | Laset(_, aa, ws, x, i) ->
      ensure_public_address env venv (L.loc x) x;
      ensure_public env venv (L.loc x) i;
      let le = content_ty ety in
      let venv =
        let l = Env.fresh2 env in
        let xty =
          match Env.get_i venv x with
          | Direct lx -> VlPairs.add_le lx l; VlPairs.add_le le l; Direct l
          | Indirect (lp, lx) -> VlPairs.add_le lx l; VlPairs.add_le le l;
              Indirect (lp, l)
        in
        Env.set_ty env venv x xty in

      if ssafe_test x aa ws i then
        msf, venv
      else (* mispeculation has necessarily occured *)
        msf, Env.corruption_speculative env venv le

  | Lasub(_, _, _, x, i) ->
      ensure_public_address env venv (L.loc x) x;
      ensure_public env venv (L.loc x) i;
      let le = content_ty ety in
      let l = Env.fresh2 env in
      let xty =
        match Env.get_i venv x with
        | Direct lx -> VlPairs.add_le lx l; VlPairs.add_le le l; Direct l
        | Indirect (lp, lx) -> VlPairs.add_le lx l; VlPairs.add_le le l;
            Indirect (lp, l)
      in
      msf, Env.set_ty env venv x xty

let ty_lvals1 env (msf_e : msf_e) xs ety : msf_e =
  List.fold_left (fun msf_e x -> ty_lval env msf_e x ety) msf_e xs

let ty_lvals env (msf_e : msf_e) xs tys : msf_e =
  List.fold_left2 (ty_lval env) msf_e xs tys


(* -------------------------------------------------------------- *)
(* declassify                                                     *)
(* TODO ensure declassify cannot occur on potentially corrupted   *)
(* stack values                                                   *)

let sdeclassify = "declassify"

let is_declassify annot =
  Annot.ensure_uniq1 sdeclassify Annot.none annot <> None

let declassify_lvl env (_, s) = (Env.public env, s)

let declassify env = function
  | Direct le          -> Direct (declassify_lvl env le)
  | Indirect (lp, le)  -> Indirect (lp, declassify_lvl env le)

let declassify_ty env annot ty = if is_declassify annot
  then declassify env ty
  else ty

let declassify_tys env annot tys = if is_declassify annot
  then List.map (declassify env) tys
  else tys


(* right now only used by syscall, which only consists of randombytes
   it is thus tailored for this specific function. *)
let ensure_public_address_expr env venv loc e =
  let ety = ty_expr env venv loc e in
  match ety with
  | Direct _ -> ()
  | Indirect (le, _) -> try VlPairs.add_le le (Env.public2 env)
      with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety (Direct (Env.public2 env))

(* --------------------------------------------------------------- *)
let move_msf ~loc env (msf, venv) mso msi =
  let mso = reg_lval ~direct:true loc mso
  and msi = reg_expr ~direct:true loc msi in
  MSF.check_msf msf msi;
  MSF.add mso msf, Env.set_ty env venv mso (Env.dpublic env)

(* --------------------------------------------------------------- *)
(* [ty_instr env msf i] return msf' such that env, msf |- i : msf' *)

let rec ty_instr is_ct_asm fenv env (msf, venv) i =
  let venv = Env.venv_forget (fst i.i_info) venv in
  let msf, venv = ty_instr_r is_ct_asm fenv env (msf, venv) i in
  msf, Env.venv_forget (snd i.i_info) venv

and ty_instr_r is_ct_asm fenv env ((msf,venv) as msf_e :msf_e) i =
  let loc = i.i_loc.L.base_loc in
  match i.i_desc with
  | Csyscall (xs, o, es) ->
    (* TODO: generalize to other syscalls *)
    assert (match o with Syscall_t.RandomBytes _ -> true);
    List.iter (ensure_public_address_expr env venv loc) es;
    (* We don't known what happen to MSF after external function call *)
    ty_lvals1 env (MSF.toinit, venv) xs (Env.dsecret env)

  | Cassgn(mso, _, _, (Pvar x as msi)) when MSF.is_msf msf x.gv ->
    move_msf ~loc env msf_e mso msi

  | Cassgn(x, _, _, e) ->
    let ety = ty_expr env venv loc e in
    ty_lval env msf_e x (declassify_ty env i.i_annot ety)

  | Copn(xs, _, o, es) ->
    begin match is_special o, xs, es with
    | Init_msf, [ms], _ ->
      let ms = reg_lval_opt ~direct:true loc ms in
      let venv = Env.set_init_msf env venv ms in
      let ms = Option.map_default MSF.exact1 MSF.toinit ms in
      ms, venv

    | Init_msf, _, _ -> assert false

    | Update_msf, [mso], [b; msi] ->
      let mso = reg_lval ~direct:true loc mso and msi = reg_expr ~direct:true loc msi in
      (* do not check b, if check_msf_trans succeed then b is public *)
      MSF.check_msf_trans msf msi b;
      let _, venv = ty_lvals1 env (msf, venv) xs (Env.dpublic env) in
      MSF.exact1 mso, venv

    | Update_msf, _, _ -> assert false

    | Mov_msf, [mso], [msi] ->
      move_msf ~loc env msf_e mso msi

    | Mov_msf, _, _ -> assert false

    | Protect, [x], [e; ms] ->
      let _ = reg_lval ~direct:false loc x and _ = reg_expr ~direct:false loc e and
          ms = reg_expr ~direct:true loc ms in
      MSF.check_msf_exact msf ms;
      let xty =
        match ty_expr env venv loc e with
        | Direct (n, _) -> Direct (n, n)
        | Indirect ((n, _), le) -> Indirect ((n, n), le) in

      ty_lval env msf_e x xty

    | Protect, _, _ -> assert false

    | Other, _, _  ->
      let public = not (CT.is_ct_sopn is_ct_asm o) in
      let ety = ty_exprs_max ~public env venv loc es in
      ty_lvals1 env msf_e xs (declassify_ty env i.i_annot ety)
    end

  | Cif(e, c1, c2) ->
    let msf1, msf2 =
      if is_inline i then
        msf, msf
      else begin
          ensure_public env venv loc e;
          MSF.enter_if msf e, MSF.enter_if msf (Papp1(Onot, e))
        end
    in
    let msf1, venv1 = ty_cmd is_ct_asm fenv env (msf1, venv) c1 in
    let msf2, venv2 = ty_cmd is_ct_asm fenv env (msf2, venv) c2 in
    let venv1 = Env.venv_forget (snd i.i_info) venv1 in
    let venv2 = Env.venv_forget (snd i.i_info) venv2 in
    MSF.max msf1 msf2, Env.max env venv1 venv2

  | Cfor(x, (_, e1, e2), c) ->
      ensure_public env venv loc e1;
      ensure_public env venv loc e2;

      (* Live set after the loop guard *)
      let live_at_c =
        let live_after_i = snd i.i_info in
        match c with i :: _ -> Sv.union live_after_i (fst i.i_info) | [] -> live_after_i in

      let msf = MSF.loop env i.i_loc msf in
      (* let w, _ = written_vars [i] in *)
      let venv1 = Env.freshen env venv in (* venv <= venv1 *)
      let msf_e = ty_lval env (msf, venv1) (Lvar x) (Env.dpublic env) in
      let (msf', venv') = ty_cmd is_ct_asm fenv env msf_e c in
      let msf' = MSF.end_loop loc msf msf' in
      let venv1 = Env.venv_forget live_at_c venv1 in
      let venv' = Env.venv_forget live_at_c venv' in
      Env.ensure_le loc venv' venv1; (* venv' <= venv1 *)
      msf', venv1

  | Cwhile(_, c1, e, (_, (_, live_at_c2)), c2) ->
    (* c1; while e do (c2; c1) *)
    (* env, msf <= env1, msf1
       env1, msf1 |- c1 : msf2, env2   env2 |- e : public
       env2, enter_if e msf2 |- c2 : env1, msf1
       --------------------------------------------------------------------------------
       env, msf |- while c1 e c2 : enter_if e msf1
     *)
    let msf1 = MSF.loop env i.i_loc msf in
    (* let w, _ = written_vars [i] in *)
    (* NOTE cannot restrict refreshed variables to local modified vars
       because of memory corruption: if loop body corrupts some stack variable
       constrained to public, the test fails, while marking all stack variables
       as secret is sufficient *)

    let venv1 = Env.freshen env venv in (* venv <= venv1 *)
    let (msf2, venv2) = ty_cmd is_ct_asm fenv env (msf1, venv1) c1 in
    ensure_public env venv2 loc e;
    let (msf', venv') = ty_cmd is_ct_asm fenv env (MSF.enter_if msf2 e, Env.venv_forget live_at_c2 venv2) c2 in
    let _ = MSF.end_loop loc msf1 msf' in
    Env.ensure_le loc venv' venv1; (* venv' <= venv1 *)
    MSF.enter_if msf2 (Papp1(Onot, e)), venv2

  | Ccall (xs, f, es) ->
    let fty = FEnv.get_fty fenv f in
    let modmsf = fty.modmsf in
    let tyout, tyin, resulting_corruption = Env.clone_for_call env fty in

    let input_ty e vfty =
      match vfty with
      | IsMsf ->
        (* we don't check that e is public, it is ensured by being msf *)
        let ms = reg_expr ~direct:true loc e in
        MSF.check_msf_exact msf ms
      | IsNormal ety' ->
        let ety = ty_expr env venv loc e in
        try match ety, ety' with
        | Direct le, Direct le' -> VlPairs.add_le le le'
        | Direct le, Indirect (_, le') -> VlPairs.add_le le le'
        | Indirect(lp, le), Direct le' -> VlPairs.add_le lp (Env.public2 env); VlPairs.add_le le le'
        | Indirect(lp, le), Indirect(lp', le') -> VlPairs.add_le lp lp'; VlPairs.add_le le le'
        with Lvl.Unsat unsat -> error_unsat loc unsat pp_expr e ety ety' in
    List.iter2 input_ty es tyin;

    (* callee function has its own effect on this function corruption *)
    let venv = Env.corruption env venv resulting_corruption in

    (* compute the resulting venv *)
    let output_ty msf_e x vfty =
      let ty =
        match vfty with
        | IsMsf -> Env.dpublic env
        | IsNormal ty -> declassify_ty env i.i_annot ty in
      let (msf, venv) = ty_lval env msf_e x ty in
      let msf = if vfty = IsMsf then MSF.add (reg_lval ~direct:true loc x) msf else msf in
      (msf, venv) in
    let msf = if is_Modified modmsf then MSF.toinit else msf in
    List.fold_left2 output_ty (msf, venv) xs tyout

and ty_cmd is_ct_asm fenv env msf_e c =
  List.fold_left (ty_instr is_ct_asm fenv env) msf_e c


(* ------------------------------------------------------------------- *)
(* Do the inference + type checking of function                        *)
(*
#nomodmsf #constraints = "l1 <= transient, l2 <= l1"
   fn f (#public #secret reg u64[1] x, #poly = l1 stack u8 c) ->
        #poly=l1 #poly=l2 u64[1]
*)

let parse_var_annot ~(msf:bool) (annot: annotations) : ulevel list =
  let module A = Annot in
  let filters =
    [spublic, (fun a -> A.none a; Public);
     ssecret, (fun a -> A.none a; Secret);
     stransient, (fun a -> A.none a; Transient)
     ] in

  let filters =
    if msf then (smsf, (fun a -> A.none a; Msf)) :: filters else filters in

  let lvls = A.process_annot filters annot in

  List.map snd lvls

exception Error_after of string * string

let parse_constraints (s:string) : (string * string) list =
  let error expected s n =
    raise (Error_after (expected, String.sub s 0 n)) in

  let is_blank c = c = ' ' || c = '\n' || c = '\t' in

  let rec check_blank expected s first last =
    if first >= last then ()
    else if is_blank s.[first] then
      check_blank expected s (first + 1) last
    else error expected s first in

  let search expected re s n =
    let first =
      try Str.search_forward re s n
      with Not_found -> error expected s n in
    let sub = Str.matched_string s in
    let next = Str.match_end () in
    check_blank expected s n first;
    sub, next in

  let rec skip_blank s n =
    if n < String.length s && is_blank s.[n] then skip_blank s (n+1) else n in

  let ident s n = search "ident" (Str.regexp "[A-Za-z0-9]+") s n in

  let is_le s n = snd (search "<=" (Str.regexp "<=") s n) in

  let is_comma s n = snd (search "," (Str.regexp ",") s n) in

  let parse_c s n =
    let l1, n = ident s n in
    let n = is_le s n in
    let l2, n = ident s n in
    (l1, l2), n in

  let rec parse_rec acc s n =
    let n = skip_blank s n in
    if String.length s <= n then acc
    else
      let n = is_comma s n in
      let c, n = parse_c s n in
      parse_rec (c::acc) s n in

  let n = skip_blank s 0 in
  if String.length s <= n then []
  else
    let c, n = parse_c s n in
    parse_rec [c] s n

let parse_user_constraints (a:annotations) : (string * string) list =
  let module A = Annot in
  let sconstraints = "constraints" in
  let on_string loc _ s =
    try parse_constraints s
    with Error_after(kw, s) ->
      A.error ~loc
        "%s expected after %s" kw s in
  let error loc =
    A.error ~loc
      "attribute for %s should be a string" sconstraints in

   List.flatten
     (List.map snd (A.process_annot [sconstraints, A.on_attribute ~on_string error] a))

let init_constraint fenv f =
  let sig_annot = SecurityAnnotations.get_sct_signature f.f_annot.f_user_annot in
  let env = Env.init () in
  let venv = Env.empty env in
  let tbl = Hashtbl.create 97 in
  Hashtbl.add tbl spublic    (Env.public2 env);
  Hashtbl.add tbl stransient (Env.transient env);
  Hashtbl.add tbl ssecret    (Env.secret2 env);
  (* export function: all input type should be at most transient and msf is not allowed *)
  (* the new version does not take that into accout, does it? *)
  let export = FInfo.is_export f.f_cc in

  let add_lvl s =
    try Hashtbl.find tbl s
    with Not_found ->
      let l = Env.fresh2 ~name:s env in
      Hashtbl.add tbl s l;
      l in

  let to_lvl = function
    | Secret -> Env.secret2 env
    | Transient -> Env.transient env
    | Public | Msf -> Env.public2 env in

  let lvl_of_sa_level =
    let open SecurityAnnotations in
    let lvl_of_simple_level get =
      function
      | Public -> Env.public env
      | Secret -> Env.secret env
      | Named s -> get (add_lvl s)
    in
    function { normal; speculative } ->
      lvl_of_simple_level fst normal, lvl_of_simple_level snd speculative
  in

  let to_vty =
    function
    | SecurityAnnotations.Msf -> Direct (to_lvl Msf)
    | Direct n -> Direct (lvl_of_sa_level n)
    | Indirect { ptr; value } -> Indirect (lvl_of_sa_level ptr, lvl_of_sa_level value)
  in

  let error_msf loc =
    error ~loc
      "%s annotation not allowed here" smsf in

  (** The [is_local] argument is true when variable [x] is a local variable as
  opposed to an argument or a returned value which inherits constraints from the call-sites. *)
  let mk_vty loc ~is_local ~(msf:bool) x ls an =
    let msf, ovty =
      match ls, an with
      | [], None -> None, None
      | [l], None ->
        if not msf && l = Msf then error_msf loc;
        Some (l = Msf), Some(Direct (to_lvl l))
      | [l1; l2], None ->
        if (l1 = Msf || l2 = Msf) then error_msf loc;
        Some false, Some(Indirect (to_lvl l1, to_lvl l2))
      | _, None ->
        error ~loc:(x.v_dloc)
          "invalid security annotations %a" pp_var x
      | [Msf], Some n ->
         if not msf then error_msf loc;
         Some true, Some (to_vty n)
      | _, Some n ->
         Some (n = SecurityAnnotations.Msf), Some (to_vty n)
    in
    let vty =
      match ovty with
      | None ->
        begin match x.v_kind with
        | Const -> Env.dpublic env
        | (Stack Direct | Reg (_, Direct)) when is_local -> Direct (Env.secret2 env)
        | (Stack Direct | Reg (_, Direct)) -> Direct (Env.fresh2 env)
        | (Stack (Pointer _) | Reg (_, Pointer _)) when is_local -> Indirect(Env.secret2 env, Env.fresh2 env)
        | (Stack (Pointer _) | Reg (_, Pointer _)) -> Indirect(Env.fresh2 env, Env.fresh2 env)
        | Inline -> Env.dpublic env
        | Global -> Env.dpublic env (* unsure *)
        end
      | Some ty -> (* this partly has the same role as Env.init_ty. Remove one occurence? *)
        begin match x.v_kind, ty with
        | Const, Direct _ -> ()
        | Stack Direct, Direct _ -> ()
        | Stack (Pointer _), Indirect _ -> ()
        | Reg (_, Direct), Direct _ -> ()
        | Reg (_, Pointer _), Indirect _ -> ()
        | Inline, Direct _ -> ()
        | Global, Direct _ -> ()
        | _ ->
          error ~loc
            "invalid security annotations for %a" pp_var x
        end; ty in
    msf, vty in

  let process_return i x annot =
    let loc = L.loc x and x = L.unloc x in
    let an = Option.bind sig_annot (SecurityAnnotations.get_nth_result i) in
    let ls = parse_var_annot ~msf:(not export) annot in
    mk_vty ~is_local:false loc ~msf:(not export) x ls an in

  (* process function outputs *)
  let tyout = List.map2i process_return f.f_ret f.f_ret_info.ret_annot in

  (* infer msf_oracle info *)
  let msfs =
    infer_msf_c ~withcheck:true fenv (Env.get_msf_oracle env) f.f_body
      (List.fold_left2 (fun s x (msf, _) ->
           if Option.default false msf then Sv.add (L.unloc x) s else s)
         Sv.empty f.f_ret tyout) in

  if export && not (Sv.is_empty msfs) then begin
    let vars_kind, pos =
      if Sv.subset msfs (Sv.of_list f.f_args)
      then "arguments", ", this is not allowed for export functions"
      else "variables", ""
    in
    error
      ~loc:f.f_loc
      "@[<h>the %s %a need to be MSFs%s.@]"
      vars_kind
      pp_vset msfs
      pos
  end;

  (* process function inputs *)
  let process_param i venv x =
    let an = Option.bind sig_annot (SecurityAnnotations.get_nth_argument i) in
    let ls = parse_var_annot ~msf:(not export) x.v_annot in
    let msf, vty = mk_vty ~is_local:false x.v_dloc ~msf:(not export) x ls an in
    let msf =
      match msf with
      | None -> Sv.mem x msfs
      | Some b ->
        if b <> Sv.mem x msfs then begin
          let loc = x.v_dloc in
          if b
          then warn ~loc:loc "%a does not need to be an MSF" pp_var x
          else error ~loc "%a should be an MSF" pp_var x
        end;
        b in
    if export then
      begin let lvls = match vty with
      | Indirect (p, v) -> [ p; v ]
      | Direct v -> [ v ]
      in List.iter begin fun l ->
          try VlPairs.add_le (Env.public env, Env.secret env) l
          with Lvl.Unsat _unsat ->
            error ~loc:(x.v_dloc)
              "security annotation for %a should be at least %s"
                 pp_var x stransient
        end lvls
      end;
    let venv = Env.add_var env venv x vty in
    let ty = if msf then IsMsf else IsNormal vty in
    venv, ty in

  let venv, tyin = List.mapi_fold process_param venv f.f_args in

  (* build the constraints *)
  let do_constraint (s1, s2) =
    let get s =
      try Hashtbl.find tbl s
      with Not_found ->
        error ~loc:f.f_loc "unbound security level %s" s in
    try VlPairs.add_le (get s1) (get s2)
    with Lvl.Unsat _unsat ->
      error ~loc:f.f_loc
        "cannot add constraint %s <= %s, it is inconsistent" s1 s2 in
  List.iter do_constraint (parse_user_constraints f.f_annot.f_user_annot);

  (* init type for local *)
  let do_local x venv =
    let _, vty = mk_vty x.v_dloc ~is_local:true ~msf:false x [] None in
    Env.add_var env venv x vty in

  let venv = Sv.fold do_local (locals f) venv in

  (* infer modmsf and check consistency with user info *)
  let modmsf = modmsf_c fenv f.f_body in
  let umodmsf =
    Annot.ensure_uniq
       ["modmsf", (fun a -> Annot.none a; true);
        "nomodmsf", (fun a -> Annot.none a; false)] f.f_annot.f_user_annot in
  begin match umodmsf with
  | None -> ()
  | Some annot ->
      if annot <> is_Modified modmsf then
        let sannot = if annot then "modmsf" else "nomodmsf" in
        error ~loc:f.f_loc "annotation %s should be %a" sannot pp_modmsf modmsf
  end;

  env, venv, tyin, tyout, modmsf


let rec ty_fun is_ct_asm fenv fn =
  try Hf.find fenv.env_ty fn
  with Not_found ->
    let fty = ty_fun_infer is_ct_asm fenv fn in
    Hf.add fenv.env_ty fn fty;
    fty

and ty_fun_infer is_ct_asm fenv fn =
  let f = FEnv.get_fun_def fenv fn in
  (* First compute all function call by f and recurse *)
  let _, called = written_vars_fc f in
  Mf.iter (fun fn _ -> ignore (ty_fun is_ct_asm fenv fn)) called;
  let env, venv, tyin, tyout, modmsf = init_constraint fenv f in
  (* init msf status *)
  let msf =
    List.fold_left2 (fun msf x ty ->
        if ty = IsMsf then MSF.add (L.mk_loc x.v_dloc x) msf else msf)
      MSF.toinit f.f_args tyin in
  (* start type checking of the body *)
  let msf, venv = ty_cmd is_ct_asm fenv env (msf, venv) f.f_body in
  (* build the resulting type *)
  let doout x (omsf, ty) =
    let le_ty ty1 ty2 =
      try
        match ty1, ty2 with
        | Direct le1, Direct le2 -> VlPairs.add_le le1 le2
        | Indirect(lp1, le1), Indirect(lp2, le2) ->
          VlPairs.add_le lp1 lp2; VlPairs.add_le le1 le2
        | _, _ -> assert false
      with Lvl.Unsat _unsat ->
        error ~loc:(L.loc x)
          "return type for %a is %a it should be less than %a"
             pp_var_i x pp_vty ty1 pp_vty ty2 in
    match omsf with
    | Some true  -> MSF.check_msf_exact msf x; IsMsf
    | Some false ->
      if MSF.is_msf_exact msf x then
        error ~loc:(L.loc x)
          "return annotation for %a should be %s" pp_var_i x smsf;
      le_ty (Env.get_i venv x) ty;
      IsNormal ty
    | None ->
      if MSF.is_msf_exact msf x then IsMsf
      else (le_ty (Env.get_i venv x) ty; IsNormal ty) in

  let tyout = List.map2 doout f.f_ret tyout in
  let resulting_corruption = Env.get_resulting_corruption venv in
  let (n1, s1) = resulting_corruption in

  let constraints = Env.constraints env in
  let add ls vty =
    match vty with
    | IsMsf -> ls
    | IsNormal (Direct (n, s)) -> n :: s :: ls
    | IsNormal (Indirect ((np, sp), (ne, se))) -> np :: sp :: ne :: se :: ls in
  let to_keep = List.fold_left add (List.fold_left add [n1; s1] tyin) tyout in

  C.prune constraints to_keep;
  let fty = { modmsf; tyin; tyout; constraints; resulting_corruption; } in
  if !Glob_options.debug then
    Format.eprintf
      "Before optimization:@.%a@."
      pp_funty
      (f.f_name.fn_name, fty);
  let tomax = List.fold_left add [] tyin in
  let tomin = List.fold_left add [n1; s1] tyout in
  C.optimize constraints ~tomin ~tomax;
  fty


let ty_prog is_ct_asm (prog:('info, 'asm) prog) fl =
  let prog = Liveness.liveness false prog in
  let prog = snd prog in
  let fenv = { env_ty = Hf.create 101; env_def = prog } in
  let fl =
    if fl = [] then
      List.rev_map (fun f -> f.f_name) prog
    else
      let get fn =
        try (List.find (fun f -> f.f_name.fn_name = fn) prog).f_name
        with Not_found ->
          hierror ~loc:Lnone ~kind:"speculative constant type checker" "unknown function %s" fn in
      List.map get fl in
  List.map (fun fn -> fn.fn_name, ty_fun is_ct_asm fenv fn) fl

(* ------------------------------------------------------------------------------- *)
(* Inference of msf_info needed by the compiler                                    *)

let compile_infer_msf (prog:('info, 'asm) prog) =
  let prog = snd prog in
  let fenv = { env_ty = Hf.create 101; env_def = prog } in

  let env = Env.init () in
  (* dummy infos *)
  let resulting_corruption = Env.public2 env in
  let notmsf = IsNormal (Direct resulting_corruption) in
  let constraints = C.init() in

  let infer_fun f =
    let sig_annot = SecurityAnnotations.get_sct_signature f.f_annot.f_user_annot in

    let process_return i annot =
      let ls = parse_var_annot ~msf:true annot in
      let an = Option.bind sig_annot (SecurityAnnotations.get_nth_result i) in
      List.mem Msf ls || an = Some SecurityAnnotations.Msf
    in

    (* process function outputs *)
    let tyout = List.mapi process_return f.f_ret_info.ret_annot in

    (* infer the set of input variables that need to be msf *)
    let msfin =
      infer_msf_c ~withcheck:false fenv (Hashtbl.create 13) f.f_body
        (List.fold_left2 (fun s x msf ->
           if msf then Sv.add (L.unloc x) s else s)
         Sv.empty f.f_ret tyout)
    in
    let mkmsf msf = if msf then IsMsf else notmsf in
    let tyout = List.map mkmsf tyout in
    let tyin  = List.map (fun x -> mkmsf (Sv.mem x msfin)) f.f_args in

   (* infer modmsf and check consistency with user info *)
   let modmsf = modmsf_c fenv f.f_body in

   let fty = {
       modmsf;
       tyin;
       tyout;
       constraints; (* dummy info *)
       resulting_corruption; (* dummy info *)
     }  in
   Hf.add fenv.env_ty f.f_name fty
  in
  List.iter infer_fun (List.rev prog);

  let do_t = function
     | IsNormal _ -> Slh_lowering.Slh_None
     | IsMsf      -> Slh_msf
  in

  let do_f _fn fty =
    let in_t = List.map do_t fty.tyin in
    let out_t = List.map do_t fty.tyout in
    (in_t, out_t)
  in

  Hf.map do_f fenv.env_ty