include "LIN/LIN.ma".
include "joint/SemanticUtils.ma".

let rec pre_erase_lin_internal_function'
  (globals: list ident) (the_program: list (lin_statement globals))
    (labels: list label) 
      on the_program: ((identifier_map ? label) × (list (lin_statement globals))) ≝
  match the_program with
  [ nil        ⇒
    match labels with
    [ nil ⇒ 〈empty_map …, [ ]〉 (* XXX: labels should be empty *)
    | _   ⇒ ⊥
    ]
  | cons hd tl ⇒
    let 〈lbl, instr〉 ≝ hd in
      match lbl with
      [ None ⇒
        let 〈maps, instructions〉 ≝ pre_erase_lin_internal_function' globals tl labels in
          match instr with
          [ sequential seq l ⇒
            match seq with
            [ COST_LABEL cost ⇒ 〈maps, instructions〉
            | _ ⇒ 〈maps, 〈None …, instr〉 :: instructions〉
            ]
          | _ ⇒ 〈maps, 〈None …, instr〉 :: instructions〉
          ]
      | Some lbl ⇒
        match instr with
        [ sequential seq l ⇒
          match seq with
          [ COST_LABEL cost ⇒ pre_erase_lin_internal_function' globals tl (lbl :: labels)
          | _ ⇒
            let 〈maps, instructions〉 ≝ pre_erase_lin_internal_function' globals tl labels in
            let maps ≝ foldr ? ? (λl. λmap. add … map l lbl) maps (lbl :: labels)
            in
              〈maps, instructions〉
          ]
        | _ ⇒
          let 〈maps, instructions〉 ≝ pre_erase_lin_internal_function' globals tl labels in
          let maps ≝ foldr ? ? (λl. λmap. add … map l lbl) maps (lbl :: labels)
          in
            〈maps, instructions〉
        ]
      ]
  ].
  @daemon
qed.

definition pre_erase_lin_internal_function ≝
  λglobals: list ident.
  λint_fun: joint_internal_function … (lin_params globals).
    pre_erase_lin_internal_function' … (joint_if_code … (lin_params …) int_fun) [ ].

definition relabel_joint_instruction:
    ∀globals: list ident.
    ∀params: params_.
      identifier_map LabelTag label → joint_instruction params globals → joint_instruction params globals ≝
  λglobals: list ident.
  λparams.
  λmap.
  λinstr.
  match instr with
  [ COND acc_a lbl ⇒
    let located ≝ lookup_def ?? map lbl lbl in
      COND params globals acc_a located
  | _ ⇒ instr
  ].

definition relabel_joint_statement:
  ∀globals: list ident.
  ∀params: params_.
    identifier_map LabelTag label → joint_statement params globals → joint_statement params globals ≝
  λglobals: list ident.
  λparams: params_.
  λmap: identifier_map LabelTag label.
  λstmt: joint_statement params globals.
  match stmt with
  [ sequential seq l ⇒
    let relabelled ≝ relabel_joint_instruction globals params map seq in
      sequential params globals relabelled l
  | RETURN ⇒ RETURN params globals
  | GOTO l ⇒
    let located ≝ lookup_def ?? map l l in
      GOTO params globals located
  ].

let rec relabel_lin_internal_function'
  (globals: list ident) (map: identifier_map LabelTag label)
    (program: list (lin_statement globals))
      on program: list (lin_statement globals) ≝
  match program with
  [ nil        ⇒ [ ]
  | cons hd tl ⇒
    let 〈label, stmt〉 ≝ hd in
    let relabelled ≝ relabel_joint_statement globals lin_params_ map stmt in
    let label ≝
      match label with
      [ None ⇒ None …
      | Some label ⇒
          Some … (lookup_def ?? map label label)
      ]
    in
      〈label, relabelled〉 :: relabel_lin_internal_function' globals map tl
  ].

definition relabel_lin_internal_function ≝
  λglobals.
  λmap.
  λint_fun: joint_internal_function globals (lin_params globals).
    relabel_lin_internal_function' globals map (joint_if_code … int_fun).

