From Coq Require Import Ascii.
From mininix Require Export res nix.operational_props.
From stdpp Require Import options.

Section val.
  Local Unset Elimination Schemes.
  Inductive val :=
    | VLit (bl : base_lit) (Hbl : base_lit_ok bl)
    | VClo (x : string) (E : gmap string (kind * thunk)) (e : expr)
    | VCloMatch (E : gmap string (kind * thunk))
                (ms : gmap string (option expr))
                (strict : bool) (e : expr)
    | VList (ts : list thunk)
    | VAttr (ts : gmap string thunk)
  with thunk :=
    | Forced (v : val) : thunk
    | Thunk (E : gmap string (kind * thunk)) (e : expr) : thunk
    | Indirect (x : string)
               (E : gmap string (kind * thunk))
               (tαs : gmap string (expr + thunk)).
End val.

Notation VLitI bl := (VLit bl I).

Notation tattr := (expr + thunk)%type.
Notation env := (gmap string (kind * thunk)).

Definition maybe_VLit (v : val) : option base_lit :=
  if v is VLit bl _ then Some bl else None.
Global Instance maybe_VList : Maybe VList := λ v,
  if v is VList ts then Some ts else None.
Global Instance maybe_VAttr : Maybe VAttr := λ v,
  if v is VAttr ts then Some ts else None.

Fixpoint interp_eq_list_body (i : nat) (ts1 ts2 : list thunk) : expr :=
  match ts1, ts2 with
  | [], [] => ELit (LitBool true)
  | _ :: ts1, _ :: ts2 =>
      EIf (EBinOp EqOp (EId' ("1" +:+ pretty i)) (EId' ("2" +:+ pretty i)))
          (interp_eq_list_body (S i) ts1 ts2) (ELit (LitBool false))
  | _, _ => ELit (LitBool false)
  end.

Definition interp_eq_list (ts1 ts2 : list thunk) : thunk :=
  Thunk (kmap (λ n : nat, String "1" (pretty n)) ((ABS,.) <$> map_seq 0 ts1) ∪
         kmap (λ n : nat, String "2" (pretty n)) ((ABS,.) <$> map_seq 0 ts2)) $
    interp_eq_list_body 0 ts1 ts2.

Fixpoint interp_lt_list_body (i : nat) (ts1 ts2 : list thunk) : expr :=
  match ts1, ts2 with
  | [], _ => ELit (LitBool true)
  | _ :: ts1, _ :: ts2 =>
      EIf (EBinOp LtOp (EId' ("1" +:+ pretty i)) (EId' ("2" +:+ pretty i)))
          (ELit (LitBool true))
          (EIf (EBinOp EqOp (EId' ("1" +:+ pretty i)) (EId' ("2" +:+ pretty i)))
            (interp_lt_list_body (S i) ts1 ts2) (ELit (LitBool false)))
  | _ :: _, [] => ELit (LitBool false)
  end.

Definition interp_lt_list (ts1 ts2 : list thunk) : thunk :=
  Thunk (kmap (λ n : nat, String "1" (pretty n)) ((ABS,.) <$> map_seq 0 ts1) ∪
         kmap (λ n : nat, String "2" (pretty n)) ((ABS,.) <$> map_seq 0 ts2)) $
    interp_lt_list_body 0 ts1 ts2.

Definition interp_eq_attr (ts1 ts2 : gmap string thunk) : thunk :=
  Thunk (kmap (String "1") ((ABS,.) <$> ts1) ∪
         kmap (String "2") ((ABS,.) <$> ts2)) $
    sem_and_attr $ map_imap (λ x _,
      Some (EBinOp EqOp (EId' ("1" +:+ x)) (EId' ("2" +:+ x)))) ts1.

Definition interp_eq (v1 v2 : val) : option thunk :=
  match v1, v2 with
  | VLit bl1 _, VLit bl2 _ =>
      Some $ Forced $ VLitI (LitBool $ sem_eq_base_lit bl1 bl2)
  | VClo _ _ _, VClo _ _ _ => None
  | VList ts1, VList ts2 => Some $
      if decide (length ts1 = length ts2) then interp_eq_list ts1 ts2
      else Forced $ VLitI (LitBool false)
  | VAttr ts1, VAttr ts2 => Some $
      if decide (dom ts1 = dom ts2) then interp_eq_attr ts1 ts2
      else Forced $ VLitI (LitBool false)
  | _, _ => Some $ Forced $ VLitI (LitBool false)
  end.

Definition type_of_val (v : val) : string :=
  match v with
  | VLit bl _ => type_of_base_lit bl
  | VClo _ _ _ | VCloMatch _ _ _ _ => "lambda"
  | VList _ => "list"
  | VAttr _ => "set"
  end.

Global Instance val_inhabited : Inhabited val := populate (VLitI inhabitant).
Global Instance thunk_inhabited : Inhabited thunk := populate (Forced inhabitant).

