(*
  This file is part of scilla.

  Copyright (c) 2018 - present Zilliqa Research Pvt. Ltd.
  
  scilla is free software: you can redistribute it and/or modify it under the
  terms of the GNU General Public License as published by the Free Software
  Foundation, either version 3 of the License, or (at your option) any later
  version.
 
  scilla is distributed in the hope that it will be useful, but WITHOUT ANY
  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 
  You should have received a copy of the GNU General Public License along with
*)

open Core
open Sexplib.Std
open Literal
open Syntax
open Datatypes
open TypeUtil
open ContractUtil
module CFLiteral = GlobalLiteral
module CFType = CFLiteral.LType
module CFIdentifier = CFType.TIdentifier
module CFName = CFIdentifier.Name

module CashflowRep (R : Rep) = struct
  type money_tag =
    | NoInfo
    | NotMoney
    | Money
    | Map of money_tag
    | Adt of
        CFName.t
        * money_tag list (* name of adt paired with tags of type params *)
    | Inconsistent
  [@@deriving sexp, to_yojson, equal]

  let rec money_tag_to_string tag =
    match tag with
    | Adt (n, ts) ->
        "(" ^ CFName.as_error_string n ^ " "
        ^ String.concat ~sep:" " (List.map ~f:money_tag_to_string ts)
        ^ ")"
    | Map t -> "(Map " ^ money_tag_to_string t ^ ")"
    | _ -> sexp_of_money_tag tag |> Sexplib.Sexp.to_string

  type rep = money_tag * R.rep [@@deriving sexp, to_yojson]

  let get_loc r = match r with _, rr -> R.get_loc rr
  let dummy_rep = (NoInfo, R.dummy_rep)
  let mk_rep (t : money_tag) (s : R.rep) = (t, s)
  let address_rep = mk_rep NoInfo R.address_rep
  let uint128_rep = mk_rep NoInfo R.uint128_rep
  let uint32_rep = mk_rep NoInfo R.uint32_rep
  let bnum_rep = mk_rep NoInfo R.bnum_rep
  let string_rep = mk_rep NoInfo R.string_rep
  let parse_rep s = (NoInfo, R.parse_rep s)
  let get_rep_str r = match r with _, rr -> R.get_rep_str rr
end

module ScillaCashflowChecker
    (SR : Rep) (ER : sig
      include Rep

      val get_type : rep -> PlainTypes.t inferred_type [@@warning "-32"]
    end) =