definition empty_graph ≝
  λa: Type[0].
    empty_map LabelTag a.

let rec pre_erase_graph_internal_function'
  (params1: params1) (globals: list ident)
    (the_graph: codeT globals (graph_params params1 globals)) (labels: list label)
      (entry_point: label) (size_counter: nat) (map: identifier_map LabelTag label)
        (visited: identifier_map LabelTag bool)
          on size_counter: ((identifier_map LabelTag label) × (identifier_map LabelTag bool) × (codeT globals (graph_params params1 globals))) ≝
  match size_counter return λ_.  ((identifier_map LabelTag label) × (identifier_map LabelTag bool) × (codeT globals (graph_params params1 globals))) with
  [ O ⇒
    match labels with
    [ nil ⇒ 〈map, visited, the_graph〉
    | _   ⇒ ⊥
    ]
  | S size_counter ⇒
    let statement ≝ lookup LabelTag ? the_graph entry_point in
    match statement return λs: option (joint_statement (graph_params params1 globals) globals). ((identifier_map LabelTag label) × (identifier_map LabelTag bool) × (codeT globals (graph_params params1 globals))) with
    [ None ⇒ ⊥ (* XXX: should not happen *)
    | Some statement ⇒
      let entered_previously ≝ lookup_def … visited entry_point false in
      match entered_previously return λe: bool. ((identifier_map LabelTag label) × (identifier_map LabelTag bool) × (codeT globals (graph_params params1 globals))) with
      [ true ⇒ 〈map, visited, the_graph〉
      | false ⇒ (* XXX: never visited here before *)
        let visited ≝ add … visited entry_point true in
        match statement return λs: joint_statement (graph_params params1 globals) globals. ((identifier_map LabelTag label) × (identifier_map LabelTag bool) × (codeT globals (graph_params params1 globals))) with
        [ sequential seq l ⇒
          match seq with
          [ COST_LABEL cost_label ⇒
            pre_erase_graph_internal_function' params1 globals the_graph (entry_point :: labels) l size_counter map visited
          | COND acc_a cond_label ⇒
            let 〈map, visited, the_graph〉 ≝ pre_erase_graph_internal_function' params1 globals the_graph [ ] l size_counter map visited in
            let 〈map, visited, the_graph〉 ≝ pre_erase_graph_internal_function' params1 globals the_graph [ ] cond_label size_counter map visited in
            let map ≝ foldr … (λl. λmap. add … map l entry_point) map (entry_point :: labels) in
              〈map, visited, the_graph〉
          | _ ⇒
            let 〈map, visited, the_graph〉 ≝ pre_erase_graph_internal_function' params1 globals the_graph [ ] l size_counter map visited in
            let map ≝ foldr … (λl. λmap. add … map l entry_point) map (entry_point :: labels) in
              〈map, visited, the_graph〉
          ]
        | RETURN ⇒
          let map ≝ foldr … (λl. λmap. add … map l entry_point) map (entry_point :: labels) in
            〈map, visited, the_graph〉
        | GOTO l ⇒
          let 〈map, visited, the_graph〉 ≝ pre_erase_graph_internal_function' params1 globals the_graph [ ] l size_counter map visited in
          let map ≝ foldr … (λl. λmap. add … map l entry_point) map (entry_point :: labels) in
            〈map, visited, the_graph〉
        ]
      ]
    ]
  ].
  cases daemon
qed.

definition pre_erase_graph_internal_function ≝
  λparams1: params1.
  λglobals: list ident.
  λint_fun: joint_internal_function globals (graph_params params1 globals).
  match joint_if_entry … int_fun with
  [ dp entry entry_proof ⇒
    let code ≝ joint_if_code … int_fun in
    let 〈map, visited, the_graph〉 ≝ pre_erase_graph_internal_function' params1 globals code [ ] entry (graph_num_nodes … code) (empty_map …) (empty_map …) in
      〈map, the_graph〉
  ].