Definition interp_bin_op (op : bin_op) (v1 : val) : option (val → option thunk) :=
  if decide (op = EqOp) then
    Some (interp_eq v1)
  else if decide (op = TypeOfOp) then
    Some $ λ v2,
      guard (maybe_VLit v2 = Some LitNull);;
      Some $ Forced $ VLitI (LitString $ type_of_val v1)
  else
    match v1 with
    | VLit (LitNum n1) Hn1 =>
        if maybe RoundOp op is Some m then
          Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            Some $ Forced $ VLit
              (LitNum $ NInt $ float_to_bounded_Z $ Float.round m $ num_to_float n1)
              (float_to_bounded_Z_ok _)
        else
          '(f ↾ Hf) ← option_to_eq_Some (sem_bin_op_num op n1);
          Some $ λ v2,
            if v2 is VLit (LitNum n2) Hn2 then
              '(bl ↾ Hbl) ← option_to_eq_Some (f n2);
              Some $ Forced $ VLit bl (sem_bin_op_num_ok Hn1 Hn2 Hf Hbl)
            else None
    | VLit (LitString s1) _ =>
        match op with
        | SingletonAttrOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            Some $ Forced $ VClo "t" ∅ (EAttr {[ s1 := AttrN (EId' "t") ]})
        | MatchStringOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            match s1 with
            | EmptyString => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool true));
                "head" := Forced (VLitI LitNull);
                "tail" := Forced (VLitI LitNull) ]}
            | String a s1 => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool false));
                "head" := Forced (VLitI (LitString (String a EmptyString)));
                "tail" := Forced (VLitI (LitString s1)) ]}
            end
        | _ =>
            '(f ↾ Hf) ← option_to_eq_Some (sem_bin_op_string op);
            Some $ λ v2,
              bl2 ← maybe_VLit v2;
              s2 ← maybe LitString bl2;
              Some $ Forced $ VLit (f s1 s2) (sem_bin_op_string_ok Hf)
        end
    | VClo _ _ _ =>
       match op with
       | FunctionArgsOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            Some (Forced (VAttr ∅))
        | _ => None
        end
    | VCloMatch _ ms _ _ =>
        match op with
        | FunctionArgsOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            Some $ Forced $ VAttr $
              (λ m, Forced $ VLitI (LitBool (from_option (λ _, true) false m))) <$> ms
        | _ => None
        end
    | VList ts1 =>
        match op with
        | LtOp => Some $ λ v2,
            ts2 ← maybe VList v2;
            Some (interp_lt_list ts1 ts2)
        | MatchListOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            match ts1 with
            | [] => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool true));
                "head" := Forced (VLitI LitNull);
                "tail" := Forced (VLitI LitNull) ]}
            | t :: ts1 => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool false));
                "head" := t;
                "tail" := Forced (VList ts1) ]}
            end
        | AppendListOp => Some $ λ v2,
            ts2 ← maybe VList v2;
            Some (Forced (VList (ts1 ++ ts2)))
        | _ => None
        end
    | VAttr ts1 =>
        match op with
        | SelectAttrOp => Some $ λ v2,
            bl ← maybe_VLit v2;
            x ← maybe LitString bl;
            ts1 !! x
        | UpdateAttrOp => Some $ λ v2,
            ts2 ← maybe VAttr v2;
            Some $ Forced $ VAttr $ ts2 ∪ ts1
        | HasAttrOp => Some $ λ v2,
            bl ← maybe_VLit v2;
            x ← maybe LitString bl;
            Some $ Forced $ VLitI (LitBool $ bool_decide (is_Some (ts1 !! x)))
        | DeleteAttrOp => Some $ λ v2,
            bl ← maybe_VLit v2;
            x ← maybe LitString bl;
            Some $ Forced $ VAttr $ delete x ts1
        | MatchAttrOp => Some $ λ v2,
            guard (maybe_VLit v2 = Some LitNull);;
            match map_minimal_key attr_le ts1 with
            | None => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool true));
                "key" := Forced (VLitI LitNull);
                "head" := Forced (VLitI LitNull);
                "tail" := Forced (VLitI LitNull) ]}
            | Some x => Some $ Forced $ VAttr {[
                "empty" := Forced (VLitI (LitBool false));
                "key" := Forced (VLitI (LitString x));
                "head" := ts1 !!! x;
                "tail" := Forced (VAttr (delete x ts1)) ]}
            end
        | _ => None
        end
    | _ => None
    end.

Definition interp_match
    (ts : gmap string thunk) (mds : gmap string (option expr))
    (strict : bool) : option (gmap string tattr) :=
  map_mapM id $ merge (λ mt mmd,
    (* Some (Some _) means keep, Some None means fail, None means skip *)
    match mt, mmd with
    | Some t, Some _ => Some $ Some (inr t)
    | None, Some (Some e) => Some $ Some (inl e)
    | None, Some _ => Some None (* bad *)
    | Some _, None => guard strict;; Some None
    | _, _ => None (* skip *)
    end) ts mds.

