(************************************************************************)
(*         *   The Coq Proof Assistant / The Coq Development Team       *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)

(* Author: Benjamin Grégoire as part of the bytecode-based virtual reduction
   machine, Oct 2004 *)
(* Extension: Arnaud Spiwack (support for native arithmetic), May 2005 *)

open Names
open Vmvalues
open Vmbytecodes
open Vmopcodes
open Mod_subst
open CPrimitives

type emitcodes = String.t

external tcode_of_code : Bytes.t -> Vmvalues.tcode = "coq_tcode_of_code"

(* Relocation information *)
type reloc_info =
  | Reloc_annot of annot_switch
  | Reloc_const of structured_constant
  | Reloc_getglobal of Names.Constant.t
  | Reloc_caml_prim of caml_prim

let eq_reloc_info r1 r2 = match r1, r2 with
| Reloc_annot sw1, Reloc_annot sw2 -> eq_annot_switch sw1 sw2
| Reloc_annot _, _ -> false
| Reloc_const c1, Reloc_const c2 -> eq_structured_constant c1 c2
| Reloc_const _, _ -> false
| Reloc_getglobal c1, Reloc_getglobal c2 -> Constant.CanOrd.equal c1 c2
| Reloc_getglobal _, _ -> false
| Reloc_caml_prim p1, Reloc_caml_prim p2 -> CPrimitives.equal (caml_prim_to_prim p1) (caml_prim_to_prim p2)
| Reloc_caml_prim _, _ -> false

let hash_reloc_info r =
  let open Hashset.Combine in
  match r with
  | Reloc_annot sw -> combinesmall 1 (hash_annot_switch sw)
  | Reloc_const c -> combinesmall 2 (hash_structured_constant c)
  | Reloc_getglobal c -> combinesmall 3 (Constant.CanOrd.hash c)
  | Reloc_caml_prim p -> combinesmall 4 (CPrimitives.hash (caml_prim_to_prim p))

module RelocTable = Hashtbl.Make(struct
  type t = reloc_info
  let equal = eq_reloc_info
  let hash = hash_reloc_info
end)

(** We use arrays for on-disk representation. On 32-bit machines, this means we
    can only have a maximum amount of about 4.10^6 relocations, which seems
    quite a lot, but potentially reachable if e.g. compiling big proofs. This
    would prevent VM computing with these terms on 32-bit architectures. Maybe
    we should use a more robust data structure? *)
type patches = {
  reloc_infos : (reloc_info * int array) array;
}

let patch_char4 buff pos c1 c2 c3 c4 =
  Bytes.unsafe_set buff pos       c1;
  Bytes.unsafe_set buff (pos + 1) c2;
  Bytes.unsafe_set buff (pos + 2) c3;
  Bytes.unsafe_set buff (pos + 3) c4

let patch1 buff pos n =
  patch_char4 buff pos
    (Char.unsafe_chr n) (Char.unsafe_chr (n asr 8))  (Char.unsafe_chr (n asr 16))
    (Char.unsafe_chr (n asr 24))

let patch_int buff reloc =
  let buff = Bytes.of_string buff in
  let iter (reloc, npos) = Array.iter (fun pos -> patch1 buff pos reloc) npos in
  let () = CArray.iter iter reloc in
  buff

let patch (buff, pl) f =
  let map (r, pos) =
    let r = f r in
    (r, pos)
  in
  let reloc = CArray.map_left map pl.reloc_infos in
  let buff = patch_int buff reloc in
  tcode_of_code buff

(* Buffering of bytecode *)

type label_definition =
    Label_defined of int
  | Label_undefined of (int * int) list

type env = {
  mutable out_buffer : Bytes.t;
  mutable out_position : int;
  mutable label_table : label_definition array;
  (* i-th table element = Label_defined n means that label i was already
     encountered and lives at offset n
     i-th table element = Label_undefined l means that the label was not
     encountered yet, first integer is the location of the value to be patched
     in the string, seconed one is its origin *)
  reloc_info : int list RelocTable.t;
}