let rec relabel_graph'
  (params1: params1) (globals: list ident)
    (the_graph: codeT globals (graph_params params1 globals)) (map: identifier_map LabelTag label)
      (visited: identifier_map LabelTag bool) (entry_point: label) (size_counter: nat)
        on size_counter: codeT globals (graph_params params1 globals) ≝
  match size_counter return λs: nat. codeT globals (graph_params params1 globals) with
  [ O ⇒ empty_graph …
  | S size_counter ⇒
    let relabelled_entry_point ≝ lookup_def … map entry_point entry_point in
    let statement ≝ lookup LabelTag ? the_graph entry_point in
    match statement return λs: option (joint_statement (graph_params params1 globals) globals). codeT globals (graph_params params1 globals) with
    [ None ⇒ ⊥ (* XXX: should not happen *)
    | Some statement ⇒
      let entered_previously ≝ lookup_def … visited entry_point false in
      match entered_previously return λe: bool. codeT globals (graph_params params1 globals) with
      [ true ⇒ the_graph
      | false ⇒ (* XXX: never visited here before *)
        let visited ≝ add … visited entry_point true in
        match statement return λs: joint_statement (graph_params params1 globals) globals. codeT globals (graph_params params1 globals) with
        [ sequential seq l ⇒
          let 〈the_graph, new_seq〉 ≝
            match seq return λs. (codeT globals (graph_params params1 globals)) × ? with
            [ COND acc_a cond_label ⇒
              let relabelled ≝ lookup_def … map cond_label cond_label in
              let tail_graph ≝ relabel_graph' params1 globals the_graph map visited cond_label size_counter in
                〈tail_graph, COND … acc_a relabelled〉
            | _ ⇒ 〈the_graph, seq〉
            ]
          in
            let tail_graph ≝ relabel_graph' params1 globals the_graph map visited l size_counter in
            let relabelled ≝ lookup_def … map l l in
              add LabelTag … tail_graph relabelled_entry_point (sequential … new_seq relabelled)
        | RETURN ⇒ the_graph
        | GOTO l ⇒
          let relabelled ≝ lookup_def … map l l in
            add LabelTag … the_graph relabelled_entry_point (GOTO … relabelled)          
        ]
      ]
    ]
  ].
  cases daemon
qed.

definition relabel_graph ≝
  λglobals: list ident.
  λparams1: params1.
  λmap: identifier_map LabelTag label.
  λint_fun: joint_internal_function globals (graph_params params1 globals).
    match joint_if_entry … int_fun with
    [ dp entry entry_prf ⇒
      let code ≝ joint_if_code … int_fun in
        relabel_graph' … code map (empty_map …) entry (graph_num_nodes … code)
    ].
    
definition pre_erase_joint_internal_function ≝
  λglobals: list ident.
  λpars: params globals.
  λerasure_function: joint_internal_function globals pars → (identifier_map LabelTag label) × (codeT … pars).
  λrelabelling_function: identifier_map LabelTag label → joint_internal_function globals pars → codeT … pars.
  λint_fun: joint_internal_function globals pars.
  match joint_if_entry … int_fun with
  [ dp entry entry_prf ⇒
    match joint_if_exit … int_fun with
    [ dp exit exit_prf ⇒
      let 〈maps, code〉 ≝ erasure_function int_fun in
      let int_fun ≝ set_joint_code globals pars int_fun code entry exit in
      let code ≝ relabelling_function maps int_fun in
        set_joint_code globals pars int_fun code (lookup_def … maps entry entry)
        (lookup_def … maps exit exit)
    ]
  ].
  cases daemon (* XXX *)
qed.

definition erase_graph_program:
    ∀params1: params1. joint_program (graph_params params1) → joint_program (graph_params params1) ≝ 
  λparams1.
  λprogram.
    transform_program … program (transf_fundef … (λvarnames. pre_erase_joint_internal_function varnames … (pre_erase_graph_internal_function varnames …) (relabel_graph varnames …))).

definition erase_lin_program: lin_program → lin_program ≝
  λprogram.
    transform_program … program (transf_fundef … (λvarnames. pre_erase_joint_internal_function varnames … (pre_erase_lin_internal_function varnames …) (relabel_lin_internal_function varnames …))).