From Velus Require Import Common.
From Velus Require Import Operators.
From Velus Require Import Clocks.
From Velus Require Import Lustre.LSyntax.

From Coq Require Import List.
Import List.ListNotations.
Open Scope list_scope.

From Coq Require Import Morphisms.

From Velus Require Import Environment.
From Velus Require Import Operators.

Lustre typing

Module Type LTYPING
       (Import Ids : IDS)
       (Import Op : OPERATORS)
       (Import Syn : LSYNTAX Ids Op).

  Inductive wt_clock : list (ident * type) -> clock -> Prop :=
  | wt_Cbase: forall vars,
      wt_clock vars Cbase
  | wt_Con: forall vars ck x b,
      In (x, bool_type) vars ->
      wt_clock vars ck ->
      wt_clock vars (Con ck x b).

  Inductive wt_nclock : list (ident * type) -> nclock -> Prop :=
  | wt_Cnamed: forall vars id ck,
      wt_clock vars ck ->
      wt_nclock vars (ck, id).

  Section WellTyped.

    Variable G : global.
    Variable vars : list (ident * type).

    Inductive wt_exp : exp -> Prop :=
    | wt_Econst: forall c,
        wt_exp (Econst c)

    | wt_Evar: forall x ty nck,
        In (x, ty) vars ->
        wt_nclock vars nck ->
        wt_exp (Evar x (ty, nck))

    | wt_Eunop: forall op e tye ty nck,
        wt_exp e ->
        typeof e = [tye] ->
        type_unop op tye = Some ty ->
        wt_nclock vars nck ->
        wt_exp (Eunop op e (ty, nck))

    | wt_Ebinop: forall op e1 e2 ty1 ty2 ty nck,
        wt_exp e1 ->
        wt_exp e2 ->
        typeof e1 = [ty1] ->
        typeof e2 = [ty2] ->
        type_binop op ty1 ty2 = Some ty ->
        wt_nclock vars nck ->
        wt_exp (Ebinop op e1 e2 (ty, nck))

    | wt_Efby: forall e0s es anns,
        Forall wt_exp e0s ->
        Forall wt_exp es ->
        typesof es = map fst anns ->
        typesof e0s = map fst anns ->
        Forall (wt_nclock vars) (map snd anns) ->
        wt_exp (Efby e0s es anns)

    | wt_Earrow: forall e0s es anns,
        Forall wt_exp e0s ->
        Forall wt_exp es ->
        typesof es = map fst anns ->
        typesof e0s = map fst anns ->
        Forall (wt_nclock vars) (map snd anns) ->
        wt_exp (Earrow e0s es anns)

    | wt_Ewhen: forall es x b tys nck,
        Forall wt_exp es ->
        typesof es = tys ->
        In (x, bool_type) vars ->
        wt_nclock vars nck ->
        wt_exp (Ewhen es x b (tys, nck))

    | wt_Emerge: forall x ets efs tys nck,
        Forall wt_exp ets ->
        Forall wt_exp efs ->
        In (x, bool_type) vars ->
        typesof ets = tys ->
        typesof efs = tys ->
        wt_nclock vars nck ->
        wt_exp (Emerge x ets efs (tys, nck))

    | wt_Eifte: forall e ets efs tys nck,
        wt_exp e ->
        Forall wt_exp ets ->
        Forall wt_exp efs ->
        typeof e = [bool_type] ->
        typesof ets = tys ->
        typesof efs = tys ->
        wt_nclock vars nck ->
        wt_exp (Eite e ets efs (tys, nck))

    | wt_Eapp: forall f es anns n,
        Forall wt_exp es ->
        find_node f G = Some n ->
        Forall2 (fun et '(_, (t, _)) => et = t) (typesof es) n.(n_in) ->
        Forall2 (fun a '(_, (t, _)) => fst a = t) anns n.(n_out) ->
        Forall (fun a => wt_nclock (vars++(idty (fresh_ins es++anon_streams anns))) (snd a)) anns ->
        wt_exp (Eapp f es None anns)

    | wt_EappReset: forall f es r anns n,
        Forall wt_exp es ->
        find_node f G = Some n ->
        Forall2 (fun et '(_, (t, _)) => et = t) (typesof es) n.(n_in) ->
        Forall2 (fun a '(_, (t, _)) => fst a = t) anns n.(n_out) ->
        Forall (fun a => wt_nclock (vars++(idty (fresh_ins es++anon_streams anns))) (snd a)) anns ->
        wt_exp r ->
        typeof r = [bool_type] ->
        wt_exp (Eapp f es (Some r) anns).

    Section wt_exp_ind2.

      Variable P : exp -> Prop.

      Hypothesis EconstCase:
        forall c : const,
          P (Econst c).

      Hypothesis EvarCase:
        forall x ty nck,
          In (x, ty) vars ->
          wt_nclock vars nck ->
          P (Evar x (ty, nck)).

      Hypothesis EunopCase:
        forall op e tye ty nck,
          wt_exp e ->
          P e ->
          typeof e = [tye] ->
          type_unop op tye = Some ty ->
          wt_nclock vars nck ->
          P (Eunop op e (ty, nck)).

      Hypothesis EbinopCase:
        forall op e1 e2 ty1 ty2 ty nck,
          wt_exp e1 ->
          P e1 ->
          wt_exp e2 ->
          P e2 ->
          typeof e1 = [ty1] ->
          typeof e2 = [ty2] ->
          type_binop op ty1 ty2 = Some ty ->
          wt_nclock vars nck ->
          P (Ebinop op e1 e2 (ty, nck)).

      Hypothesis EfbyCase:
        forall e0s es anns,
          Forall wt_exp e0s ->
          Forall wt_exp es ->
          Forall P e0s ->
          Forall P es ->
          typesof es = map fst anns ->
          typesof e0s = map fst anns ->
          Forall (wt_nclock vars) (map snd anns) ->
          P (Efby e0s es anns).

      Hypothesis EarrowCase:
        forall e0s es anns,
          Forall wt_exp e0s ->
          Forall wt_exp es ->
          Forall P e0s ->
          Forall P es ->
          typesof es = map fst anns ->
          typesof e0s = map fst anns ->
          Forall (wt_nclock vars) (map snd anns) ->
          P (Earrow e0s es anns).

      Hypothesis EwhenCase:
        forall es x b tys nck,
          Forall wt_exp es ->
          Forall P es ->
          typesof es = tys ->
          In (x, bool_type) vars ->
          wt_nclock vars nck ->
          P (Ewhen es x b (tys, nck)).

      Hypothesis EmergeCase:
        forall x ets efs tys nck,
          Forall wt_exp ets ->
          Forall P ets ->
          Forall wt_exp efs ->
          Forall P efs ->
          In (x, bool_type) vars ->
          typesof ets = tys ->
          typesof efs = tys ->
          wt_nclock vars nck ->
          P (Emerge x ets efs (tys, nck)).

      Hypothesis EiteCase:
        forall e ets efs tys nck,
          wt_exp e ->
          P e ->
          Forall wt_exp ets ->
          Forall P ets ->
          Forall wt_exp efs ->
          Forall P efs ->
          typeof e = [bool_type] ->
          typesof ets = tys ->
          typesof efs = tys ->
          wt_nclock vars nck ->
          P (Eite e ets efs (tys, nck)).

      Hypothesis EappCase:
        forall f es anns n,
          Forall wt_exp es ->
          Forall P es ->
          find_node f G = Some n ->
          Forall2 (fun et '(_, (t, _)) => et = t) (typesof es) n.(n_in) ->
          Forall2 (fun a '(_, (t, _)) => fst a = t) anns n.(n_out) ->
          Forall (fun a => wt_nclock (vars++(idty (fresh_ins es++anon_streams anns))) (snd a)) anns ->
          P (Eapp f es None anns).

      Hypothesis EappResetCase:
        forall f es r anns n,
          Forall wt_exp es ->
          Forall P es ->
          find_node f G = Some n ->
          Forall2 (fun et '(_, (t, _)) => et = t) (typesof es) n.(n_in) ->
          Forall2 (fun a '(_, (t, _)) => fst a = t) anns n.(n_out) ->
          Forall (fun a => wt_nclock (vars++(idty (fresh_ins es++anon_streams anns))) (snd a)) anns ->
          wt_exp r ->
          P r ->
          typeof r = [bool_type] ->
          P (Eapp f es (Some r) anns).

      Fixpoint wt_exp_ind2 (e: exp) (H: wt_exp e) {struct H} : P e.

    End wt_exp_ind2.

    Definition wt_equation (eq : equation) : Prop :=
      let (xs, es) := eq in
      Forall wt_exp es
      /\ Forall2 (fun x ty=> In (x, ty) vars) xs (typesof es).

  End WellTyped.

  Definition wt_clocks (vars: list (ident * (type * clock)))
                           : list (ident * (type * clock)) -> Prop :=
    Forall (fun '(_, (_, ck)) => wt_clock (idty vars) ck).

  Definition wt_node (G: global) (n: node) : Prop
    := wt_clocks n.(n_in) n.(n_in)
       /\ wt_clocks (n.(n_in) ++ n.(n_out)) n.(n_out)
       /\ wt_clocks (n.(n_in) ++ n.(n_out) ++ n.(n_vars)) n.(n_vars)
       /\ Forall (wt_equation G (idty (n.(n_in) ++ n.(n_vars) ++ n.(n_out)))) n.(n_eqs).

  Inductive wt_global : global -> Prop :=
  | wtg_nil:
      wt_global []
  | wtg_cons: forall n ns,
      wt_global ns ->
      wt_node ns n ->
      Forall (fun n'=> n.(n_name) <> n'.(n_name) :> ident) ns ->
      wt_global (n::ns).

  Hint Constructors wt_clock wt_nclock wt_exp wt_global : ltyping.
  Hint Unfold wt_equation : ltyping.

  Lemma wt_global_NoDup:
    forall g,
      wt_global g ->
      NoDup (map n_name g).

  Lemma wt_global_app:
    forall G G',
      wt_global (G' ++ G) ->
      wt_global G.

  Lemma wt_find_node:
    forall G f n,
      wt_global G ->
      find_node f G = Some n ->
      exists G', wt_node G' n.

  Lemma wt_clock_add:
    forall x v env ck,
      ~InMembers x env ->
      wt_clock env ck ->
      wt_clock ((x, v) :: env) ck.

  Instance wt_clock_Proper:
    Proper (@Permutation.Permutation (ident * type) ==> @eq clock ==> iff)
           wt_clock.

  Instance wt_nclock_Proper:
    Proper (@Permutation.Permutation (ident * type) ==> @eq nclock ==> iff)
           wt_nclock.

  Instance wt_nclock_pointwise_Proper:
    Proper (@Permutation.Permutation (ident * type)
                          ==> pointwise_relation nclock iff) wt_nclock.

  Instance wt_exp_Proper:
    Proper (@eq global ==> @Permutation.Permutation (ident * type)
                ==> @eq exp ==> iff)
           wt_exp.

  Instance wt_exp_pointwise_Proper:
    Proper (@eq global ==> @Permutation.Permutation (ident * type)
                                     ==> pointwise_relation exp iff) wt_exp.

  Instance wt_equation_Proper:
    Proper (@eq global ==> @Permutation.Permutation (ident * type)
                ==> @eq equation ==> iff)
           wt_equation.


  Section incl.

    Fact wt_clock_incl : forall vars vars' cl,
      incl vars vars' ->
      wt_clock vars cl ->
      wt_clock vars' cl.
    Local Hint Resolve wt_clock_incl.

    Fact wt_nclock_incl : forall vars vars' cl,
        incl vars vars' ->
        wt_nclock vars cl ->
        wt_nclock vars' cl.
    Local Hint Resolve wt_nclock_incl.

    Lemma wt_exp_incl : forall G vars vars' e,
        incl vars vars' ->
        wt_exp G vars e ->
        wt_exp G vars' e.

    Lemma wt_equation_incl : forall G vars vars' eq,
        incl vars vars' ->
        wt_equation G vars eq ->
        wt_equation G vars' eq.

  End incl.


  Section global_incl.
    Fact wt_exp_global_incl : forall G G' vars e,
      incl G G' ->
      NoDup (map n_name G) ->
      NoDup (map n_name G') ->
      wt_exp G vars e ->
      wt_exp G' vars e.

    Fact wt_equation_global_incl : forall G G' vars e,
      incl G G' ->
      NoDup (map n_name G) ->
      NoDup (map n_name G') ->
      wt_equation G vars e ->
      wt_equation G' vars e.

    Fact wt_node_global_incl : forall G G' e,
      incl G G' ->
      NoDup (map n_name G) ->
      NoDup (map n_name G') ->
      wt_node G e ->
      wt_node G' e.

    Lemma wt_global_Forall : forall G,
        wt_global G ->
        Forall (wt_node G) G.
  End global_incl.

  Local Hint Resolve wt_clock_incl incl_appl incl_refl.
  Lemma wt_exp_clockof:
    forall G env e,
      wt_exp G env e ->
      Forall (wt_clock (env++idty (fresh_in e))) (clockof e).

Validation

  Module MyOpAux := OperatorsAux Op.

  Fixpoint check_clock (venv : Env.t type) (ck : clock) : bool :=
    match ck with
    | Cbase => true
    | Con ck' x b =>
      match Env.find x venv with
      | None => false
      | Some xt => (xt ==b bool_type) && check_clock venv ck'
      end
    end.

  Definition check_nclock (venv : Env.t type) (nck : nclock) : bool :=
    check_clock venv (fst nck).

  Lemma check_clock_correct:
    forall venv ck,
      check_clock venv ck = true ->
      wt_clock (Env.elements venv) ck.

  Lemma check_nclock_correct:
    forall venv nck,
      check_nclock venv nck = true ->
      wt_nclock (Env.elements venv) nck.
  Local Hint Resolve check_nclock_correct.

  Section ValidateExpression.

    Variable G : global.
    Variable venv : Env.t type.

    Open Scope option_monad_scope.

    Definition check_paired_types2 (t1 : type) (tc : ann) : bool :=
      let '(t, c) := tc in
      (t1 ==b t) && (check_nclock venv c).

    Definition check_paired_types3 (t1 t2 : type) (tc : ann) : bool :=
      let '(t, c) := tc in
      (t1 ==b t) && (t2 ==b t) && (check_nclock venv c).

    Definition check_reset (rt : option (option (list type))) : bool :=
      match rt with
      | None => true
      | Some (Some [ty]) => ty ==b bool_type
      | _ => false
      end.

    Function check_var (x : ident) (ty : type) : bool :=
      match Env.find x venv with
      | None => false
      | Some xt => ty ==b xt
      end.

    Fixpoint check_exp (e : exp) : option (list type) :=
      match e with
      | Econst c => Some ([type_const c])

      | Evar x (xt, nck) =>
        if check_var x xt && check_nclock venv nck then Some [xt] else None

      | Eunop op e (xt, nck) =>
        do te <- assert_singleton (check_exp e);
        do t <- type_unop op te;
        if (xt ==b t) && check_nclock venv nck then Some [xt] else None

      | Ebinop op e1 e2 (xt, nck) =>
        do te1 <- assert_singleton (check_exp e1);
        do te2 <- assert_singleton (check_exp e2);
        do t <- type_binop op te1 te2;
        if (xt ==b t) && check_nclock venv nck then Some [xt] else None

      | Efby e0s es anns =>
        do t0s <- oconcat (map check_exp e0s);
        do ts <- oconcat (map check_exp es);
        if forall3b check_paired_types3 t0s ts anns
        then Some (map fst anns) else None

      | Earrow e0s es anns =>
        do t0s <- oconcat (map check_exp e0s);
        do ts <- oconcat (map check_exp es);
        if forall3b check_paired_types3 t0s ts anns
        then Some (map fst anns) else None

      | Ewhen es x b (tys, nck) =>
        do ts <- oconcat (map check_exp es);
        if check_var x bool_type && (forall2b equiv_decb ts tys) && (check_nclock venv nck)
        then Some tys else None

      | Emerge x e1s e2s (tys, nck) =>
        do t1s <- oconcat (map check_exp e1s);
        do t2s <- oconcat (map check_exp e2s);
        if check_var x bool_type
             && (forall3b equiv_decb3 t1s t2s tys)
             && (check_nclock venv nck)
        then Some tys else None

      | Eite e e1s e2s (tys, nck) =>
        do t1s <- oconcat (map check_exp e1s);
        do t2s <- oconcat (map check_exp e2s);
        do xt <- assert_singleton (check_exp e);
        if (xt ==b bool_type)
             && (forall3b equiv_decb3 t1s t2s tys)
             && (check_nclock venv nck)
        then Some tys else None

      | Eapp f es ro anns =>
        do n <- find_node f G;
        do ts <- oconcat (map check_exp es);
        if (forall2b (fun et '(_, (t, _)) => et ==b t) ts n.(n_in))
             && (forall2b (fun '(ot, oc) '(_, (t, _)) =>
                             check_nclock (Env.adds' (idty (fresh_ins es++anon_streams anns)) venv) oc
                             && (ot ==b t)) anns n.(n_out))
             && check_reset (option_map check_exp ro)
        then Some (map fst anns)
        else None
      end.

    Definition check_rhs (e : exp) : option (list type) :=
      match e with
      | Eapp f es ro anns =>
        do n <- find_node f G;
        do ts <- oconcat (map check_exp es);
        if (forall2b (fun et '(_, (t, _)) => et ==b t) ts n.(n_in))
             && (forall2b (fun '(ot, oc) '(_, (t, _)) =>
                             check_nclock (Env.adds' (idty (fresh_ins es)) venv) oc
                             && (ot ==b t)) anns n.(n_out))
             && check_reset (option_map check_exp ro)
        then Some (map fst anns)
        else None
      | e => check_exp e
      end.

    Definition check_equation (eq : equation) : bool :=
      let '(xs, es) := eq in
      match oconcat (map check_rhs es) with
      | None => false
      | Some tys => forall2b check_var xs tys
      end.

    Lemma check_var_correct:
      forall x ty,
        check_var x ty = true <-> In (x, ty) (Env.elements venv).
    Local Hint Resolve check_var_correct.

    Lemma check_paired_types2_correct:
      forall tys1 anns,
        forall2b check_paired_types2 tys1 anns = true ->
        tys1 = map fst anns
        /\ Forall (wt_nclock (Env.elements venv)) (map snd anns).

    Lemma check_paired_types3_correct:
      forall tys1 tys2 anns,
        forall3b check_paired_types3 tys1 tys2 anns = true ->
        tys1 = map fst anns
        /\ tys2 = map fst anns
        /\ Forall (wt_nclock (Env.elements venv)) (map snd anns).

    Lemma oconcat_map_check_exp':
      forall {f} es tys,
        (forall e tys,
            In e es ->
            NoDupMembers (Env.elements venv++idty (fresh_in e)) ->
            f e = Some tys ->
            wt_exp G (Env.elements venv) e /\ typeof e = tys) ->
        NoDupMembers (Env.elements venv++idty (fresh_ins es)) ->
        oconcat (map f es) = Some tys ->
        Forall (wt_exp G (Env.elements venv)) es
        /\ typesof es = tys.

    Import Permutation.
    Local Hint Constructors wt_exp.
    Lemma check_exp_correct:
      forall e tys,
        NoDupMembers (Env.elements venv++(idty (fresh_in e))) ->
        check_exp e = Some tys ->
        wt_exp G (Env.elements venv) e
        /\ typeof e = tys.

    Corollary oconcat_map_check_exp:
      forall es tys,
        NoDupMembers (Env.elements venv++idty (fresh_ins es)) ->
        oconcat (map check_exp es) = Some tys ->
        Forall (wt_exp G (Env.elements venv)) es
        /\ typesof es = tys.

    Lemma check_rhs_correct:
      forall e tys,
        NoDupMembers (Env.elements venv++(idty (anon_in e))) ->
        check_rhs e = Some tys ->
        wt_exp G (Env.elements venv) e
        /\ typeof e = tys.

    Corollary oconcat_map_check_rhs:
      forall es tys,
        NoDupMembers (Env.elements venv++idty (anon_ins es)) ->
        oconcat (map check_rhs es) = Some tys ->
        Forall (wt_exp G (Env.elements venv)) es
        /\ typesof es = tys.

    Lemma check_equation_correct:
      forall eq,
        NoDupMembers (Env.elements venv++(idty (anon_in_eq eq))) ->
        check_equation eq = true ->
        wt_equation G (Env.elements venv) eq.

  End ValidateExpression.

  Section ValidateGlobal.
    Definition check_node (G : global) (n : node) :=
      forallb (check_clock (Env.from_list (idty (n_in n)))) (List.map (fun '(_, (_, ck)) => ck) (n_in n)) &&
      forallb (check_clock (Env.from_list (idty (n_in n ++ n_out n)))) (List.map (fun '(_, (_, ck)) => ck) (n_out n)) &&
      forallb (check_clock (Env.from_list (idty (n_in n ++ n_out n ++ n_vars n)))) (List.map (fun '(_, (_, ck)) => ck) (n_vars n)) &&
      forallb (check_equation G (Env.from_list (idty (n_in n ++ n_vars n ++ n_out n)))) (n_eqs n).

    Definition check_global (G : global) :=
      check_nodup (List.map n_name G) &&
      (fix aux G := match G with
                    | [] => true
                    | hd::tl => check_node tl hd && aux tl
                    end) G.

    Lemma check_node_correct : forall G n,
        check_node G n = true ->
        wt_node G n.

    Lemma check_global_correct : forall G,
        check_global G = true ->
        wt_global G.
  End ValidateGlobal.

  Section interface_eq.

    Hint Constructors wt_exp.
    Fact iface_eq_wt_exp : forall G G' vars e,
        global_iface_eq G G' ->
        wt_exp G vars e ->
        wt_exp G' vars e.

    Fact iface_eq_wt_equation : forall G G' vars equ,
        global_iface_eq G G' ->
        wt_equation G vars equ ->
        wt_equation G' vars equ.

    Lemma iface_eq_wt_node : forall G G' n,
        global_iface_eq G G' ->
        wt_node G n ->
        wt_node G' n.

  End interface_eq.

wt implies wl

  Hint Constructors wl_exp.
  Fact wt_exp_wl_exp : forall G vars e,
      wt_exp G vars e ->
      wl_exp G e.
  Hint Resolve wt_exp_wl_exp.

  Corollary Forall_wt_exp_wl_exp : forall G vars es,
      Forall (wt_exp G vars) es ->
      Forall (wl_exp G) es.
  Hint Resolve Forall_wt_exp_wl_exp.

  Fact wt_equation_wl_equation : forall G vars equ,
      wt_equation G vars equ ->
      wl_equation G equ.
  Hint Resolve wt_equation_wl_equation.

  Fact wt_node_wl_node : forall G n,
      wt_node G n ->
      wl_node G n.
  Hint Resolve wt_node_wl_node.

  Fact wt_global_wl_global : forall G,
      wt_global G ->
      wl_global G.
  Hint Resolve wt_global_wl_global.


  Lemma wt_clock_Is_free_in : forall x env ck,
      wt_clock env ck ->
      Is_free_in_clock x ck ->
      InMembers x env.

  Corollary wt_nclock_Is_free_in : forall x env name ck,
      wt_nclock env (ck, name) ->
      Is_free_in_clock x ck ->
      InMembers x env.

End LTYPING.

Module LTypingFun
       (Ids : IDS)
       (Op : OPERATORS)
       (Syn : LSYNTAX Ids Op)
       <: LTYPING Ids Op Syn.
  Include LTYPING Ids Op Syn.
End LTypingFun.