Module Tailrec

Require Import Coqlib Wfsimpl Maps Errors Integers Values.
Require Import AST Linking Globalenvs.
Require Import Op Registers RTL.

Fixpoint copy_to_temps srcregs tmpreg pc (prec_code : code) : (node * code) :=
  match srcregs with
  | nil => (pc, prec_code)
  | h::t =>
      copy_to_temps t (Pos.succ tmpreg) (Pos.succ pc)
        (PTree.set pc (Iop Omove (h::nil) tmpreg (Pos.succ pc)) prec_code)
  end.

Fixpoint copy_from_temps dstregs tmpreg pc (prec_code : code) : (node * code):=
  match dstregs with
  | nil => (pc, prec_code)
  | h::t =>
      copy_from_temps t (Pos.succ tmpreg) (Pos.succ pc)
        (PTree.set pc (Iop Omove (tmpreg::nil) h (Pos.succ pc)) prec_code)
  end.


Definition build_jump start tmpreg srcregs dstregs pc0
           (code0 : code): (node * code) :=
  let (pc1, code1) := copy_to_temps (List.rev srcregs) tmpreg pc0 code0 in
  let (pc2, code2) := copy_from_temps (List.rev dstregs) tmpreg pc1 code1 in
  ((Pos.succ pc2), (PTree.set pc2 (Inop start) code2)).

Definition process_instr id start params tmpreg new_pc new_code pc instr :=
  match instr with
  | Itailcall sig (inr id') srcregs =>
      if (QualIdent.eq id id') &&
           (Nat.eq_dec (List.length srcregs) (List.length params))
      then build_jump start tmpreg srcregs params new_pc
                      (PTree.set pc (Inop new_pc) new_code)
      else (new_pc, new_code)
  | _ => (new_pc, new_code)
  end.

Definition successors_inside f :=
  let limit := max_pc_function f in
  List.forallb (fun (pc_instr : node*instruction) =>
             let (pc, instr) := pc_instr in
             List.forallb (fun x => Pos.leb x limit)
                          (successors_instr instr))
               (PTree.elements (fn_code f)).

Definition transf_function' (id: qualident) (f: function) :=
      (let start := fn_entrypoint f in
       let tmpreg := Pos.succ (RTL.max_reg_function f) in
       let params := fn_params f in
       let code' := snd
                      (PTree.fold
                         (fun (new_pc_code : node * code) (pc : node) instr =>
                            let (new_pc, new_code) := new_pc_code in
                            process_instr id start params tmpreg new_pc new_code pc instr)
                         (fn_code f)
                         ((Pos.succ (RTL.max_pc_function f)), (fn_code f))) in
       OK (mkfunction f.(fn_sig)
                          params
                          f.(fn_stacksize)
                              code'
                        start)).
  
Definition transf_function (ge: Genv.t fundef unit) (id : qualident) (f : function) :=
  if option_eq_dec eq_fundef (find_function ge (inr id) (Regmap.init Vundef)) (Some (Internal f))
  then
    if successors_inside f
    then transf_function' id f
    else Error (msg "Tailrec: successor points outside")
  else Error (msg "Tailrec: functions don't match; report and pass -fno-tailrec until fixed").
    
Definition transf_fundef (ge: Genv.t fundef unit) (id: qualident) (fd: fundef) : Errors.res fundef :=
  AST.transf_partial_fundef (transf_function ge id) fd.

Definition transf_program (p: program): Errors.res program :=
  AST.transform_partial_program2 (transf_fundef (Genv.globalenv p)) (fun i v => OK v) p.