let out_word env b1 b2 b3 b4 =
  let p = env.out_position in
  if p >= Bytes.length env.out_buffer then begin
    let len = Bytes.length env.out_buffer in
    let new_len =
      if len <= Sys.max_string_length / 2
      then 2 * len
      else
        if len = Sys.max_string_length
        then invalid_arg "String.create"  (* Not the right exception... *)
        else Sys.max_string_length in
    let new_buffer = Bytes.create new_len in
    Bytes.blit env.out_buffer 0 new_buffer 0 len;
    env.out_buffer <- new_buffer
  end;
  patch_char4 env.out_buffer p (Char.unsafe_chr b1)
   (Char.unsafe_chr b2) (Char.unsafe_chr b3) (Char.unsafe_chr b4);
  env.out_position <- p + 4

let out env opcode =
  out_word env opcode 0 0 0

let is_immed i = Uint63.le (Uint63.of_int i) Uint63.maxuint31

(* Detect whether the current value of the accu register is no longer
   needed (i.e., the register is written before being read). If so, the
   register can be used freely; no need to save and restore it. *)
let is_accu_dead = function
  | [] -> false
  | c :: _ ->
      match c with
      | Kacc _ | Kenvacc _ | Kconst _ | Koffsetclosure _ | Kgetglobal _ -> true
      | _ -> false

let out_int env n =
  out_word env n (n asr 8) (n asr 16) (n asr 24)

(* Handling of local labels and backpatching *)

let extend_label_table env needed =
  let new_size = ref(Array.length env.label_table) in
  while needed >= !new_size do new_size := 2 * !new_size done;
  let new_table = Array.make !new_size (Label_undefined []) in
  Array.blit env.label_table 0 new_table 0 (Array.length env.label_table);
  env.label_table <- new_table

let backpatch env (pos, orig) =
  let displ = (env.out_position - orig) asr 2 in
  Bytes.set env.out_buffer  pos    @@ Char.unsafe_chr displ;
  Bytes.set env.out_buffer (pos+1) @@ Char.unsafe_chr (displ asr 8);
  Bytes.set env.out_buffer (pos+2) @@ Char.unsafe_chr (displ asr 16);
  Bytes.set env.out_buffer (pos+3) @@ Char.unsafe_chr (displ asr 24)

let define_label env lbl =
  if lbl >= Array.length env.label_table then extend_label_table env lbl;
  match (env.label_table).(lbl) with
    Label_defined _ ->
      raise(Failure "CEmitcode.define_label")
  | Label_undefined patchlist ->
      List.iter (fun p -> backpatch env p) patchlist;
      (env.label_table).(lbl) <- Label_defined env.out_position

let out_label_with_orig env orig lbl =
  if lbl >= Array.length env.label_table then extend_label_table env lbl;
  match (env.label_table).(lbl) with
    Label_defined def ->
      out_int env ((def - orig) asr 2)
  | Label_undefined patchlist ->
        (env.label_table).(lbl) <-
          Label_undefined((env.out_position, orig) :: patchlist);
      out_int env 0

let out_label env l = out_label_with_orig env env.out_position l

(* Relocation information *)

let enter env info =
  let pos = env.out_position in
  let old = try RelocTable.find env.reloc_info info with Not_found -> [] in
  RelocTable.replace env.reloc_info info (pos :: old)

let slot_for_const env c =
  enter env (Reloc_const c);
  out_int env 0

let slot_for_annot env a =
  enter env (Reloc_annot a);
  out_int env 0

let slot_for_getglobal env p =
  enter env (Reloc_getglobal p);
  out_int env 0

let slot_for_caml_prim env op =
  enter env (Reloc_caml_prim op);
  out_int env 0

(* Emission of one instruction *)

