Library Malfunction.Serialize

From MetaRocq.Utils Require Import bytestring ReflectEq utils.
From MetaRocq.Common Require Kernames.

Require Import String Ascii Bool Arith.
Require Import Malfunction.Malfunction.

Set Warnings "-masking-absolute-name".
Require Import Ceres.Ceres.
Require Import Ceres.CeresString.
Require Import Malfunction.Ceres.CeresFormat Malfunction.Ceres.CeresSerialize.

Local Open Scope sexp.
Local Open Scope string.


Fixpoint _escape_ident (_end s : String.t) : String.t :=
  match s with
  | ""%bs_end
  | String.String c s'
       if (c == "'"%byte) || (c == " "%byte) || (c == "."%byte) then String.String "_" (_escape_ident _end s')
       else match s' with
            | String.String c2 s''
                if (String.String c (String.String c2 String.EmptyString)) == "Γ"%bs
                   then ("Gamma" ++ _escape_ident _end s'')%bs
                else if (String.String c (String.String c2 String.EmptyString)) == "φ"%bs
                   then ("Phi" ++ _escape_ident _end s'')%bs
                else if (String.String c (String.String c2 String.EmptyString)) == "Δ"%bs
                   then ("Delta" ++ _escape_ident _end s'')%bs
                else if (String.String c (String.String c2 String.EmptyString)) == "π"%bs
                   then ("pi" ++ _escape_ident _end s'')%bs
                else if (String.String c (String.String c2 String.EmptyString)) == "ρ"%bs
                   then ("rho" ++ _escape_ident _end s'')%bs
                else if (String.String c (String.String c2 String.EmptyString)) == "Σ"%bs
                   then ("Sigma" ++ _escape_ident _end s'')%bs
                else String.String c (_escape_ident _end s')
            | _String.String c (_escape_ident _end s')
            end
  end.

#[export] Instance Serialize_Ident : Serialize Ident.t :=
  fun aAtom (append "$" (bytestring.String.to_string (_escape_ident ""%bs a))).

Require Sint63.

Definition sint_to_Z := Sint63.to_Z.

#[export] Instance Serialize_int : Serialize int :=
   fun ito_sexp (sint_to_Z i).

Import SpecFloat.
From Stdlib Require Numbers.HexadecimalString.
Definition string_of_specfloat (f : SpecFloat.spec_float) : string :=
  match f with
  | S754_zero signif sign then "-0.0" else "0.0"
  | S754_infinity signif sign then "neg_infinity" else "infinity"
  | S754_nan ⇒ "nan"
  | S754_finite sign p z
    let abs :=
    "0x" ++ HexadecimalString.NilZero.string_of_uint (Pos.to_hex_uint p) ++ "p" ++
      DecimalString.NilZero.string_of_int (Z.to_int z)
    in
    if sign then "-" ++ abs else abs
  end.

#[export] Instance Serialize_numconst : Serialize numconst :=
  fun amatch a with
        | numconst_Int ito_sexp (sint_to_Z i)
        | numconst_Bigint xAtom (append (CeresString.string_of_Z x) ".ibig")
        | numconst_Float64 xAtom (append (string_of_specfloat (FloatOps.Prim2SF x)) ".f64")
        end.

Definition Cons x (l : sexp) :=
  match l with
  | List lList (x :: l)
  | xx
  end.

Definition App (l1 : sexp) (l2 : sexp) :=
  match l1, l2 with
  | List l1, List l2List (l1 ++ l2)
  | _x, yy
  end.

Definition rawapp (s : sexp) (a : string) :=
  match s with
  | Atom_ (Raw s) ⇒ Atom (Raw (append s a))
  | xx
  end.

#[export] Instance Serialize_case : Serialize case :=
  fun amatch a with
        | Tag tag[Atom "tag"; Atom (sint_to_Z tag)]
        | Deftag[Atom "tag"; Atom "_"]
        | Intrange (i1, i2)if Uint63.leb i1 i2 then [ to_sexp i1 ; to_sexp i2 ] else Atom "_"
        end.

#[export] Instance Serialize_unary_num_op : Serialize unary_num_op :=
  fun amatch a with NegAtom "neg" | NotAtom "not" end.

Definition numtype_to_string (n : numtype) :=
  match n with
  | embed_inttype xmatch x with
                      | Int ⇒ ""
                      | Int32 ⇒ ".i32"
                      | Int64 ⇒ ".i64"
                      | Bigint ⇒ ".ibig"
                      end
  | Float64 ⇒ ""
  end.

Definition vector_type_to_string (n : vector_type) :=
  match n with
  | Array ⇒ ""
  | Bytevec ⇒ ".byte"
  end.

#[export] Instance Serialize_binary_arith_op : Serialize binary_arith_op :=
  fun aAtom match a with
        | Add ⇒ "+"
        | Sub ⇒ "-"
        | Mul ⇒ "*"
        | Div ⇒ "/"
        | Mod ⇒ "%"
        end.

#[export] Instance Serialize_binary_bitwise_op : Serialize binary_bitwise_op :=
  fun aAtom match a with
             | And ⇒ "&"
             | Or ⇒ "|"
             | Xor ⇒ "^"
             | Lsl ⇒ "<<"
             | Lsr ⇒ ">>"
             | Asr ⇒ "a>>"
             end.

#[export] Instance Serialize_binary_comparison : Serialize binary_comparison :=
  fun aAtom match a with
             | Lt ⇒ "<"
             | Gt ⇒ ">"
             | Lte ⇒ "<="
             | Gte ⇒ ">="
             | Eq ⇒ "=="
             end.

#[export] Instance Serialize_binary_num_op : Serialize binary_num_op :=
  fun amatch a with
        | embed_binary_arith_op xto_sexp x
        | embed_binary_bitwise_op xto_sexp x
        | embed_binary_comparison xto_sexp x
        end.

Definition Serialize_singleton_list {A} `{Serialize A} : Serialize (list A)
  := fun xsmatch xs with cons x nilto_sexp x | xsList (List.map to_sexp xs) end.

Fixpoint split_dot accl accw (s : string) :=
  match s with
  | EmptyString(string_reverse accl, string_reverse accw)
  | String c s
      if (c =? ".")%char2 then
        let accl' := match accl with EmptyStringaccw
                                | accl ⇒ (accw ++ "." ++ accl)
                     end in
        split_dot accl' EmptyString s
      else
        split_dot accl (String c accw) s
  end.
Definition before_dot s := fst (split_dot EmptyString EmptyString s).
Definition after_dot s := snd (split_dot EmptyString EmptyString s).

Fixpoint to_sexp_t (a : t) : sexp :=
  match a with
  | Mvar xto_sexp x
  | Mlambda (ids, x)[ Atom "lambda" ; to_sexp ids ; to_sexp_t x ]
  | Mapply (x, args)List (Atom "apply" :: to_sexp_t x :: List.map to_sexp_t args)
  | Mlet (binds, x)List (Atom "let" :: List.map to_sexp_binding binds ++ (to_sexp_t x :: nil))
  | Mnum xto_sexp x
  | Mstring xAtom (Str (bytestring.String.to_string x))
  | Mglobal xto_sexp ("def_" ++ x)%bs
  | Mswitch (x, sels)
      
      
      
      
      Cons (Atom "switch") (Cons (to_sexp_t x) (@Serialize_list _ (fun '(sel, t)App (to_sexp sel) ([to_sexp_t t]) ) sels))
  | Mnumop1 (op, num, x)[ rawapp (to_sexp op) (numtype_to_string num) ; to_sexp_t x ]
  | Mnumop2 (op, num, x1, x2)[ rawapp (to_sexp op) (numtype_to_string num) ; to_sexp_t x1 ; to_sexp_t x2 ]
  | Mconvert (from, to, x)[rawapp (rawapp (Atom "convert") (numtype_to_string from)) (numtype_to_string to) ; to_sexp_t x]
  | Mvecnew (ty, x1, x2)[ Atom (append "makevec" (vector_type_to_string ty)) ; to_sexp_t x1 ; to_sexp_t x2 ]
  | Mvecget (ty, x1, x2)[ Atom (append "load" (vector_type_to_string ty)) ; to_sexp_t x1 ; to_sexp_t x2 ]
  | Mvecset (ty, x1, x2, x3)[ Atom (append "store" (vector_type_to_string ty)) ; to_sexp_t x1 ; to_sexp_t x2; to_sexp_t x3 ]
  | Mveclen (ty, x)[ Atom (append "load" (vector_type_to_string ty)) ; to_sexp_t x ]
  | Mlazy x[Atom "lazy"; to_sexp_t x]
  | Mforce x[Atom "force"; to_sexp_t x]
  | Mblock (tag, xs)List (Atom "block" :: [Atom "tag"; Atom (sint_to_Z tag)] :: List.map to_sexp_t xs)
  | Mfield (i, x)[ Atom "field"; to_sexp i; to_sexp_t x]
  end
with
to_sexp_binding (a : binding) : sexp :=
  match a with
  | Unnamed x[ Atom "_"; to_sexp_t x ]
  | Named (id, x)[ to_sexp id ; to_sexp_t x ]
  | Recursive xCons (Atom "rec") (@Serialize_list _ (@Serialize_product _ _ _ to_sexp_t) x)
  end.

#[export] Instance Serialize_t : Serialize t := to_sexp_t.
#[export] Instance Serialize_binding : Serialize binding := to_sexp_binding.



Definition uncapitalize_char (c : Byte.byte) : Byte.byte :=
  let n := Byte.to_nat c in
  if (65 <=? n)%nat && (n <=? 90)%nat then match Byte.of_nat (n + 32) with Some cc | _c end
  else c.

Definition uncapitalize (s : bytestring.string) : bytestring.string :=
  match s with
  | bytestring.String.EmptyStringbytestring.String.EmptyString
  | bytestring.String.String c sbytestring.String.String (uncapitalize_char c) s
  end.

Definition encode_name (s : bytestring.string) : bytestring.string :=
  _escape_ident ""%bs s.

Definition exports (m : list (Ident.t × option t)) : list (Ident.t × option t) :=
  List.map (fun '(x, v)(("def_" ++ encode_name x)%bs, Some (Mglobal x))) m.

Definition bytestring_atom s :=
  ("$" :: bytestring.String.to_string s).

Fixpoint find_prim (id : Ident.t) (prims : primitives) : option (prim_def string) :=
  match prims with
  | nil%listNone
  | ((kn, primdef) :: prims)%list
    if ReflectEq.eqb id kn then
      match primdef with
      | Global modname labelSome (Global (bytestring_atom modname) (bytestring_atom label))
      | Primitive symbol aritySome (Primitive symbol arity)
      | ErasedSome Erased
      end
    else find_prim id prims
  end.

Section binders.
  Context (x : bytestring.string).

  Definition add_suffix n := (x ++ MRString.string_of_nat n)%bs.

  Fixpoint binders n acc :=
    match n with
    | 0 ⇒ acc
    | S nbinders n (add_suffix n :: acc)%list
    end.
End binders.

Definition mk_eta_exp n s :=
  let binders := binders "x"%bs n nil in
  [ Atom "lambda" ; to_sexp binders ; List (Atom s :: List.map to_sexp binders) ].

Definition global_serializer (prims : primitives) : Serialize (Ident.t × option t) :=
  fun '(i, b)
  match b with
  | Some xto_sexp ("def_" ++ i, x)%bs
  | None
    match find_prim i prims with
    | Some (Global modname label) ⇒
      let na := bytestring.String.to_string (uncapitalize ("def_" ++ encode_name i)%bs) in
      List ( Atom (Raw ("$" :: na)) :: [Atom "global" ; Atom (Raw modname) ; Atom (Raw label)] :: nil)
    | Some (Primitive symbol arity) ⇒
      let na := bytestring.String.to_string (uncapitalize ("def_" ++ encode_name i)%bs) in
      List ( Atom (Raw ("$" :: na)) ::
      mk_eta_exp arity (Raw (bytestring.String.to_string symbol)) :: nil)
    | Some Erased
    | None
    let na := bytestring.String.to_string (uncapitalize ("def_" ++ encode_name i)%bs) in
      List ( Atom (Raw ("$" :: na)) :: [Atom "global" ; Atom (Raw ("$Axioms")) ; Atom (Raw ("$" :: na)) ]
             :: nil)
    end
  end.

Fixpoint filter_erased_prims prims (l : list (Ident.t × option t)) : list (Ident.t × option t) :=
  match l with
  | nilnil
  | cons ((id, Some _) as x) xsx :: filter_erased_prims prims xs
  | cons ((id, None) as x) xs
    match find_prim id prims with
    | Some Erasedfilter_erased_prims prims xs
    | _x :: filter_erased_prims prims xs
    end
  end.

Fixpoint thename a (s : bytestring.String.t) :=
  match s with
  | String.EmptyStringbytestring.String.of_string (string_of_list_byte (List.rev a))
  | String.String b sif b == "."%byte
                        then thename nil s
                        else thename (b :: a)%list s
  end.

Variant program_type : Set :=
  | Standalone
  | Shared_lib (modname : bytestring.String.t) (register : bytestring.String.t).

Definition shared_lib_register modname label '(name, export) :=
  let code := (List [Atom "apply"; List [Atom "global" ;
  Atom ("$" ++ String.to_string modname)%string ; Atom ("$" ++ String.to_string label)%string]%list;
  Atom (Str (bytestring.String.to_string name));
  Atom ("$" ++ bytestring.String.to_string export)]) in
  Cons (Atom "_") (List (code :: nil)).


Definition Serialize_module prims (pt : program_type) (names : list bytestring.string): Serialize program :=
  fun '(m, x)
    let name : Ident.t := match m with
                           | (x :: l)%listfst x
                           | nil ⇒ ""%bs
                           end in
    let main := "main"%bs in
    let names := (names ++ ["main"%bs])%list in
    let shortnames : list Ident.t := List.map (fun nameuncapitalize (thename nil name)) names in
    let longnames : list sexp := List.map (fun name ⇒ (to_sexp ("def_" ++ name)%bs)) names in
    let allnames := List.combine shortnames longnames in
    let exports : list sexp := List.map (fun shortnameAtom ("$" ++ String.to_string shortname)%string)
      shortnames in
    let m := filter_erased_prims prims m in
    let linkopt :=
      match pt return list sexp with
      | Standalonenil
      | Shared_lib modname label
        List.map (shared_lib_register modname label) (List.combine names shortnames)
      end
    in
    match
      Cons (Atom "module") (@Serialize_list _ (global_serializer prims)
        (List.rev ((main, Some x) :: m)%list))
    with
      List l
        List (l
              ++ List.map (fun '(shortname,longname)Cons (Atom ("$" ++ String.to_string shortname)%string)
                             (List (longname :: nil))) allnames
              ++ linkopt
              ++ (Cons (Atom "export") (List exports) :: nil))%list
    | xx
    end.