Definition force_deep1
    (force_deep : val → res val)
    (interp_thunk : thunk → res val) (v : val) : res val :=
  match v with
  | VList ts => VList ∘ fmap Forced <$>
      mapM (mbind force_deep ∘ interp_thunk) ts
  | VAttr ts => VAttr ∘ fmap Forced <$>
      map_mapM_sorted attr_le (mbind force_deep ∘ interp_thunk) ts
  | _ => mret v
  end.

Definition indirects_env (E : env) (tαs : gmap string tattr) : env :=
  map_imap (λ y _, Some (ABS, Indirect y E tαs)) tαs ∪ E.

Definition attr_to_tattr (E : env) (α : attr) : tattr :=
  from_attr inl (inr ∘ Thunk E) α.

Definition interp1
    (force_deep : val → res val)
    (interp : env → expr → res val)
    (interp_thunk : thunk → res val)
    (interp_app : val → thunk → res val)
    (E : env) (e : expr) : res val :=
  match e with
  | ELit bl =>
      bl_ok ← guard (base_lit_ok bl);
      mret (VLit bl bl_ok)
  | EId x mke =>
      '(_,t) ← Res $ union_kinded (E !! x) (prod_map id (Thunk ∅) <$> mke);
      interp_thunk t
  | EAbs x e => mret (VClo x E e)
  | EAbsMatch ms strict e => mret (VCloMatch E ms strict e)
  | EApp e1 e2 =>
      v1 ← interp E e1;
      interp_app v1 (Thunk E e2)
  | ESeq μ' e1 e2 =>
      v ← interp E e1;
      (if μ' is DEEP then force_deep else mret) v;;
      interp E e2
  | EList es => mret (VList (Thunk E <$> es))
  | EAttr αs =>
      let E' := indirects_env E (attr_to_tattr E <$> αs) in
      mret (VAttr (from_attr (Thunk E') (Thunk E) <$> αs))
  | ELetAttr k e1 e2 =>
      v1 ← interp E e1;
      ts ← Res (maybe VAttr v1);
      interp (union_kinded ((k,.) <$> ts) E) e2
  | EBinOp op e1 e2 =>
      v1 ← interp E e1;
      f ← Res (interp_bin_op op v1);
      v2 ← interp E e2;
      t2 ← Res (f v2);
      interp_thunk t2
  | EIf e1 e2 e3 =>
      v1 ← interp E e1;
      '(b : bool) ← Res (maybe_VLit v1 ≫= maybe LitBool);
      interp E (if b then e2 else e3)
  end.

Definition interp_thunk1
    (interp : env → expr → res val)
    (interp_thunk : thunk → res val)
    (t : thunk) : res val :=
  match t with
  | Forced v => mret v
  | Thunk E e => interp E e
  | Indirect x E tαs =>
      tα ← Res $ tαs !! x;
      match tα with
      | inl e => interp (indirects_env E tαs) e
      | inr t => interp_thunk t
      end
  end.

Definition interp_app1
    (interp : env → expr → res val)
    (interp_thunk : thunk → res val)
    (interp_app : val → thunk → res val)
    (v1 : val) (t2 : thunk) : res val :=
  match v1 with
  | VClo x E e =>
      interp (<[x:=(ABS, t2)]> E) e
  | VCloMatch E ms strict e =>
      vt ← interp_thunk t2;
      ts ← Res (maybe VAttr vt);
      tαs ← Res $ interp_match ts ms strict;
      interp (indirects_env E tαs) e
  | VAttr ts =>
      t ← Res (ts !! "__functor");
      vt ← interp_thunk t;
      v ← interp_app vt (Forced v1);
      interp_app v t2
  | _ => mfail
  end.

Fixpoint force_deep (n : nat) (v : val) : res val :=
  match n with
  | O => NoFuel
  | S n => force_deep1 (force_deep n) (interp_thunk n) v
  end
with interp (n : nat) (E : env) (e : expr) : res val :=
  match n with
  | O => NoFuel
  | S n => interp1 (force_deep n) (interp n) (interp_thunk n) (interp_app n) E e
  end
with interp_thunk (n : nat) (t : thunk) : res val :=
  match n with
  | O => NoFuel
  | S n => interp_thunk1 (interp n) (interp_thunk n) t
  end
with interp_app (n : nat) (v1 : val) (t2 : thunk) : res val :=
  match n with
  | O => NoFuel
  | S n => interp_app1 (interp n) (interp_thunk n) (interp_app n) v1 t2
  end.

Definition force_deep' (n : nat) (μ : mode) : val → res val :=
  match μ with SHALLOW => mret | DEEP => force_deep n end.

Definition interp' (n : nat) (μ : mode) (E : env) (e : expr) : res val :=
  interp n E e ≫= force_deep' n μ.

Global Opaque force_deep interp interp_thunk interp_app.