let check_prim_op = function
  | Int63head0     -> opCHECKHEAD0INT63
  | Int63tail0     -> opCHECKTAIL0INT63
  | Int63add       -> opCHECKADDINT63
  | Int63sub       -> opCHECKSUBINT63
  | Int63mul       -> opCHECKMULINT63
  | Int63div       -> opCHECKDIVINT63
  | Int63mod       -> opCHECKMODINT63
  | Int63divs      -> opCHECKDIVSINT63
  | Int63mods      -> opCHECKMODSINT63
  | Int63lsr       -> opCHECKLSRINT63
  | Int63lsl       -> opCHECKLSLINT63
  | Int63asr       -> opCHECKASRINT63
  | Int63land      -> opCHECKLANDINT63
  | Int63lor       -> opCHECKLORINT63
  | Int63lxor      -> opCHECKLXORINT63
  | Int63addc      -> opCHECKADDCINT63
  | Int63subc      -> opCHECKSUBCINT63
  | Int63addCarryC -> opCHECKADDCARRYCINT63
  | Int63subCarryC -> opCHECKSUBCARRYCINT63
  | Int63mulc      -> opCHECKMULCINT63
  | Int63diveucl   -> opCHECKDIVEUCLINT63
  | Int63div21     -> opCHECKDIV21INT63
  | Int63addMulDiv -> opCHECKADDMULDIVINT63
  | Int63eq        -> opCHECKEQINT63
  | Int63lt        -> opCHECKLTINT63
  | Int63le        -> opCHECKLEINT63
  | Int63lts       -> opCHECKLTSINT63
  | Int63les       -> opCHECKLESINT63
  | Int63compare   -> opCHECKCOMPAREINT63
  | Int63compares  -> opCHECKCOMPARESINT63
  | Float64opp     -> opCHECKOPPFLOAT
  | Float64abs     -> opCHECKABSFLOAT
  | Float64eq      -> opCHECKEQFLOAT
  | Float64lt      -> opCHECKLTFLOAT
  | Float64le      -> opCHECKLEFLOAT
  | Float64compare -> opCHECKCOMPAREFLOAT
  | Float64equal   -> opCHECKEQUALFLOAT
  | Float64classify -> opCHECKCLASSIFYFLOAT
  | Float64add     -> opCHECKADDFLOAT
  | Float64sub     -> opCHECKSUBFLOAT
  | Float64mul     -> opCHECKMULFLOAT
  | Float64div     -> opCHECKDIVFLOAT
  | Float64sqrt    -> opCHECKSQRTFLOAT
  | Float64ofUint63 -> opCHECKFLOATOFINT63
  | Float64normfr_mantissa -> opCHECKFLOATNORMFRMANTISSA
  | Float64frshiftexp -> opCHECKFRSHIFTEXP
  | Float64ldshiftexp -> opCHECKLDSHIFTEXP
  | Float64next_up    -> opCHECKNEXTUPFLOAT
  | Float64next_down  -> opCHECKNEXTDOWNFLOAT
  | Arraymake | Arrayget | Arrayset | Arraydefault | Arraycopy | Arraylength ->
    assert false

let check_caml_prim_op = function
| CAML_Arraymake -> opCHECKCAMLCALL2_1
| CAML_Arrayget -> opCHECKCAMLCALL2
| CAML_Arrayset -> opCHECKCAMLCALL3_1
| CAML_Arraydefault | CAML_Arraycopy | CAML_Arraylength -> opCHECKCAMLCALL1

let inplace_prim_op = function
  | Float64next_up | Float64next_down -> true
  | _ -> false

let check_prim_op_inplace = function
  | Float64next_up   -> opCHECKNEXTUPFLOATINPLACE
  | Float64next_down -> opCHECKNEXTDOWNFLOATINPLACE
  | _ -> assert false