struct
  module SCFR = SR
  module ECFR = CashflowRep (ER)
  module TypedSyntax = ScillaSyntax (SR) (ER) (CFLiteral)
  module CFSyntax = ScillaSyntax (SCFR) (ECFR) (CFLiteral)
  open TypedSyntax
  open ECFR

  (*******************************************************)
  (*     Initial traversal: Set every tag to NoInfo      *)
  (*******************************************************)

  let add_tag_to_id s tag =
    let open CFIdentifier in
    mk_id (get_id s) (tag, get_rep s)

  (* Lift Ident (n, rep) to Ident (n, (NoInfo, rep)) *)
  let add_noinfo_to_ident i =
    let open CFIdentifier in
    mk_id (get_id i) (ECFR.mk_rep NoInfo (get_rep i))

  let add_money_or_mapmoney_to_ident i typ =
    let open CFType in
    let rec create_money_tag typ =
      match typ with
      | MapType (_, vtyp) -> ECFR.Map (create_money_tag vtyp)
      | _ -> ECFR.Money
    in
    add_tag_to_id i (create_money_tag typ)

  let add_noinfo_to_builtin (op, rep) = (op, (ECFR.NoInfo, rep))

  let rec cf_init_gas_charge gc =
    let open TypedSyntax.SGasCharge in
    match gc with
    | StaticCost i -> CFSyntax.SGasCharge.StaticCost i
    | SizeOf v -> CFSyntax.SGasCharge.SizeOf v
    | ValueOf v -> CFSyntax.SGasCharge.ValueOf v
    | LengthOf v -> CFSyntax.SGasCharge.LengthOf v
    | MapSortCost m -> CFSyntax.SGasCharge.MapSortCost m
    | SumOf (g1, g2) ->
        CFSyntax.SGasCharge.SumOf (cf_init_gas_charge g1, cf_init_gas_charge g2)
    | ProdOf (g1, g2) ->
        CFSyntax.SGasCharge.ProdOf (cf_init_gas_charge g1, cf_init_gas_charge g2)
    | MinOf (g1, g2) ->
        CFSyntax.SGasCharge.MinOf (cf_init_gas_charge g1, cf_init_gas_charge g2)
    | DivCeil (g1, g2) -> CFSyntax.SGasCharge.DivCeil (cf_init_gas_charge g1, g2)
    | LogOf g -> CFSyntax.SGasCharge.LogOf (cf_init_gas_charge g)

  let rec cf_init_tag_pattern p =
    match p with
    | Wildcard -> CFSyntax.Wildcard
    | Binder x -> CFSyntax.Binder (add_noinfo_to_ident x)
    | Constructor (cn, ps) ->
        CFSyntax.Constructor (cn, List.map ~f:cf_init_tag_pattern ps)

  let cf_init_tag_payload p =
    match p with
    | MLit l -> CFSyntax.MLit l
    | MVar v -> CFSyntax.MVar (add_noinfo_to_ident v)

  let cf_init_tag_bcinfo p =
    match p with
    | CurBlockNum -> CFSyntax.CurBlockNum
    | ChainID -> CFSyntax.ChainID
    | Timestamp v -> CFSyntax.Timestamp (add_noinfo_to_ident v)
    | ReplicateContr (addr, iparams) ->
        CFSyntax.ReplicateContr
          (add_noinfo_to_ident addr, add_noinfo_to_ident iparams)

  let rec cf_init_tag_expr erep =
    let e, rep = erep in
    let res_e =
      match e with
      | Literal l -> CFSyntax.Literal l
      | Var i -> CFSyntax.Var (add_noinfo_to_ident i)
      | Fun (arg, t, body) ->
          CFSyntax.Fun (add_noinfo_to_ident arg, t, cf_init_tag_expr body)
      | GasExpr (g, body) ->
          CFSyntax.GasExpr (cf_init_gas_charge g, cf_init_tag_expr body)
      | App (f, actuals) ->
          CFSyntax.App
            (add_noinfo_to_ident f, List.map ~f:add_noinfo_to_ident actuals)
      | Builtin (op, targs, actuals) ->
          CFSyntax.Builtin
            ( add_noinfo_to_builtin op,
              targs,
              List.map ~f:add_noinfo_to_ident actuals )
      | Let (i, topt, lhs, rhs) ->
          CFSyntax.Let
            ( add_noinfo_to_ident i,
              topt,
              cf_init_tag_expr lhs,
              cf_init_tag_expr rhs )
      | Constr (cname, ts, actuals) ->
          CFSyntax.Constr (cname, ts, List.map ~f:add_noinfo_to_ident actuals)
      | MatchExpr (x, clauses) ->
          CFSyntax.MatchExpr
            ( add_noinfo_to_ident x,
              List.map
                ~f:(fun (p, e) -> (cf_init_tag_pattern p, cf_init_tag_expr e))
                clauses )
      | Fixpoint (f, t, body) ->
          CFSyntax.Fixpoint (add_noinfo_to_ident f, t, cf_init_tag_expr body)
      | TFun (tvar, body) ->
          CFSyntax.TFun (add_noinfo_to_ident tvar, cf_init_tag_expr body)
      | TApp (tf, arg_types) -> CFSyntax.TApp (add_noinfo_to_ident tf, arg_types)
      | Message bs ->
          CFSyntax.Message
            (List.map ~f:(fun (s, p) -> (s, cf_init_tag_payload p)) bs)
    in
    (res_e, (ECFR.NoInfo, rep))

  let rec cf_init_tag_stmt srep =
    let s, rep = srep in
    let res_s =
      match s with
      | Load (x, y) ->
          CFSyntax.Load (add_noinfo_to_ident x, add_noinfo_to_ident y)
      | RemoteLoad (x, adr, y) ->
          CFSyntax.RemoteLoad
            ( add_noinfo_to_ident x,
              add_noinfo_to_ident adr,
              add_noinfo_to_ident y )
      | Store (x, y) ->
          CFSyntax.Store (add_noinfo_to_ident x, add_noinfo_to_ident y)
      | Bind (x, e) -> CFSyntax.Bind (add_noinfo_to_ident x, cf_init_tag_expr e)
      | MapUpdate (m, ks, v) ->
          CFSyntax.MapUpdate
            ( add_noinfo_to_ident m,
              List.map ~f:add_noinfo_to_ident ks,
              match v with
              | None -> None
              | Some v' -> Some (add_noinfo_to_ident v') )
      | MapGet (x, m, ks, retrieve) ->
          CFSyntax.MapGet
            ( add_noinfo_to_ident x,
              add_noinfo_to_ident m,
              List.map ~f:add_noinfo_to_ident ks,
              retrieve )
      | RemoteMapGet (x, adr, m, ks, retrieve) ->
          CFSyntax.RemoteMapGet
            ( add_noinfo_to_ident x,
              add_noinfo_to_ident adr,
              add_noinfo_to_ident m,
              List.map ~f:add_noinfo_to_ident ks,
              retrieve )
      | MatchStmt (x, pss) ->
          CFSyntax.MatchStmt
            ( add_noinfo_to_ident x,
              List.map
                ~f:(fun (p, ss) ->
                  (cf_init_tag_pattern p, List.map ~f:cf_init_tag_stmt ss))
                pss )
      | ReadFromBC (x, s) ->
          CFSyntax.ReadFromBC (add_noinfo_to_ident x, cf_init_tag_bcinfo s)
      | TypeCast (x, r, t) ->
          CFSyntax.TypeCast (add_noinfo_to_ident x, add_noinfo_to_ident r, t)
      | AcceptPayment -> CFSyntax.AcceptPayment
      | Return i -> CFSyntax.Return (add_noinfo_to_ident i)
      | SendMsgs x -> CFSyntax.SendMsgs (add_noinfo_to_ident x)
      | CreateEvnt x -> CFSyntax.CreateEvnt (add_noinfo_to_ident x)
      | CallProc (id_opt, p, args) ->
          let id = Option.map id_opt ~f:(fun id -> add_noinfo_to_ident id) in
          CFSyntax.CallProc (id, p, List.map args ~f:add_noinfo_to_ident)
      | Iterate (l, p) -> CFSyntax.Iterate (add_noinfo_to_ident l, p)
      | Throw xopt -> (
          match xopt with
          | Some x -> CFSyntax.Throw (Some (add_noinfo_to_ident x))
          | None -> CFSyntax.Throw None)
      | GasStmt g -> CFSyntax.GasStmt (cf_init_gas_charge g)
    in
    (res_s, rep)

  let cf_init_tag_component component =
    let { comp_type; comp_name; comp_params; comp_body; comp_return } =
      component
    in
    {
      CFSyntax.comp_type;
      CFSyntax.comp_name;
      CFSyntax.comp_params =
        List.map ~f:(fun (x, t) -> (add_noinfo_to_ident x, t)) comp_params;
      CFSyntax.comp_body = List.map ~f:cf_init_tag_stmt comp_body;
      CFSyntax.comp_return;
    }

  let cf_init_tag_contract contract token_fields =
    let { cname; cparams; cconstraint; cfields; ccomps } = contract in
    let token_fields_contains x =
      List.mem token_fields (CFIdentifier.as_string x) ~equal:String.( = )
    in
    {
      CFSyntax.cname;
      CFSyntax.cparams =
        List.map
          ~f:(fun (x, t) ->
            ( (if token_fields_contains x then add_money_or_mapmoney_to_ident x t
              else add_noinfo_to_ident x),
              t ))
          cparams;
      CFSyntax.cconstraint = cf_init_tag_expr cconstraint;
      CFSyntax.cfields =
        List.map
          ~f:(fun (x, t, e) ->
            ( (if token_fields_contains x then add_money_or_mapmoney_to_ident x t
              else add_noinfo_to_ident x),
              t,
              cf_init_tag_expr e ))
          cfields;
      CFSyntax.ccomps = List.map ~f:cf_init_tag_component ccomps;
    }

  let cf_init_tag_type_def tdef =
    let { cname; c_arg_types } = tdef in
    { CFSyntax.cname = add_noinfo_to_ident cname; CFSyntax.c_arg_types }

  let cf_init_tag_library lib =
    let { lname; lentries } = lib in
    let init_tag_entry entry =
      match entry with
      | LibVar (lname, ltype, lexp) ->
          CFSyntax.LibVar
            (add_noinfo_to_ident lname, ltype, cf_init_tag_expr lexp)
      | LibTyp (lname, type_defs) ->
          CFSyntax.LibTyp
            ( add_noinfo_to_ident lname,
              List.map ~f:cf_init_tag_type_def type_defs )
    in
    { CFSyntax.lname; CFSyntax.lentries = List.map ~f:init_tag_entry lentries }

  let cf_init_tag_module cmod token_fields =
    let { smver; libs; elibs; contr } = cmod in
    let res_libs =
      match libs with None -> None | Some l -> Some (cf_init_tag_library l)
    in
    {
      CFSyntax.smver;
      CFSyntax.libs = res_libs;
      CFSyntax.elibs;
      CFSyntax.contr = cf_init_tag_contract contr token_fields;
    }

  (*******************************************************)
  (*                  Find fixpoint                      *)
  (* Strategy:                                           *)
  (*  - Expressions have an expected tag, which is       *)
  (* unified with the tag that can be extrapolated from  *)
  (* the usage.                                          *)
  (*  - Statement lists are first traversed to collect   *)
  (* all declared local variable and their tags. Then    *)
  (* they are traversed again to extrapolate new tags.   *)
  (*******************************************************)
  open CFSyntax

  [@@@ocamlformat "disable"]

  (* Least upper bound in the money_tag lattice. *)
  let rec lub_tags t1 t2 =
    match t1, t2 with
    | Inconsistent , _
    | _            , Inconsistent  -> Inconsistent
    | NoInfo       , x
    | x            , NoInfo        -> x
    | Map x        , Map y         -> Map (lub_tags x y)
    | Adt (n1, ts1), Adt (n2, ts2)
      when [%equal : CFName.t ] n1 n2 ->
        (match List.map2 ts1 ts2 ~f:lub_tags with
         | Ok res -> Adt (n1, res)
         | Unequal_lengths -> Inconsistent)
    | Money        , Money         -> Money
    | NotMoney     , NotMoney      -> NotMoney
    | _            , _             -> Inconsistent

  (* Greatest lower bound in the money_tag lattice. *)
  let rec glb_tags t1 t2 =
    match t1, t2 with
    | Inconsistent , x
    | x            , Inconsistent  -> x
    | NoInfo       , _
    | _            , NoInfo        -> NoInfo
    | Map x        , Map y         -> Map (glb_tags x y)
    | Adt (n1, ts1), Adt (n2, ts2)
      when [%equal : CFName.t ] n1 n2 ->
        (match List.map2 ts1 ts2 ~f:glb_tags with
         | Ok res -> Adt (n1, res)
         | Unequal_lengths -> Inconsistent)
    | Money        , Money         -> Money
    | NotMoney     , NotMoney      -> NotMoney
    | _            , _             -> NoInfo

  [@@@ocamlformat "enable"]

  (*******************************************************)
  (*           Helper functions for ADTs                 *)
  (*******************************************************)

  let ctr_to_adt_tag ctr_name arg_tags =
    let open DataTypeDictionary in
    match lookup_constructor ctr_name with
    | Error _ ->
        (* We don't allow failures at this stage of the checker *)
        Inconsistent
    | Ok (adt, _) -> (
        match adt.tparams with
        | [] ->
            (* No type parameters - check for special cases *)
            (* Case 1 (Bool case): No constructor takes arguments : NotMoney *)
            if
              List.for_all adt.tmap ~f:(fun (_, arg_typs) ->
                  match arg_typs with [] -> true | _ -> false)
            then NotMoney
              (* Case 2 (Nat case): 2 constructors. One constructor takes
                 1 argument of same type, other constructor
                 takes no argument : NoInfo *)
            else if
              List.length adt.tconstr = 2
              && List.exists adt.tmap ~f:(fun (_, arg_typs) ->
                     List.is_empty arg_typs)
              && List.exists adt.tmap ~f:(fun (_, arg_typs) ->
                     match arg_typs with
                     | [ ADT (arg_typ_name, _) ] ->
                         [%equal: CFName.t]
                           (CFIdentifier.get_id arg_typ_name)
                           adt.tname
                     | _ -> false)
            then NoInfo
            else Adt (adt.tname, [])
        | _ -> (
            (* Deduce mapping from type arguments to tags based on
               constructor argument tags *)
            match constr_tmap adt ctr_name with
            | None ->
                (* No tmap for constructor = doesn't take arguments *)
                Adt (adt.tname, List.map adt.tparams ~f:(fun _ -> NoInfo))
            | Some arg_typs ->
                (* Default mapping : 'A -> NoInfo *)
                let init_targ_to_tag_map =
                  List.map adt.tparams ~f:(fun tparam -> (tparam, NoInfo))
                in
                let update_targ_tag targ new_tag map =
                  let current_tag =
                    match List.Assoc.find map ~equal:String.( = ) targ with
                    | None -> Inconsistent
                    | Some t -> t
                  in
                  List.Assoc.add map ~equal:String.( = ) targ
                    (lub_tags current_tag new_tag)
                in
                let rec match_arg_tag_with_typ arg_typ arg_tag targ_tag_map =
                  let open CFType in
                  match arg_typ with
                  | TypeVar v -> update_targ_tag v arg_tag targ_tag_map
                  | MapType (_, vt) | FunType (_, vt) -> (
                      match arg_tag with
                      | Map vtag -> match_arg_tag_with_typ vt vtag targ_tag_map
                      | NoInfo -> match_arg_tag_with_typ vt NoInfo targ_tag_map
                      | _ -> targ_tag_map)
                  | ADT (_, adt_params) -> (
                      match arg_tag with
                      | Adt (_, adt_param_tags) ->
                          let new_map, _ =
                            List.fold_left adt_params
                              ~init:(targ_tag_map, adt_param_tags)
                              ~f:(fun (map, adt_param_tags) adt_param ->
                                match adt_param_tags with
                                | [] -> (map, [])
                                | t :: ts ->
                                    (match_arg_tag_with_typ adt_param t map, ts))
                          in
                          new_map
                      | _ -> targ_tag_map)
                  | PrimType _ | PolyFun (_, _) | Unit -> targ_tag_map
                  | Address _ -> (* TODO *) targ_tag_map
                in
                let tvar_tag_map, _ =
                  List.fold_left arg_typs ~init:(init_targ_to_tag_map, arg_tags)
                    ~f:(fun (map, arg_tags) arg_typ ->
                      match arg_tags with
                      | [] -> (map, [])
                      | t :: ts -> (match_arg_tag_with_typ arg_typ t map, ts))
                in
                let final_adt_arg_tags =
                  List.map adt.tparams ~f:(fun tparam ->
                      match
                        List.Assoc.find tvar_tag_map ~equal:String.( = ) tparam
                      with
                      | None -> Inconsistent
                      | Some t -> t)
                in
                Adt (adt.tname, final_adt_arg_tags)))

  let ctr_pattern_to_subtags ctr_name expected_tag =
    let open DataTypeDictionary in
    match lookup_constructor ctr_name with
    | Error _ ->
        (* Catch errors as We don't allow failures at this stage of the checker *)
        None
    | Ok (adt, _) -> (
        match constr_tmap adt ctr_name with
        | None ->
            (* No tmap = No constructor arguments *)
            Some []
        | Some tmap ->
            let tvar_tag_map =
              let zipped_tvar_tags =
                match expected_tag with
                | Adt (exp_typ_name, arg_tags)
                  when [%equal: CFName.t] exp_typ_name adt.tname ->
                    List.zip adt.tparams arg_tags
                | NoInfo (* Nothing known *)
                | Money (* Nat case *)
                | NotMoney (* Nat case *) ->
                    Ok
                      (List.map adt.tparams ~f:(fun tparam ->
                           (tparam, expected_tag)))
                | _ ->
                    (* Don't let inconsistent tags infect subpatterns *)
                    Ok
                      (List.map adt.tparams ~f:(fun tparam -> (tparam, NoInfo)))
              in
              match zipped_tvar_tags with
              | Unequal_lengths ->
                  (* Can only happen if arg_tags in Adt case has wrong length *)
                  List.map adt.tparams ~f:(fun tparam -> (tparam, Inconsistent))
              | Ok map -> map
            in
            let rec tag_tmap t =
              let open CFType in
              match t with
              | PrimType _ ->
                  (* TODO: Fixed constructor argument types
                           should be analysed *)
                  NoInfo
              | MapType (_, vt) | FunType (_, vt) -> Map (tag_tmap vt)
              | ADT (adt_name, arg_typs) ->
                  Adt
                    (CFIdentifier.get_id adt_name, List.map arg_typs ~f:tag_tmap)
              | TypeVar tvar -> (
                  match
                    List.Assoc.find tvar_tag_map ~equal:String.( = ) tvar
                  with
                  | Some tag -> tag
                  | None -> Inconsistent)
              | _ -> Inconsistent
            in
            Some (List.map tmap ~f:tag_tmap))

  let ctr_arg_filter targ =
    let open CFType in
    match targ with
    | PrimType _ | MapType _ | FunType _ -> true
    | ADT _ (* TODO: Detect induction, and ignore only when inductive *)
    | TypeVar _ (* TypeVars tagged at type level *) | PolyFun _ (* Ignore *)
    | Unit ->
        (* Ignore *)
        false
    | Address _ -> (* TODO *) false

  let init_ctr_tag_map () =
    let open DataTypeDictionary in
    let all_ctrs = get_all_ctrs () in
    List.fold_left all_ctrs ~init:[] ~f:(fun acc (adt, ctr) ->
        match constr_tmap adt ctr.cname with
        | None ->
            (* No constructor arguments - ignore *)
            acc
        | Some targs ->
            let tag_list =
              List.map targs ~f:(fun t ->
                  if ctr_arg_filter t then Some NoInfo else None)
            in
            if List.exists tag_list ~f:Option.is_some then
              List.Assoc.add acc ~equal:String.( = )
                (CFName.as_string ctr.cname)
                tag_list
            else acc)

  let update_ctr_tag_map ctr_tag_map ctr_name arg_tags =
    match List.Assoc.find ctr_tag_map ~equal:String.( = ) ctr_name with
    | None ->
        (* Ignored constructor *)
        None
    | Some arg_map_tags -> (
        match
          List.map2 arg_map_tags arg_tags ~f:(fun arg_map_tag arg_tag ->
              Option.map arg_map_tag ~f:(lub_tags arg_tag))
        with
        | Ok new_tags
          when not
               @@ [%equal: ECFR.money_tag Option.t list] new_tags arg_map_tags
          ->
            Some
              (List.Assoc.add ctr_tag_map ~equal:String.( = ) ctr_name new_tags)
        | _ -> None)

  let option_adt_tag t = Adt (CFName.parse_simple_name "Option", [ t ])

  (*******************************************************)
  (*      Helper functions for local variables           *)
  (*******************************************************)

  let get_id_tag id = match CFIdentifier.get_rep id with tag, _ -> tag

  let update_id_tag id new_tag =
    let open CFIdentifier in
    match get_rep id with _, rep -> mk_id (get_id id) (new_tag, rep)

  let lookup_var_tag i env =
    let open CFIdentifier in
    match AssocDictionary.lookup (as_string i) env with
    | Some t -> t
    | None -> get_id_tag i

  let lookup_var_tag2 i env1 env2 =
    let open CFIdentifier in
    match AssocDictionary.lookup (as_string i) env1 with
    | Some t -> t
    | None -> lookup_var_tag i env2

  let update_var_tag i t env =
    AssocDictionary.update (CFIdentifier.as_string i) t env

  let update_var_tag2 i t env1 env2 =
    match AssocDictionary.lookup (CFIdentifier.as_string i) env1 with
    | Some _ -> (update_var_tag i t env1, env2)
    | None -> (env1, update_var_tag i t env2)

  let update_vars_tags2 ids env1 env2 =
    List.fold_left ids ~init:(env1, env2) ~f:(fun (acc_env1, acc_env2) id ->
        update_var_tag2 id (get_id_tag id) acc_env1 acc_env2)

  let update_ids_tags2 ids env1 env2 =
    List.map ids ~f:(fun i ->
        let i_tag = lookup_var_tag2 i env1 env2 in
        update_id_tag i i_tag)

  let update_ids_tags ids env =
    let dummy_env = AssocDictionary.make_dict () in
    update_ids_tags2 ids env dummy_env

  (*******************************************************)
  (*      Helper functions for builtin functions         *)
  (*******************************************************)

  (* Calculate the signature of a builtin function.

     Step 1: Calculate candidate signatures based on the
     desired result tag and each argument tag.

     - For each tag t, pick every least upper bound of that
     tag that makes sense for that result/argument.
     Call these bounds b_t.

     - For each b_t, find all sets of tags satisfying that
     the use of those tags in the other argument/result positions
     is the greatest lower bound of a consistent use of tags
     satisfying b_t. These sets along with b_t are considered
     the candidate sigantures for t, called C_t.

     Step 2: Consider the elements of C_t1 x C_t2 x ..., i.e.,
     the cartesian product of the candidate signature sets for
     each t.

     For each element, calculate the least upper bound of all
     the tags in the signatures of the element. Call the
     resulting set of candidate signatures C.

     Step 3: Calculate the greatest lower bound of C. *)
  [@@@ocamlformat "disable"]

  let builtin_signature f res_tag args_tags =
    let lub_sigs c_rs c_ass =
      let c =
        List.fold_left c_ass ~init:c_rs 
          ~f:(fun partial_c c_t ->
             List.fold_left partial_c ~init:[]
               ~f:(fun acc_partial_c (partial_c_res_tag, partial_c_args_tags) ->
                  List.fold_left c_t ~init:acc_partial_c
                    ~f:(fun acc_lub_sigs (c_t_res_tag, c_t_args_tags) ->
                        let res_sigs =
                          match List.map2 ~f:lub_tags partial_c_args_tags c_t_args_tags with
                          | Ok res          -> (lub_tags partial_c_res_tag c_t_res_tag, res)
                          | Unequal_lengths -> (Inconsistent, [NoInfo]) in
                        
                        res_sigs :: acc_lub_sigs))) in
      List.fold_left c ~init:( Inconsistent , List.map ~f:(fun _ -> Inconsistent ) args_tags )
        ~f:(fun (glb_res_tag, glb_args_tags) (c_res_tag, c_args_tags) ->
            ( glb_tags glb_res_tag c_res_tag,
              match List.map2 ~f:glb_tags glb_args_tags c_args_tags with
              | Ok res -> res
              | Unequal_lengths -> [Inconsistent]))
        in
    let (c_r, c_as) =
      match fst f with
      | Builtin_put ->
          let c_r_sigs =
            match res_tag with
            | Map t  -> [ ( Map t        , [ Map t      ; NotMoney ; t      ] ) ]
            | NoInfo -> [ ( Map NoInfo   , [ Map NoInfo ; NotMoney ; NoInfo ] ) ]
            | _      -> [ ( Inconsistent , [ Map NoInfo ; NotMoney ; NoInfo ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ; k ; v ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( Map t      , [ Map t        ; NotMoney ; t      ] ) ] 
                  | NoInfo -> [ ( Map NoInfo , [ Map NoInfo   ; NotMoney ; NoInfo ] ) ] 
                  | _      -> [ ( Map NoInfo , [ Inconsistent ; NotMoney ; NoInfo ] ) ] in
                let k_sig =
                  match k with
                  | NotMoney
                  | NoInfo   -> [ ( Map NoInfo , [ Map NoInfo ; NotMoney     ; NoInfo ] ) ]
                  | _        -> [ ( Map NoInfo , [ Map NoInfo ; Inconsistent ; NoInfo ] ) ] in
                let v_sig =
                  match v with
                  | _        -> [ ( Map v , [ Map v ; NotMoney ; v ] ) ] in
                [ m_sig ; k_sig ; v_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_remove ->
          let c_r_sigs =
            match res_tag with
            | Map t  -> [ ( Map t        , [ Map t      ; NotMoney ] ) ]
            | NoInfo -> [ ( Map NoInfo   , [ Map NoInfo ; NotMoney ] ) ]
            | _      -> [ ( Inconsistent , [ Map NoInfo ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ; k ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( Map t      , [ Map t        ; NotMoney ] ) ] 
                  | NoInfo -> [ ( Map NoInfo , [ Map NoInfo   ; NotMoney ] ) ] 
                  | _      -> [ ( Map NoInfo , [ Inconsistent ; NotMoney ] ) ] in
                let k_sig =
                  match k with
                  | NotMoney
                  | NoInfo   -> [ ( Map NoInfo , [ Map NoInfo ; NotMoney     ] ) ]
                  | _        -> [ ( Map NoInfo , [ Map NoInfo ; Inconsistent ] ) ] in
                [ m_sig ; k_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_get ->
          let c_r_sigs =
            match res_tag with
            | Adt (adt_name, [t])
              when is_option_adt_name adt_name ->
                                     [ (option_adt_tag t     , [ Map t      ; NotMoney ] ) ]
            | NoInfo              -> [ (option_adt_tag NoInfo, [ Map NoInfo ; NotMoney ] ) ]
            | _                   -> [ ( Inconsistent        , [ Map NoInfo ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ; k ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( option_adt_tag t     , [ Map t        ; NotMoney ] ) ] 
                  | NoInfo -> [ ( option_adt_tag NoInfo, [ Map NoInfo   ; NotMoney ] ) ] 
                  | _      -> [ ( option_adt_tag NoInfo, [ Inconsistent ; NotMoney ] ) ] in
                let k_sig =
                  match k with
                  | NotMoney
                  | NoInfo   -> [ ( option_adt_tag NoInfo, [ Map NoInfo ; NotMoney     ] ) ]
                  | _        -> [ ( option_adt_tag NoInfo, [ Map NoInfo ; Inconsistent ] ) ] in
                [ m_sig ; k_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_contains ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ Map NoInfo ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ Map NoInfo ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ; k ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( NotMoney , [ Map t        ; NotMoney ] ) ] 
                  | NoInfo -> [ ( NotMoney , [ Map NoInfo   ; NotMoney ] ) ] 
                  | _      -> [ ( NotMoney , [ Inconsistent ; NotMoney ] ) ] in
                let k_sig =
                  match k with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ Map NoInfo ; NotMoney     ] ) ]
                  | _        -> [ ( NotMoney , [ Map NoInfo ; Inconsistent ] ) ] in
                [ m_sig ; k_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_to_list ->
          (* Lists not supported, so use Inconsistent *)
          let c_r_sigs =
            [ ( Inconsistent , [ Map NoInfo ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( Inconsistent , [ Map t        ] ) ] 
                  | NoInfo -> [ ( Inconsistent , [ Map NoInfo   ] ) ] 
                  | _      -> [ ( Inconsistent , [ Inconsistent ] ) ] in
                [ m_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_size ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ Map NoInfo ] ) ]
            | _        -> [ ( Inconsistent , [ Map NoInfo ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ m ] ->
                let m_sig =
                  match m with
                  | Map t  -> [ ( NotMoney , [ Map t        ] ) ] 
                  | NoInfo -> [ ( NotMoney , [ Map NoInfo   ] ) ] 
                  | _      -> [ ( NotMoney , [ Inconsistent ] ) ] in
                [ m_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_eq
      | Builtin_lt ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ;
                            ( NotMoney     , [ Money    ; Money    ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ;
                            ( Inconsistent , [ Money    ; Money    ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | Money    -> [ ( NotMoney , [ Money        ; Money    ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ;
                                  ( NotMoney , [ Money        ; Money    ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ; NoInfo   ] ) ] in
                let v2_sig =
                  match v2 with
                  | Money    -> [ ( NotMoney , [ Money    ; Money        ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ;
                                  ( NotMoney , [ Money    ; Money        ] ) ]
                  | _        -> [ ( NotMoney , [ NoInfo   ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_add
      | Builtin_sub ->
          let c_r_sigs =
            match res_tag with
            | NotMoney -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ]
            | Money    -> [ ( Money        , [ Money    ; Money    ] ) ]
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ;
                            ( Money        , [ Money    ; Money    ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ;
                            ( Inconsistent , [ Money    ; Money    ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | Money    -> [ ( Money    , [ Money        ; Money    ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ;
                                  ( Money    , [ Money        ; Money    ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ; NotMoney ] ) ;
                                  ( Money    , [ Inconsistent ; Money    ] ) ] in
                let v2_sig =
                  match v2 with
                  | Money    -> [ ( Money    , [ Money    ; Money        ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ;
                                  ( Money    , [ Money    ; Money        ] ) ]
                  | _        -> [ ( NotMoney , [ NotMoney ; Inconsistent ] ) ;
                                  ( Money    , [ Money    ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_mul ->
          let c_r_sigs =
            match res_tag with
            | NotMoney -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ]
            | Money    -> [ ( Money        , [ NotMoney ; Money    ] ) ;
                            ( Money        , [ Money    ; NotMoney ] ) ]
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ;
                            ( Money        , [ NotMoney ; Money    ] ) ;
                            ( Money        , [ Money    ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ;
                            ( Inconsistent , [ NotMoney ; Money    ] ) ;
                            ( Inconsistent , [ Money    ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | Money    -> [ ( Money    , [ Money        ; NotMoney ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ;
                                  ( Money    , [ NotMoney     ; Money    ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ;
                                  ( Money    , [ NotMoney     ; Money    ] ) ;
                                  ( Money    , [ Money        ; NotMoney ] ) ]
                  | _        -> [ ( Money    , [ Inconsistent ; Money    ] ) ;
                                  ( Money    , [ Inconsistent ; NotMoney ] ) ;
                                  ( NotMoney , [ Inconsistent ; NotMoney ] ) ] in
                let v2_sig =
                  match v2 with
                  | Money    -> [ ( Money    , [ NotMoney ; Money        ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ;
                                  ( Money    , [ Money    ; NotMoney     ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ;
                                  ( Money    , [ NotMoney ; Money        ] ) ;
                                  ( Money    , [ Money    ; NotMoney     ] ) ]
                  | _        -> [ ( Money    , [ Money    ; Inconsistent ] ) ;
                                  ( Money    , [ NotMoney ; Inconsistent ] ) ;
                                  ( NotMoney , [ NotMoney ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_pow ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ; NotMoney ] ) ] in
                let v2_sig =
                  match v2 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ]
                  | _        -> [ ( NotMoney , [ NotMoney ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)          
      | Builtin_div
      | Builtin_rem ->
          let c_r_sigs =
            match res_tag with
            | NotMoney -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ]
            | Money    -> [ ( Money        , [ Money    ; NotMoney ] ) ]
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ;
                            ( Money        , [ Money    ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ;
                            ( Inconsistent , [ Money    ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | Money    -> [ ( Money    , [ Money        ; NotMoney ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ;
                                  ( Money    , [ Money        ; NotMoney ] ) ]
                  | _        -> [ ( Money    , [ Inconsistent ; NotMoney ] ) ;
                                  ( NotMoney , [ Inconsistent ; NotMoney ] ) ] in
                let v2_sig =
                  match v2 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ;
                                  ( Money    , [ Money    ; NotMoney     ] ) ]
                  | _        -> [ ( Money    , [ Money    ; Inconsistent ] ) ;
                                  ( NotMoney , [ NotMoney ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_to_int32
      | Builtin_to_int64
      | Builtin_to_int128
      | Builtin_to_int256
      | Builtin_to_string
      | Builtin_to_nat    ->
          let c_r_sigs =
            match res_tag with
            | NotMoney -> [ ( NotMoney     , [ NotMoney ] ) ]
            | Money    -> [ ( Money        , [ Money    ] ) ]
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ] ) ;
                            ( Money        , [ Money    ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ] ) ;
                            ( Inconsistent , [ Money    ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ] ->
                let v1_sig =
                  match v1 with
                  | Money    -> [ ( Money    , [ Money        ] ) ] 
                  | NotMoney -> [ ( NotMoney , [ NotMoney     ] ) ]
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ] ) ;
                                  ( Money    , [ Money        ] ) ]
                  | _        -> [ ( Money    , [ Inconsistent ] ) ;
                                  ( NotMoney , [ Inconsistent ] ) ] in
                [ v1_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_sha256hash
      | Builtin_keccak256hash
      | Builtin_ripemd160hash
      | Builtin_schnorr_get_address
      | Builtin_strlen
      | Builtin_to_bystr ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ] ) ] in
          let c_as_sigs = 
            match args_tags with
            | [ v1 ] ->
                let v1_sig =
                  match v1 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ] ) ] in
                [ v1_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_concat
      | Builtin_bech32_to_bystr20
      | Builtin_bystr20_to_bech32
      | Builtin_blt
      | Builtin_badd
      | Builtin_bsub ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ] ->
                let v1_sig =
                  match v1 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ; NotMoney ] ) ] in
                let v2_sig =
                  match v2 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ] ) ]
                  | _        -> [ ( NotMoney , [ NotMoney ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | Builtin_substr
      | Builtin_schnorr_verify ->
          let c_r_sigs =
            match res_tag with
            | NotMoney
            | NoInfo   -> [ ( NotMoney     , [ NotMoney ; NotMoney ; NotMoney ] ) ]
            | _        -> [ ( Inconsistent , [ NotMoney ; NotMoney ; NotMoney ] ) ] in
          let c_as_sigs =
            match args_tags with
            | [ v1 ; v2 ; v3 ] ->
                let v1_sig =
                  match v1 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney     ; NotMoney ; NotMoney ] ) ]
                  | _        -> [ ( NotMoney , [ Inconsistent ; NotMoney ; NotMoney ] ) ] in
                let v2_sig =
                  match v2 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney     ; NotMoney ] ) ]
                  | _        -> [ ( NotMoney , [ NotMoney ; Inconsistent ; NotMoney ] ) ] in
                let v3_sig =
                  match v3 with
                  | NotMoney
                  | NoInfo   -> [ ( NotMoney , [ NotMoney ; NotMoney ; NotMoney     ] ) ]
                  | _        -> [ ( NotMoney , [ NotMoney ; NotMoney ; Inconsistent ] ) ] in
                [ v1_sig ; v2_sig ; v3_sig ]
            | _             ->
                (* Error *)
                [[ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ]] in
          (c_r_sigs, c_as_sigs)
      | _ -> 
          (* Error *)
          let c_r_sigs =
            [ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ] in
          let c_as_sigs =
            [ [ ( Inconsistent , List.map ~f:(fun _ -> Inconsistent) args_tags ) ] ] in
          (c_r_sigs, c_as_sigs) in
    lub_sigs c_r c_as

  (*******************************************************)
  (*            Helper functions for patterns            *)
  (*******************************************************)

  [@@@ocamlformat "enable"]

  let rec get_pattern_vars p acc =
    match p with
    | Wildcard -> acc
    | Binder x -> x :: acc
    | Constructor (_, ps) -> List.fold_right ps ~init:acc ~f:get_pattern_vars

  let update_pattern_vars_tags_from_usage_and_remove_from_env p local_env
      ctr_tag_map =
    let rec walk p local_env ctr_tag_map =
      match p with
      | Wildcard -> (Wildcard, NoInfo, local_env, ctr_tag_map, false)
      | Binder x ->
          let new_x_tag = lookup_var_tag x local_env in
          let new_x = update_id_tag x new_x_tag in
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          ( Binder new_x,
            new_x_tag,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] (get_id_tag x) (get_id_tag new_x) )
      | Constructor (s, ps) ->
          let new_ps, new_ps_tags, new_local_env, ps_ctr_tag_map, ps_changes =
            List.fold_right ps ~init:([], [], local_env, ctr_tag_map, false)
              ~f:(fun
                   p
                   ( acc_ps,
                     acc_ps_tags,
                     acc_local_env,
                     acc_ctr_tag_map,
                     acc_changes )
                 ->
                let new_p, new_p_tag, p_local_env, p_ctr_tag_map, p_changes =
                  walk p acc_local_env acc_ctr_tag_map
                in
                ( new_p :: acc_ps,
                  new_p_tag :: acc_ps_tags,
                  p_local_env,
                  p_ctr_tag_map,
                  acc_changes || p_changes ))
          in
          let new_ctr_tag_map, ctr_tag_map_changes =
            update_ctr_tag_map ps_ctr_tag_map (CFIdentifier.as_string s)
              new_ps_tags
            |> Option.value_map ~default:(ps_ctr_tag_map, false) ~f:(fun map ->
                   (map, true))
          in
          let ctr_tag = ctr_to_adt_tag (CFIdentifier.get_id s) new_ps_tags in
          ( Constructor (s, new_ps),
            ctr_tag,
            new_local_env,
            new_ctr_tag_map,
            ctr_tag_map_changes || ps_changes )
    in
    let new_p, _, new_local_env, new_ctr_tag_map, changes =
      walk p local_env ctr_tag_map
    in
    (new_p, new_local_env, new_ctr_tag_map, changes)

  let update_pattern_vars_tags_from_scrutinee p scrutinee_tag =
    let rec walk p expected_tag =
      match p with
      | Wildcard -> (Wildcard, false)
      | Binder x ->
          let new_x_tag = lub_tags expected_tag (get_id_tag x) in
          let new_x = update_id_tag x new_x_tag in
          (Binder new_x, not ([%equal: ECFR.money_tag] new_x_tag (get_id_tag x)))
      | Constructor (s, ps) ->
          let expected_subtags =
            match
              ctr_pattern_to_subtags (CFIdentifier.get_id s) expected_tag
            with
            | Some ts -> ts
            | None -> List.map ps ~f:(fun _ -> NoInfo)
          in
          let new_subpatterns_with_tags, _ =
            List.zip_with_remainder ps expected_subtags
          in
          let new_ps, changes =
            List.fold_right new_subpatterns_with_tags ~init:([], false)
              ~f:(fun (p, exp_tag) (acc_ps, acc_changes) ->
                let new_p, change = walk p exp_tag in
                (new_p :: acc_ps, acc_changes || change))
          in
          (Constructor (s, new_ps), changes)
    in
    walk p scrutinee_tag

  let insert_pattern_vars_into_env p local_env =
    let pattern_vars = get_pattern_vars p [] in
    List.fold_left pattern_vars ~init:local_env ~f:(fun l_env x ->
        AssocDictionary.insert (CFIdentifier.as_string x) (get_id_tag x) l_env)

  (* Find least upper bound of scrutinee based on patterns and pattern tags *)
  let lub_pattern_tags ps =
    let rec walk acc_tag p =
      match p with
      | Wildcard -> acc_tag
      | Binder x -> lub_tags (get_id_tag x) acc_tag
      | Constructor (s, ps) ->
          let expected_subtags =
            match ctr_pattern_to_subtags (CFIdentifier.get_id s) acc_tag with
            | Some ts -> ts
            | None -> List.map ps ~f:(fun _ -> NoInfo)
          in
          let subpattern_tags =
            match List.map2 expected_subtags ps ~f:walk with
            | Ok tps -> tps
            | Unequal_lengths -> []
          in
          let new_tag =
            ctr_to_adt_tag (CFIdentifier.get_id s) subpattern_tags
          in
          lub_tags new_tag acc_tag
    in
    List.fold_left ps ~init:NoInfo ~f:walk

  (*******************************************************)
  (*            Main cashflow analyzer                   *)
  (*******************************************************)

  let update_var_tag_payload p local_env =
    match p with
    | MLit l -> CFSyntax.MLit l
    | MVar v -> (
        let tag =
          match AssocDictionary.lookup (CFIdentifier.as_string v) local_env with
          | None -> Inconsistent (* Should not happen *)
          | Some t -> t
        in
        match v with
        | Ident (name, (_, rep)) ->
            CFSyntax.MVar (CFIdentifier.mk_id name (tag, rep)))

  let rec cf_tag_expr erep expected_tag param_env local_env ctr_tag_map =
    let lub t = lub_tags expected_tag t in
    let e, (tag, rep) = erep in
    let ( new_e,
          new_e_tag,
          new_param_env,
          new_local_env,
          new_ctr_tag_map,
          new_changes ) =
      match e with
      | Literal _ ->
          (* No need to deduce tag from type.
             If the literal is a number, then nothing can be deduced,
             and if it is not a number, the tag of relevant variables
             will be deduced from their usage. *)
          (e, tag, param_env, local_env, ctr_tag_map, false)
      | Var i ->
          let new_i_tag = lub (lookup_var_tag2 i local_env param_env) in
          let new_i = update_id_tag i new_i_tag in
          let new_local_env, new_param_env =
            update_var_tag2 i new_i_tag local_env param_env
          in
          ( Var new_i,
            new_i_tag,
            new_param_env,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] new_i_tag (get_id_tag i) )
      | Fun (arg, t, body) ->
          (* Using Map tag to represent functions as well as maps *)
          let body_expected_tag =
            match expected_tag with
            | Map x -> x
            | NoInfo -> NoInfo
            | _ -> Inconsistent
          in
          let body_local_env =
            AssocDictionary.insert
              (CFIdentifier.as_string arg)
              (get_id_tag arg) local_env
          in
          let ( ((_, (new_body_tag, _)) as new_body),
                res_param_env,
                res_body_local_env,
                res_ctr_tag_map,
                body_changes ) =
            cf_tag_expr body body_expected_tag param_env body_local_env
              ctr_tag_map
          in
          let res_arg_tag = lookup_var_tag arg res_body_local_env in
          let new_local_env =
            AssocDictionary.remove
              (CFIdentifier.as_string arg)
              res_body_local_env
          in
          ( Fun (update_id_tag arg res_arg_tag, t, new_body),
            Map new_body_tag,
            res_param_env,
            new_local_env,
            res_ctr_tag_map,
            body_changes
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag arg) res_arg_tag) )
      | App (f, args) ->
          let new_args =
            List.map
              ~f:(fun arg ->
                update_id_tag arg (lookup_var_tag2 arg local_env param_env))
              args
          in
          let args_changes =
            match
              List.exists2
                ~f:(fun arg new_arg ->
                  not
                  @@ [%equal: ECFR.money_tag] (get_id_tag arg)
                       (get_id_tag new_arg))
                args new_args
            with
            | Ok res -> res
            | Unequal_lengths -> false
          in
          let f_tag =
            lub_tags (lookup_var_tag f local_env) (Map expected_tag)
          in
          let new_f = update_id_tag f f_tag in
          let new_local_env = update_var_tag f f_tag local_env in
          let new_e_tag =
            match f_tag with Map t -> t | NoInfo -> NoInfo | _ -> Inconsistent
          in
          ( App (new_f, new_args),
            new_e_tag,
            param_env,
            new_local_env,
            ctr_tag_map,
            args_changes
            || (not @@ [%equal: ECFR.money_tag] f_tag (get_id_tag f)) )
      | Builtin (f, targs, args) ->
          let args_tags =
            List.map
              ~f:(fun arg -> lookup_var_tag2 arg local_env param_env)
              args
          in
          let res_tag, args_tags_usage =
            builtin_signature f expected_tag args_tags
          in
          let final_args, final_param_env, final_local_env, changes =
            let tags_list =
              match List.zip args args_tags_usage with
              | Unequal_lengths -> []
              | Ok res -> res
            in
            List.fold_right tags_list ~init:([], param_env, local_env, false)
              ~f:(fun
                   (arg, arg_tag)
                   (acc_args, acc_param_env, acc_local_env, acc_changes)
                 ->
                let new_local_env, new_param_env =
                  update_var_tag2 arg arg_tag acc_local_env acc_param_env
                in
                ( update_id_tag arg arg_tag :: acc_args,
                  new_param_env,
                  new_local_env,
                  acc_changes
                  || (not @@ [%equal: ECFR.money_tag] (get_id_tag arg) arg_tag)
                ))
          in
          let f_tag = lub_tags (Map res_tag) (Map expected_tag) in

          let op, (tag, r) = f in
          let new_f = (op, (f_tag, r)) in

          ( Builtin (new_f, targs, final_args),
            res_tag,
            final_param_env,
            final_local_env,
            ctr_tag_map,
            changes || (not @@ [%equal: ECFR.money_tag] f_tag tag) )
      | Let (i, topt, lhs, rhs) ->
          let ( ((_, (new_lhs_tag, _)) as new_lhs),
                lhs_param_env,
                lhs_local_env,
                lhs_ctr_tag_map,
                lhs_changes ) =
            cf_tag_expr lhs (get_id_tag i) param_env local_env ctr_tag_map
          in
          let updated_lhs_local_env =
            AssocDictionary.insert (CFIdentifier.as_string i) new_lhs_tag
              lhs_local_env
          in
          let ( ((_, (new_rhs_tag, _)) as new_rhs),
                rhs_param_env,
                rhs_local_env,
                rhs_ctr_tag_map,
                rhs_changes ) =
            cf_tag_expr rhs expected_tag lhs_param_env updated_lhs_local_env
              lhs_ctr_tag_map
          in
          let new_i_tag = lookup_var_tag i rhs_local_env in
          let new_i = update_id_tag i new_i_tag in
          let res_local_env =
            AssocDictionary.remove (CFIdentifier.as_string i) rhs_local_env
          in
          ( Let (new_i, topt, new_lhs, new_rhs),
            new_rhs_tag,
            rhs_param_env,
            res_local_env,
            rhs_ctr_tag_map,
            lhs_changes || rhs_changes
            || (not @@ [%equal: ECFR.money_tag] new_i_tag (get_id_tag i)) )
      | Constr (cname, ts, args) ->
          let new_args =
            List.map
              ~f:(fun arg ->
                update_id_tag arg (lookup_var_tag2 arg local_env param_env))
              args
          in
          let args_changes =
            match
              List.exists2
                ~f:(fun arg new_arg ->
                  not
                  @@ [%equal: ECFR.money_tag] (get_id_tag arg)
                       (get_id_tag new_arg))
                args new_args
            with
            | Ok res -> res
            | Unequal_lengths -> false
          in
          let new_ctr_tag_map, ctr_tag_map_changes =
            update_ctr_tag_map ctr_tag_map
              (CFIdentifier.as_string cname)
              (List.map new_args ~f:get_id_tag)
            |> Option.value_map ~default:(ctr_tag_map, false) ~f:(fun map ->
                   (map, true))
          in
          let tag =
            ctr_to_adt_tag
              (CFIdentifier.get_id cname)
              (List.map new_args ~f:get_id_tag)
          in
          ( Constr (cname, ts, new_args),
            tag,
            param_env,
            local_env,
            new_ctr_tag_map,
            ctr_tag_map_changes || args_changes )
      | MatchExpr (x, clauses) ->
          let ( res_clauses,
                res_tag,
                new_param_env,
                new_local_env,
                new_ctr_tag_map,
                res_clause_changes ) =
            List.fold_right clauses
              ~init:([], expected_tag, param_env, local_env, ctr_tag_map, false)
              ~f:(fun
                   (p, ep)
                   ( acc_clauses,
                     acc_res_tag,
                     acc_param_env,
                     acc_local_env,
                     acc_ctr_tag_map,
                     acc_changes )
                 ->
                let sub_local_env =
                  insert_pattern_vars_into_env p acc_local_env
                in
                let ( ((_, (new_e_tag, _)) as new_e),
                      new_param_env,
                      new_local_env,
                      e_ctr_tag_map,
                      new_changes ) =
                  cf_tag_expr ep expected_tag acc_param_env sub_local_env
                    acc_ctr_tag_map
                in
                let new_p, new_local_env, new_ctr_tag_map, p_changes =
                  update_pattern_vars_tags_from_usage_and_remove_from_env p
                    new_local_env e_ctr_tag_map
                in
                ( (new_p, new_e) :: acc_clauses,
                  lub_tags acc_res_tag new_e_tag,
                  new_param_env,
                  new_local_env,
                  new_ctr_tag_map,
                  acc_changes || new_changes || p_changes ))
          in
          let x_res_clauses, clause_changes =
            List.fold_right res_clauses ~init:([], res_clause_changes)
              ~f:(fun (p, e) (acc_ps, acc_changes) ->
                let new_p, p_change =
                  update_pattern_vars_tags_from_scrutinee p (get_id_tag x)
                in
                ((new_p, e) :: acc_ps, p_change || acc_changes))
          in
          let x_tag_usage =
            lub_pattern_tags (List.map ~f:(fun (p, _) -> p) x_res_clauses)
          in
          let new_x_tag =
            lub_tags (lookup_var_tag2 x local_env new_param_env) x_tag_usage
          in
          let new_x = update_id_tag x new_x_tag in
          let res_local_env, res_param_env =
            update_var_tag2 x new_x_tag new_local_env new_param_env
          in
          ( MatchExpr (new_x, x_res_clauses),
            res_tag,
            res_param_env,
            res_local_env,
            new_ctr_tag_map,
            clause_changes
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag x) new_x_tag) )
      | Fixpoint (_f, _t, _body) ->
          (* Library functions not handled. *)
          (e, Inconsistent, param_env, local_env, ctr_tag_map, false)
      | TFun (tvar, body) ->
          (* Ignore polymorphism. *)
          let ( ((_, (new_body_tag, _)) as new_body),
                new_param_env,
                new_local_env,
                new_ctr_tag_map,
                changes ) =
            cf_tag_expr body expected_tag param_env local_env ctr_tag_map
          in
          ( TFun (tvar, new_body),
            new_body_tag,
            new_param_env,
            new_local_env,
            new_ctr_tag_map,
            changes )
      | TApp (tf, arg_types) ->
          let new_tf_tag = lookup_var_tag2 tf local_env param_env in
          let new_tf = update_id_tag tf new_tf_tag in
          ( TApp (new_tf, arg_types),
            new_tf_tag,
            param_env,
            local_env,
            ctr_tag_map,
            false )
      | Message bs ->
          (* Find initializers and update env as appropriate *)
          let new_bs, new_param_env, new_local_env, changes =
            List.fold_right bs ~init:([], param_env, local_env, false)
              ~f:(fun (s, p) (acc_bs, acc_param_env, acc_local_env, acc_changes)
                 ->
                match p with
                | MLit _ ->
                    ((s, p) :: acc_bs, acc_param_env, acc_local_env, acc_changes)
                | MVar x ->
                    let usage_tag =
                      match s with
                      | x when String.(x = MessagePayload.amount_label) -> Money
                      | x
                        when String.(
                               x = MessagePayload.tag_label
                               || x = MessagePayload.recipient_label
                               || x = MessagePayload.eventname_label
                               || x = MessagePayload.exception_label) ->
                          NotMoney
                      | _ -> NoInfo
                    in
                    let old_env_tag =
                      lookup_var_tag2 x acc_local_env acc_param_env
                    in
                    let new_x_tag = lub_tags usage_tag old_env_tag in
                    let new_x = update_id_tag x new_x_tag in
                    let new_local_env, new_param_env =
                      update_var_tag2 x new_x_tag acc_local_env acc_param_env
                    in
                    ( (s, MVar new_x) :: acc_bs,
                      new_param_env,
                      new_local_env,
                      acc_changes
                      || not
                         @@ [%equal: ECFR.money_tag] (get_id_tag x) new_x_tag ))
          in
          ( Message new_bs,
            NotMoney,
            new_param_env,
            new_local_env,
            ctr_tag_map,
            changes )
      | GasExpr (g, e) ->
          let ( ((_, (new_e_tag, _)) as new_e),
                e_param_env,
                e_local_env,
                e_ctr_tag_map,
                e_changes ) =
            cf_tag_expr e expected_tag param_env local_env ctr_tag_map
          in
          ( GasExpr (g, new_e),
            new_e_tag,
            e_param_env,
            e_local_env,
            e_ctr_tag_map,
            e_changes )
    in
    let e_tag = lub new_e_tag in
    ( (new_e, (e_tag, rep)),
      new_param_env,
      new_local_env,
      new_ctr_tag_map,
      new_changes || (not @@ [%equal: ECFR.money_tag] tag e_tag) )

  (* Helper function for Load and Store - ensure field and local have same tag *)
  let cf_update_tag_for_field_assignment f x param_env field_env local_env =
    let x_tag = lookup_var_tag2 x local_env param_env in
    let f_tag = lookup_var_tag f field_env in
    let new_tag = lub_tags x_tag f_tag in
    let new_x = update_id_tag x new_tag in
    let new_f = update_id_tag f new_tag in
    let new_field_env = update_var_tag f new_tag field_env in
    let new_local_env, new_param_env =
      update_var_tag2 x new_tag local_env param_env
    in
    (new_f, new_x, new_param_env, new_field_env, new_local_env)

  let rec cf_tag_stmt (srep : CFSyntax.stmt_annot) param_env field_env local_env
      ctr_tag_map =
    let s, rep = srep in
    let ( new_s,
          new_param_env,
          new_field_env,
          new_local_env,
          new_ctr_tag_map,
          changes ) =
      match s with
      | Load (x, f) ->
          let new_f, new_x, new_param_env, new_field_env, tmp_local_env =
            cf_update_tag_for_field_assignment f x param_env field_env local_env
          in
          (* x is no longer in scope, so remove from local_env *)
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) tmp_local_env
          in
          ( Load (new_x, new_f),
            new_param_env,
            new_field_env,
            new_local_env,
            ctr_tag_map,
            (not @@ [%equal: ECFR.money_tag] (get_id_tag new_x) (get_id_tag x))
            || not
               @@ [%equal: ECFR.money_tag] (get_id_tag new_f) (get_id_tag f) )
      | RemoteLoad (x, adr, f) ->
          (* TODO - see Load case for inspiration *)
          (* x is no longer in scope, so remove from local_env *)
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          ( RemoteLoad (x, adr, f),
            param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            false )
      | Store (f, x) ->
          let new_f, new_x, new_param_env, new_field_env, new_local_env =
            cf_update_tag_for_field_assignment f x param_env field_env local_env
          in
          ( Store (new_f, new_x),
            new_param_env,
            new_field_env,
            new_local_env,
            ctr_tag_map,
            (not @@ [%equal: ECFR.money_tag] (get_id_tag new_x) (get_id_tag x))
            || not
               @@ [%equal: ECFR.money_tag] (get_id_tag new_f) (get_id_tag f) )
      | Bind (x, e) ->
          let x_tag = lookup_var_tag x local_env in
          let e_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          let ( ((_, (new_e_tag, _)) as new_e),
                new_param_env,
                new_local_env,
                new_ctr_tag_map,
                e_changes ) =
            cf_tag_expr e x_tag param_env e_local_env ctr_tag_map
          in
          let new_x_tag = lub_tags x_tag new_e_tag in
          let new_x = update_id_tag x new_x_tag in
          ( Bind (new_x, new_e),
            new_param_env,
            field_env,
            new_local_env,
            new_ctr_tag_map,
            e_changes
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag x) new_x_tag) )
      | MapUpdate (m, ks, v_opt) ->
          let v_tag =
            match v_opt with
            | None -> NoInfo
            | Some v -> lookup_var_tag2 v local_env param_env
          in
          let m_tag_usage =
            List.fold_left ks ~init:v_tag ~f:(fun acc _ -> Map acc)
          in
          let m_tag = lub_tags m_tag_usage (lookup_var_tag m field_env) in
          let new_m = update_id_tag m m_tag in
          let new_field_env = update_var_tag m m_tag field_env in
          let new_ks =
            List.map ks ~f:(fun k ->
                let new_k_tag =
                  lub_tags NotMoney (lookup_var_tag2 k local_env param_env)
                in
                update_id_tag k new_k_tag)
          in
          let ks_local_env, ks_param_env =
            update_vars_tags2 new_ks local_env param_env
          in
          let new_v_opt, new_param_env, new_local_env =
            match v_opt with
            | None -> (None, ks_param_env, ks_local_env)
            | Some v ->
                let v_tag_usage =
                  List.fold_left ks ~init:m_tag ~f:(fun acc_tag _ ->
                      match acc_tag with Map t -> t | _ -> Inconsistent)
                in
                let new_v_tag = lub_tags v_tag_usage v_tag in
                let new_v = update_id_tag v new_v_tag in
                let new_local_env, new_param_env =
                  update_var_tag2 v new_v_tag ks_local_env ks_param_env
                in
                (Some new_v, new_param_env, new_local_env)
          in
          ( MapUpdate (new_m, new_ks, new_v_opt),
            new_param_env,
            new_field_env,
            new_local_env,
            ctr_tag_map,
            (not @@ [%equal: ECFR.money_tag] (get_id_tag m) m_tag)
            || (not @@ [%equal: _] new_v_opt v_opt)
            || (not @@ [%equal: _] new_ks ks) )
      | MapGet (x, m, ks, fetch) ->
          let x_tag = lookup_var_tag x local_env in
          let val_tag =
            if fetch then
              match x_tag with
              | Adt (tname, [ t ]) when is_option_adt_name tname -> t
              | NoInfo -> NoInfo
              | _ -> Inconsistent
            else NoInfo
          in
          let m_tag_usage =
            List.fold_left ks ~init:val_tag ~f:(fun acc _ -> Map acc)
          in
          let m_tag = lub_tags m_tag_usage (lookup_var_tag m field_env) in
          let new_m = update_id_tag m m_tag in
          let new_field_env = update_var_tag m m_tag field_env in
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          let new_ks =
            List.map ks ~f:(fun k ->
                let new_k_tag =
                  lub_tags NotMoney (lookup_var_tag2 k new_local_env param_env)
                in
                update_id_tag k new_k_tag)
          in
          let ks_local_env, ks_param_env =
            update_vars_tags2 new_ks new_local_env param_env
          in
          let v_tag_usage =
            List.fold_left ks ~init:m_tag ~f:(fun acc_tag _ ->
                match acc_tag with Map t -> t | _ -> Inconsistent)
          in
          let new_val_tag = lub_tags val_tag v_tag_usage in
          let new_x_tag =
            if fetch then lub_tags x_tag (option_adt_tag new_val_tag)
            else NotMoney
            (* Bool *)
          in
          let new_x = update_id_tag x new_x_tag in
          ( MapGet (new_x, new_m, new_ks, fetch),
            ks_param_env,
            new_field_env,
            ks_local_env,
            ctr_tag_map,
            (not @@ [%equal: ECFR.money_tag] (get_id_tag x) new_x_tag)
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag m) m_tag)
            || (not @@ [%equal: _] new_ks ks) )
      | RemoteMapGet (x, adr, m, ks, fetch) ->
          (* TODO - see MapGet case for inspiration *)
          (* x is no longer in scope, so remove from local_env *)
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          ( RemoteMapGet (x, adr, m, ks, fetch),
            param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            false )
      | MatchStmt (x, clauses) ->
          let ( res_clauses,
                new_param_env,
                new_field_env,
                new_local_env,
                new_ctr_tag_map,
                res_clause_changes ) =
            List.fold_right clauses
              ~init:([], param_env, field_env, local_env, ctr_tag_map, false)
              ~f:(fun
                   (p, sp)
                   ( acc_clauses,
                     acc_param_env,
                     acc_field_env,
                     acc_local_env,
                     acc_ctr_tag_map,
                     acc_changes )
                 ->
                let sub_local_env =
                  insert_pattern_vars_into_env p acc_local_env
                in
                let ( new_stmts,
                      new_param_env,
                      new_field_env,
                      s_local_env,
                      s_ctr_tag_map,
                      s_changes ) =
                  cf_tag_stmts sp acc_param_env acc_field_env sub_local_env
                    acc_ctr_tag_map
                in
                let new_p, new_local_env, new_ctr_tag_map, p_changes =
                  update_pattern_vars_tags_from_usage_and_remove_from_env p
                    s_local_env s_ctr_tag_map
                in
                ( (new_p, new_stmts) :: acc_clauses,
                  new_param_env,
                  new_field_env,
                  new_local_env,
                  new_ctr_tag_map,
                  acc_changes || s_changes || p_changes ))
          in
          let x_res_clauses, clause_changes =
            List.fold_right res_clauses ~init:([], res_clause_changes)
              ~f:(fun (p, sp) (acc_ps, acc_changes) ->
                let new_p, p_change =
                  update_pattern_vars_tags_from_scrutinee p (get_id_tag x)
                in
                ((new_p, sp) :: acc_ps, p_change || acc_changes))
          in
          let x_tag_usage =
            lub_pattern_tags (List.map ~f:(fun (p, _) -> p) res_clauses)
          in
          let new_x_tag =
            lub_tags (lookup_var_tag2 x new_local_env new_param_env) x_tag_usage
          in
          let new_x = update_id_tag x new_x_tag in
          let res_local_env, res_param_env =
            update_var_tag2 x new_x_tag new_local_env new_param_env
          in
          ( MatchStmt (new_x, x_res_clauses),
            res_param_env,
            new_field_env,
            res_local_env,
            new_ctr_tag_map,
            clause_changes
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag x) new_x_tag) )
      | ReadFromBC (x, s) ->
          let x_tag = lub_tags NotMoney (lookup_var_tag x local_env) in
          let new_x = update_id_tag x x_tag in
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          ( ReadFromBC (new_x, s),
            param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] (get_id_tag x) x_tag )
      | TypeCast (x, r, t) ->
          let x_tag = lub_tags NotMoney (lookup_var_tag x local_env) in
          let new_x = update_id_tag x x_tag in
          let r_tag = lub_tags NotMoney (lookup_var_tag r local_env) in
          let new_r = update_id_tag r r_tag in
          let new_local_env =
            AssocDictionary.remove (CFIdentifier.as_string x) local_env
          in
          ( TypeCast (new_x, new_r, t),
            param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            not
              ([%equal: ECFR.money_tag] (get_id_tag x) x_tag
              || [%equal: ECFR.money_tag] (get_id_tag r) r_tag) )
      | AcceptPayment ->
          (AcceptPayment, param_env, field_env, local_env, ctr_tag_map, false)
      | Return i ->
          let i_tag = lub_tags NoInfo (lookup_var_tag2 i local_env param_env) in
          let new_i = update_id_tag i i_tag in
          let new_local_env, new_param_env =
            update_var_tag2 i i_tag local_env param_env
          in
          ( Return new_i,
            new_param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] (get_id_tag i) i_tag )
      | GasStmt g ->
          (GasStmt g, param_env, field_env, local_env, ctr_tag_map, false)
      | SendMsgs m ->
          let m_tag =
            lub_tags NotMoney (lookup_var_tag2 m local_env param_env)
          in
          let new_m = update_id_tag m m_tag in
          let new_local_env, new_param_env =
            update_var_tag2 m m_tag local_env param_env
          in
          ( SendMsgs new_m,
            new_param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] (get_id_tag m) m_tag )
      | CreateEvnt e ->
          let e_tag = lub_tags NotMoney (lookup_var_tag e local_env) in
          let new_e = update_id_tag e e_tag in
          let new_local_env = update_var_tag e e_tag local_env in
          ( CreateEvnt new_e,
            param_env,
            field_env,
            new_local_env,
            ctr_tag_map,
            not @@ [%equal: ECFR.money_tag] (get_id_tag e) e_tag )
      | CallProc (id_opt, p, args) ->
          (* TODO: Bindings from procedure calls are not taken into account in
                   the cash flow analysis. *)
          let new_args =
            List.map args ~f:(fun arg ->
                update_id_tag arg (lookup_var_tag2 arg local_env param_env))
          in
          let args_changes =
            match
              List.exists2
                ~f:(fun arg new_arg ->
                  not
                  @@ [%equal: ECFR.money_tag] (get_id_tag arg)
                       (get_id_tag new_arg))
                args new_args
            with
            | Ok res -> res
            | Unequal_lengths -> false
          in
          ( CallProc (id_opt, p, new_args),
            param_env,
            field_env,
            local_env,
            ctr_tag_map,
            args_changes )
      | Iterate (l, p) ->
          let l' = update_id_tag l (lookup_var_tag2 l local_env param_env) in
          let l_change =
            not ([%equal: ECFR.money_tag] (get_id_tag l) (get_id_tag l'))
          in
          ( Iterate (l', p),
            param_env,
            field_env,
            local_env,
            ctr_tag_map,
            l_change )
      | Throw xopt -> (
          match xopt with
          | Some x ->
              let x_tag = lub_tags NotMoney (lookup_var_tag x local_env) in
              let new_x = update_id_tag x x_tag in
              let new_local_env = update_var_tag x x_tag local_env in
              ( Throw (Some new_x),
                param_env,
                field_env,
                new_local_env,
                ctr_tag_map,
                not @@ [%equal: ECFR.money_tag] (get_id_tag x) x_tag )
          | None ->
              (Throw None, param_env, field_env, local_env, ctr_tag_map, false))
    in
    ( (new_s, rep),
      new_param_env,
      new_field_env,
      new_local_env,
      new_ctr_tag_map,
      changes )

  and cf_tag_stmts ss param_env field_env local_env ctr_tag_map =
    let init_local_env =
      List.fold_left ss ~init:local_env ~f:(fun acc_env srep ->
          let s, _ = srep in
          match s with
          | Load (x, _) | Bind (x, _) | MapGet (x, _, _, _) | ReadFromBC (x, _)
            ->
              AssocDictionary.insert (CFIdentifier.as_string x) (get_id_tag x)
                acc_env
          | _ -> acc_env)
    in
    List.fold_right ss
      ~init:([], param_env, field_env, init_local_env, ctr_tag_map, false)
      ~f:(fun
           s
           ( acc_ss,
             acc_param_env,
             acc_field_env,
             acc_local_env,
             acc_ctr_tag_map,
             acc_changes )
         ->
        let ( new_s,
              new_param_env,
              new_field_env,
              new_local_env,
              new_ctr_tag_map,
              new_changes ) =
          cf_tag_stmt s acc_param_env acc_field_env acc_local_env
            acc_ctr_tag_map
        in
        ( new_s :: acc_ss,
          new_param_env,
          new_field_env,
          new_local_env,
          new_ctr_tag_map,
          new_changes || acc_changes ))

  let cf_tag_component t param_env field_env ctr_tag_map =
    let { comp_type; comp_name; comp_params; comp_body; comp_return } = t in
    let empty_local_env = AssocDictionary.make_dict () in
    let implicit_local_env =
      AssocDictionary.insert MessagePayload.amount_label Money
        (AssocDictionary.insert MessagePayload.origin_label NotMoney
           (AssocDictionary.insert MessagePayload.sender_label NotMoney
              empty_local_env))
    in
    let init_local_env =
      List.fold_left comp_params ~init:implicit_local_env
        ~f:(fun acc_env (p, _) ->
          AssocDictionary.insert (CFIdentifier.as_string p) (get_id_tag p)
            acc_env)
    in
    let ( new_comp_body,
          new_param_env,
          new_field_env,
          new_local_env,
          new_ctr_tag_map,
          body_changes ) =
      cf_tag_stmts comp_body param_env field_env init_local_env ctr_tag_map
    in
    let new_params, new_changes =
      List.fold_right comp_params ~init:([], body_changes)
        ~f:(fun (p, typ) (acc_ps, acc_changes) ->
          let new_tag = lookup_var_tag p new_local_env in
          ( (update_id_tag p new_tag, typ) :: acc_ps,
            acc_changes
            || (not @@ [%equal: ECFR.money_tag] (get_id_tag p) new_tag) ))
    in
    ( {
        comp_type;
        comp_name;
        comp_params = new_params;
        comp_body = new_comp_body;
        comp_return;
      },
      new_param_env,
      new_field_env,
      new_ctr_tag_map,
      new_changes )

  let cf_tag_contract c =
    let { cname; cparams; cconstraint; cfields; ccomps } = c in
    let empty_env = AssocDictionary.make_dict () in
    let init_param_env =
      List.fold_left cparams ~init:empty_env ~f:(fun acc_env (p, _) ->
          AssocDictionary.insert (CFIdentifier.as_string p) (get_id_tag p)
            acc_env)
    in
    let implicit_field_env =
      AssocDictionary.insert (CFName.as_string balance_label) Money empty_env
    in
    let ctr_tag_map = init_ctr_tag_map () in
    let init_field_env =
      List.fold_left cfields ~init:implicit_field_env
        ~f:(fun acc_env (f, _, e) ->
          let (_, (e_tag, _)), _, _, _, _ =
            cf_tag_expr e
              (lub_tags (get_id_tag f) NoInfo)
              (AssocDictionary.make_dict ())
              (AssocDictionary.make_dict ())
              ctr_tag_map
          in
          AssocDictionary.insert (CFIdentifier.as_string f) e_tag acc_env)
    in
    let rec tagger cconstraint components param_env field_env ctr_tag_map =
      let new_constraint, tmp_param_env, _, tmp_ctr_tag_map, constraint_changes
          =
        cf_tag_expr cconstraint NotMoney param_env
          (AssocDictionary.make_dict ())
          ctr_tag_map
      in
      let new_ts, new_param_env, new_field_env, tmp_ctr_tag_map, changes =
        List.fold_right components
          ~init:
            ([], tmp_param_env, field_env, tmp_ctr_tag_map, constraint_changes)
          ~f:(fun
               t
               ( acc_ts,
                 acc_param_env,
                 acc_field_env,
                 acc_ctr_tag_map,
                 acc_changes )
             ->
            let new_t, new_param_env, new_field_env, new_ctr_tag_map, t_changes
                =
              cf_tag_component t acc_param_env acc_field_env acc_ctr_tag_map
            in
            ( new_t :: acc_ts,
              new_param_env,
              new_field_env,
              new_ctr_tag_map,
              acc_changes || t_changes ))
      in
      if changes then
        tagger new_constraint new_ts new_param_env new_field_env tmp_ctr_tag_map
      else
        (new_constraint, new_ts, new_param_env, new_field_env, tmp_ctr_tag_map)
    in
    let ( new_constraint,
          new_ccomps,
          new_param_env,
          new_field_env,
          final_ctr_tag_map ) =
      tagger cconstraint ccomps init_param_env init_field_env ctr_tag_map
    in
    let new_params =
      List.fold_right cparams ~init:[] ~f:(fun (p, t) acc_params ->
          let new_tag = lookup_var_tag p new_param_env in
          (update_id_tag p new_tag, t) :: acc_params)
    in
    let new_fields =
      List.fold_right cfields ~init:[] ~f:(fun (f, t, e) acc_fields ->
          let new_tag = lookup_var_tag f new_field_env in
          (update_id_tag f new_tag, t, e) :: acc_fields)
    in
    ( {
        cname;
        cparams = new_params;
        cconstraint = new_constraint;
        cfields = new_fields;
        ccomps = new_ccomps;
      },
      final_ctr_tag_map )

  let cf_tag_module m =
    let { smver; libs; elibs; contr } = m in
    let new_contr, ctr_tag_map = cf_tag_contract contr in
    ({ smver; libs; elibs; contr = new_contr }, ctr_tag_map)

  (*******************************************************)
  (*                Main entry function                  *)
  (*******************************************************)

  let main cmod token_fields =
    let init_mod = cf_init_tag_module cmod token_fields in
    let new_mod, ctr_tag_map = cf_tag_module init_mod in
    let param_field_tags =
      List.map
        ~f:(fun (p, _) -> (CFIdentifier.as_string p, get_id_tag p))
        new_mod.contr.cparams
      @ List.map
          ~f:(fun (f, _, _) -> (CFIdentifier.as_string f, get_id_tag f))
          new_mod.contr.cfields
    in
    let ctr_tags =
      let all_adts = DataTypeDictionary.get_all_adts () in
      List.filter_map all_adts ~f:(fun adt ->
          match
            List.filter_map adt.tconstr ~f:(fun ctr ->
                match
                  List.Assoc.find ctr_tag_map ~equal:String.( = )
                    (CFName.as_string ctr.cname)
                with
                | Some arg_tag_opts -> Some (ctr.cname, arg_tag_opts)
                | None -> None)
          with
          | [] ->
              (* No interesting constructors found *)
              None
          | x -> Some (adt.tname, x))
    in
    (param_field_tags, ctr_tags)
end