let emit_instr env = function
  | Klabel lbl -> define_label env lbl
  | Kacc n ->
      if n < 8 then out env(opACC0 + n) else (out env opACC; out_int env n)
  | Kenvacc n ->
      if n >= 0 && n <= 3
      then out env(opENVACC0 + n)
      else (out env opENVACC; out_int env n)
  | Koffsetclosure ofs ->
      if Int.equal ofs 0 || Int.equal ofs 1
      then out env (opOFFSETCLOSURE0 + ofs)
      else (out env opOFFSETCLOSURE; out_int env ofs)
  | Kpush ->
      out env opPUSH
  | Kpop n ->
      out env opPOP; out_int env n
  | Kpush_retaddr lbl ->
      out env opPUSH_RETADDR; out_label env lbl
  | Kshort_apply n ->
      assert (1 <= n && n <= 4);
      out env(opAPPLY1 + n - 1)
  | Kapply n ->
      out env opAPPLY; out_int env n
  | Kappterm(n, sz) ->
      if n < 4 then (out env(opAPPTERM1 + n - 1); out_int env sz)
               else (out env opAPPTERM; out_int env n; out_int env sz)
  | Kreturn n ->
      out env opRETURN; out_int env n
  | Kjump ->
      out env opRETURN; out_int env 0
  | Krestart ->
      out env opRESTART
  | Kgrab n ->
      out env opGRAB; out_int env n
  | Kgrabrec(rec_arg) ->
      out env opGRABREC; out_int env rec_arg
  | Kclosure(lbl, n) ->
      out env opCLOSURE; out_int env n; out_label env lbl
  | Kclosurerec(nfv,init,lbl_types,lbl_bodies) ->
      out env opCLOSUREREC;out_int env (Array.length lbl_bodies);
      out_int env nfv; out_int env init;
      let org = env.out_position in
      Array.iter (out_label_with_orig env org) lbl_types;
      let org = env.out_position in
      Array.iter (out_label_with_orig env org) lbl_bodies
  | Kclosurecofix(nfv,init,lbl_types,lbl_bodies) ->
      out env opCLOSURECOFIX;out_int env (Array.length lbl_bodies);
      out_int env nfv; out_int env init;
      let org = env.out_position in
      Array.iter (out_label_with_orig env org) lbl_types;
      let org = env.out_position in
      Array.iter (out_label_with_orig env org) lbl_bodies
  | Kgetglobal q ->
      out env opGETGLOBAL; slot_for_getglobal env q
  | Kconst (Const_b0 i) when is_immed i ->
      if i >= 0 && i <= 3
          then out env (opCONST0 + i)
          else (out env opCONSTINT; out_int env i)
  | Kconst c ->
      out env opGETGLOBAL; slot_for_const env c
  | Kmakeblock(n, t) ->
      if 0 < n && n < 4 then (out env(opMAKEBLOCK1 + n - 1); out_int env t)
      else (out env opMAKEBLOCK; out_int env n; out_int env t)
  | Kmakeswitchblock(typlbl,swlbl,annot,sz) ->
      out env opMAKESWITCHBLOCK;
      out_label env typlbl; out_label env swlbl;
      slot_for_annot env annot;out_int env sz
  | Kswitch (tbl_const, tbl_block) ->
      let lenb = Array.length tbl_block in
      let lenc = Array.length tbl_const in
      assert (lenb < 0x100 && lenc < 0x1000000);
      out env opSWITCH;
      out_word env lenc (lenc asr 8) (lenc asr 16) (lenb);
(*      out_int env (Array.length tbl_const + (Array.length tbl_block lsl 23)); *)
      let org = env.out_position in
      Array.iter (out_label_with_orig env org) tbl_const;
      Array.iter (out_label_with_orig env org) tbl_block
  | Kpushfields n ->
      out env opPUSHFIELDS;out_int env n
  | Kfield n ->
      if n <= 1 then out env (opGETFIELD0+n)
      else (out env opGETFIELD;out_int env n)
  | Ksetfield n ->
      out env opSETFIELD; out_int env n
  | Ksequence _ -> invalid_arg "Vmemitcodes.emit_instr"
  | Kproj p -> out env opPROJ; out_int env (Projection.Repr.arg p)
  | Kensurestackcapacity size -> out env opENSURESTACKCAPACITY; out_int env size
  | Kbranch lbl -> out env opBRANCH; out_label env lbl
  | Kprim (op, (q,_u)) ->
      out env (check_prim_op op);
      slot_for_getglobal env q

  | Kcamlprim (op,lbl) ->
    out env (check_caml_prim_op op);
    out_label env lbl;
    slot_for_caml_prim env op

  | Kstop -> out env opSTOP

(* Emission of a current list and remaining lists of instructions. Include some peephole optimization. *)

let rec emit env insns remaining = match insns with
  | [] ->
     (match remaining with
       [] -> ()
     | (first::rest) -> emit env first rest)
  (* Peephole optimizations *)
  | Kpush :: Kacc n :: c ->
      if n = 0 then out env opPUSH
      else if n < 8 then out env (opPUSHACC1 + n - 1)
      else (out env opPUSHACC; out_int env n);
      emit env c remaining
  | Kpush :: Kenvacc n :: c ->
      if n >= 0 && n <= 3
      then out env(opPUSHENVACC0 + n)
      else (out env opPUSHENVACC; out_int env n);
      emit env c remaining
  | Kpush :: Koffsetclosure ofs :: c ->
      if Int.equal ofs 0 || Int.equal ofs 1
      then out env(opPUSHOFFSETCLOSURE0 + ofs)
      else (out env opPUSHOFFSETCLOSURE; out_int env ofs);
      emit env c remaining
  | Kpush :: Kgetglobal id :: c ->
      out env opPUSHGETGLOBAL; slot_for_getglobal env id; emit env c remaining
  | Kpush :: Kconst (Const_b0 i) :: c when is_immed i ->
      if i >= 0 && i <= 3
      then out env (opPUSHCONST0 + i)
      else (out env opPUSHCONSTINT; out_int env i);
      emit env c remaining
  | Kpush :: Kconst const :: c ->
      out env opPUSHGETGLOBAL; slot_for_const env const;
      emit env c remaining
  | Kpushfields 1 :: c when is_accu_dead c ->
      out env opGETFIELD0;
      emit env (Kpush :: c) remaining
  | Kpop n :: Kjump :: c ->
      out env opRETURN; out_int env n; emit env c remaining
  | Ksequence c1 :: c ->
      emit env c1 (c :: remaining)
  | Kprim (op1, (q1, _)) :: Kprim (op2, (q2, _)) :: c when inplace_prim_op op2 ->
      out env (check_prim_op op1);
      slot_for_getglobal env q1;
      out env (check_prim_op_inplace op2);
      slot_for_getglobal env q2;
      emit env c remaining
  (* Default case *)
  | instr :: c ->
      emit_instr env instr; emit env c remaining

(* Initialization *)

type to_patch = emitcodes * patches

(* Substitution *)
let subst_strcst s sc =
  match sc with
  | Const_sort _ | Const_b0 _ | Const_univ_level _ | Const_val _ | Const_uint _
  | Const_float _ -> sc
  | Const_ind ind -> let kn,i = ind in Const_ind (subst_mind s kn, i)

let subst_annot _ (a : annot_switch) = a

let subst_reloc s ri =
  match ri with
  | Reloc_annot a -> Reloc_annot (subst_annot s a)
  | Reloc_const sc -> Reloc_const (subst_strcst s sc)
  | Reloc_getglobal kn -> Reloc_getglobal (subst_constant s kn)
  | Reloc_caml_prim _ -> ri

let subst_patches subst p =
  let infos = CArray.map (fun (r, pos) -> (subst_reloc subst r, pos)) p.reloc_infos in
  { reloc_infos = infos; }

let subst_to_patch s (code, pl) =
  (code, subst_patches s pl)

type body_code =
  | BCdefined of to_patch * fv
  | BCalias of Names.Constant.t
  | BCconstant

let subst_body_code s = function
| BCdefined (tp, fv) -> BCdefined (subst_to_patch s tp, fv)
| BCalias cu -> BCalias (subst_constant s cu)
| BCconstant -> BCconstant

let to_memory code =
  let env = {
    out_buffer = Bytes.create 1024;
    out_position = 0;
    label_table = Array.make 16 (Label_undefined []);
    reloc_info = RelocTable.create 91;
  } in
  emit env code [];
  (** Later uses of this string are all purely functional *)
  let code = Bytes.sub_string env.out_buffer 0 env.out_position in
  let code = CString.hcons code in
  let fold reloc npos accu = (reloc, Array.of_list npos) :: accu in
  let reloc = RelocTable.fold fold env.reloc_info [] in
  let reloc = { reloc_infos = CArray.of_list reloc } in
  Array.iter (fun lbl ->
    (match lbl with
      Label_defined _ -> assert true
    | Label_undefined patchlist ->
        assert (patchlist = []))) env.label_table;
  (code, reloc)
