Require Import Coq.Program.Equality.
Require Import FSets.
Require Import Imp.
Require Import Semantics.
Require Import Deadcode.
Import VSdecide.

In this chapter: a study of register allocation viewed as renaming the variables of an IMP program in a semantic-preserving manner.

1. Renaming variables in a program

Recognition of expressions that are just variables.

Definition expr_is_var (a: aexp): option id :=
match a with AId x => Some x | _ => None end.

Lemma expr_is_var_correct:
forall a x, expr_is_var a = Some x -> a = AId x.

Proof.

unfold expr_is_var; intros. destruct a; congruence.
Qed.

Section RENAMING.

Variable f: id -> id.

Renaming variables in an expression.

Fixpoint rename_aexp (a: aexp) : aexp :=
  match a with
  | ANum n => ANum n
  | AId x => AId (f x)
  | APlus a1 a2 => APlus (rename_aexp a1) (rename_aexp a2)
  | AMinus a1 a2 => AMinus (rename_aexp a1) (rename_aexp a2)
  | AMult a1 a2 => AMult (rename_aexp a1) (rename_aexp a2)
  end.

Fixpoint rename_bexp (b: bexp) : bexp :=
  match b with
  | BTrue => BTrue
  | BFalse => BFalse
  | BEq a1 a2 => BEq (rename_aexp a1) (rename_aexp a2)
  | BLe a1 a2 => BLe (rename_aexp a1) (rename_aexp a2)
  | BNot b1 => BNot (rename_bexp b1)
  | BAnd b1 b2 => BAnd (rename_bexp b1) (rename_bexp b2)
  end.

Transformation of a command:

renaming the variables
dead code elimination as in the previous chapter
coalescing: useless variable-to-variable copies x ::= y where f x = f yare turned into SKIP.

Fixpoint transf_com (c: com) (L: VS.t): com :=
  match c with
  | SKIP => SKIP
  | x ::= a =>
      if VS.mem x L then
        match expr_is_var a with
        | None => (f x) ::= (rename_aexp a)
        | Some y => if beq_id (f x) (f y) then SKIP else (f x) ::= (rename_aexp a)
        end
      else SKIP
  | (c1; c2) =>
      (transf_com c1 (live c2 L); transf_com c2 L)
  | IFB b THEN c1 ELSE c2 FI =>
      IFB rename_bexp b THEN transf_com c1 L ELSE transf_com c2 L FI
  | WHILE b DO c END =>
      WHILE rename_bexp b DO transf_com c (live (WHILE b DO c END) L) END
  end.

2. Semantic preservation

We adapt the notion of agreement between two states over a given set of live variables that was introduced in chapter Deadcode.

Definition agree (L: VS.t) (s1 s2: state) : Prop :=
forall x, VS.In x L -> s1 x = s2 (f x).

Monotonicity property.

Lemma agree_mon:
forall L L' s1 s2,
agree L' s1 s2 -> VS.Subset L L' -> agree L s1 s2.

Proof.

unfold VS.Subset, agree; intros. auto.
Qed.

Lemma agree_extensional:
forall L s1 s2 s3,
agree L s1 s2 -> (forall x, s2 x = s3 x) -> agree L s1 s3.

Proof.

unfold agree; intros. transitivity (s2 (f x)); auto.
Qed.

Agreement on the free variables of an expression implies that this expression and its renaming evaluates identically in both states.

Lemma aeval_agree:
  forall L s1 s2, agree L s1 s2 ->
  forall a, VS.Subset (fv_aexp a) L ->
  aeval s1 a = aeval s2 (rename_aexp a).

Proof.

  induction a; simpl; intros.
  auto.
  apply H. generalize H0; fsetdec.
  f_equal. apply IHa1. generalize H0; fsetdec. apply IHa2. generalize H0; fsetdec.
  f_equal. apply IHa1. generalize H0; fsetdec. apply IHa2. generalize H0; fsetdec.
  f_equal. apply IHa1. generalize H0; fsetdec. apply IHa2. generalize H0; fsetdec.
Qed.

Lemma beval_agree:
  forall L s1 s2, agree L s1 s2 ->
  forall b, VS.Subset (fv_bexp b) L ->
  beval s1 b = beval s2 (rename_bexp b).

Proof.

  induction b; simpl; intros.
  auto.
  auto.
  repeat rewrite (aeval_agree L s1 s2); auto; generalize H0; fsetdec.
  repeat rewrite (aeval_agree L s1 s2); auto; generalize H0; fsetdec.
  f_equal; auto.
  f_equal. apply IHb1. generalize H0; fsetdec. apply IHb2. generalize H0; fsetdec.
Qed.

Agreement is preserved by unilateral assignment to a dead variable.

Lemma agree_update_dead:
  forall s1 s2 L x v,
  agree L s1 s2 -> ~VS.In x L ->
  agree L (update s1 x v) s2.

Proof.

  intros; red; intros. unfold update. remember (beq_id x x0). destruct b.
  assert (x = x0). apply beq_id_eq. auto. subst x0. contradiction.
  apply H; auto.
Qed.

Agreement is preserved by simultaneous assignment to a live variable x and its renaming f x, provided none of the other live variables z is renamed to f x (non-interference property).

Lemma agree_update_live:
  forall s1 s2 L x v,
  agree (VS.remove x L) s1 s2 ->
  (forall z, VS.In z L -> z <> x -> f z <> f x) ->
  agree L (update s1 x v) (update s2 (f x) v).

Proof.

  intros; red; intros. unfold update. remember (beq_id x x0). destruct b.
Case "x = x0".
  assert (x = x0). apply beq_id_eq. auto. subst x0.
  rewrite <- beq_id_refl. auto.
Case "x <> x0".
  assert (x <> x0). apply beq_id_false_not_eq; auto.
  assert (f x0 <> f x). apply H0; auto.
  rewrite not_eq_beq_id_false.
  apply H. apply VS.remove_2; auto.
  auto.
Qed.

A special case of agree_update_live where the value being assigned is not arbitrary, but is the value of another variable y. In this case, we can relax the non-interference hypothesis (Chaitin's observation).

Lemma agree_update_move:
  forall s1 s2 L x y,
  agree (VS.union (VS.remove x L) (VS.singleton y)) s1 s2 ->
  (forall z, VS.In z L -> z <> x -> z <> y -> f z <> f x) ->
  agree L (update s1 x (s1 y)) (update s2 (f x) (s2 (f y))).

Proof.

  intros; red; intros. unfold update. remember (beq_id x x0). destruct b.
Case "x = x0".
  assert (x = x0). apply beq_id_eq. auto. subst x0.
  rewrite <- beq_id_refl. apply H. fsetdec.
Case "x <> x0".
  assert (x <> x0). apply beq_id_false_not_eq; auto.
  remember (beq_id x0 y). destruct b.
  SCase "x0 = y".
    assert (x0 = y). apply beq_id_eq. auto. subst x0.
    transitivity (s2 (f y)).
    apply H. fsetdec.
    destruct (beq_id (f x) (f y)); auto.
  SCase "x0 <> y".
    assert (x0 <> y). apply beq_id_false_not_eq; auto.
    assert (f x0 <> f x). apply H0; auto.
    rewrite not_eq_beq_id_false.
    apply H. apply VS.union_2. apply VS.remove_2; auto.
    auto.
Qed.

A special case of agree_update_move where f x = f y. In this case, the assignment in the transformed program can be optimized away.

Lemma agree_update_coalesced_move:
  forall s1 s2 L x y,
  agree (VS.union (VS.remove x L) (VS.singleton y)) s1 s2 ->
  (forall z, VS.In z L -> z <> x -> z <> y -> f z <> f x) ->
  f x = f y ->
  agree L (update s1 x (s1 y)) s2.

Proof.

  intros. apply agree_extensional with (update s2 (f x) (s2 (f y))).
  apply agree_update_move; auto.
  intros. rewrite H1. unfold update. remember (beq_id (f y) x0). destruct b.
  assert (f y = x0). apply beq_id_eq; auto.
  congruence.
  auto.
Qed.

Gathering together the various non-interference conditions above, we obtain the following correctness criterion for a candidate renaming f. Note that this is a boolean-valued function, not a predicate; therefore, it can be used as a validator for a renaming computed by an untrusted heuristic.

Fixpoint correct_allocation (c: com) (L: VS.t) : bool :=
  match c with
  | SKIP =>
      true
  | x ::= a =>
      if VS.mem x L then
        match expr_is_var a with
        | Some y =>
            VS.for_all (fun z => beq_id z x || beq_id z y || negb (beq_id (f z) (f x))) L
        | None =>
            VS.for_all (fun z => beq_id z x || negb (beq_id (f z) (f x))) L
        end
      else true
  | (c1; c2) =>
      correct_allocation c1 (live c2 L) && correct_allocation c2 L
  | IFB b THEN c1 ELSE c2 FI =>
      correct_allocation c1 L && correct_allocation c2 L
  | WHILE b DO c END =>
      correct_allocation c (live (WHILE b DO c END) L)
  end.

Remark correct_allocation_assign_1:
  forall x y L,
  VS.for_all (fun z => beq_id z x || beq_id z y || negb (beq_id (f z) (f x))) L = true ->
  forall z, VS.In z L -> z <> x -> z <> y -> f z <> f x.

Proof.

  intros.
  set (g := fun z => beq_id z x || beq_id z y || negb (beq_id (f z) (f x))) in *.
  assert (VS.For_all (fun x => g x = true) L).
    apply VS.for_all_2.
    red; intros. congruence.
    auto.
  red in H3. generalize (H3 z H0). unfold g.
  rewrite not_eq_beq_id_false; auto.
  rewrite not_eq_beq_id_false; auto.
  simpl. remember (beq_id (f z) (f x)). destruct b; simpl; intros.
  congruence.
  apply beq_id_false_not_eq; auto.
Qed.

Remark correct_allocation_assign_2:
  forall x L,
  VS.for_all (fun z => beq_id z x || negb (beq_id (f z) (f x))) L = true ->
  forall z, VS.In z L -> z <> x -> f z <> f x.

Proof.

  intros.
  set (g := fun z => beq_id z x || negb (beq_id (f z) (f x))) in *.
  assert (VS.For_all (fun x => g x = true) L).
    apply VS.for_all_2.
    red; intros. congruence.
    auto.
  red in H2. generalize (H2 z H0). unfold g.
  rewrite not_eq_beq_id_false; auto.
  simpl. remember (beq_id (f z) (f x)). destruct b; simpl; intros.
  congruence.
  apply beq_id_false_not_eq; auto.
Qed.

The proof of semantic preservation is a straightforward adaptation of the proof for dead code elimination. There are more cases to be considered in the assignment case x ::= a, corresponding to variable-to-variable moves x ::= y, either coalesced or not.

Theorem transf_correct_terminating:
  forall st c st', c / st ==> st' ->
  forall L st1,
  agree (live c L) st st1 ->
  correct_allocation c L = true ->
  exists st1', transf_com c L / st1 ==> st1' /\ agree L st' st1'.

Proof.

  induction 1; intros L st1 AG CORR; simpl; simpl in CORR.

Case "skip".
  exists st1; split. constructor. auto.

Case ":=".
  rename l into x. simpl in AG. remember (VS.mem x L) as is_live. destruct is_live.
  SCase "x is live after".
    assert (aeval st1 (rename_aexp a1) = n).
      rewrite <- H. symmetry. eapply aeval_agree. eauto. fsetdec.
    remember (expr_is_var a1) as is_var. destruct is_var as [y | ].
    SSCase "move x := y".
      assert (a1 = AId y). apply expr_is_var_correct; auto. subst a1. simpl in *.
      remember (beq_id (f x) (f y)) as coalesced. destruct coalesced.
      SSSCase "coalesced move".
        assert (f x = f y). apply beq_id_eq; auto.
        exists st1; split.
        apply E_Skip.
        rewrite <- H. apply agree_update_coalesced_move; auto.
        apply correct_allocation_assign_1; auto.
      SSSCase "preserved move".
        exists (update st1 (f x) n); split.
        apply E_Ass. simpl. auto.
        rewrite <- H at 1. rewrite <- H0.
        apply agree_update_move; auto.
        apply correct_allocation_assign_1; auto.
    SSCase "not a move".
      exists (update st1 (f x) n); split.
      apply E_Ass. auto.
      apply agree_update_live.
      eapply agree_mon; eauto. fsetdec.
      apply correct_allocation_assign_2; auto.
  SCase "x is dead after".
    exists st1; split.
    apply E_Skip.
    apply agree_update_dead. auto.
    red; intros. assert (VS.mem x L = true). apply VS.mem_1; auto. congruence.

Case "seq".
  simpl in AG. destruct (andb_true _ _ CORR).
  destruct (IHceval1 _ _ AG H1) as [st1' [E1 A1]].
  destruct (IHceval2 _ _ A1 H2) as [st2' [E2 A2]].
  exists st2'; split.
  apply E_Seq with st1'; auto.
  auto.

Case "if true".
  simpl in AG. destruct (andb_true _ _ CORR).
  assert (beval st1 (rename_bexp b1) = true).
    rewrite <- H. symmetry. eapply beval_agree; eauto. fsetdec.
  destruct (IHceval L st1) as [st1' [E A]].
    eapply agree_mon; eauto. fsetdec.
    auto.
  exists st1'; split.
  apply E_IfTrue; auto.
  auto.

Case "if false".
  simpl in AG. destruct (andb_true _ _ CORR).
  assert (beval st1 (rename_bexp b1) = false).
    rewrite <- H. symmetry. eapply beval_agree; eauto. fsetdec.
  destruct (IHceval L st1) as [st1' [E A]].
    eapply agree_mon; eauto. fsetdec.
    auto.
  exists st1'; split.
  apply E_IfFalse; auto.
  auto.

Case "while end".
  destruct (live_while_charact b1 c1 L) as [P [Q R]].
  assert (beval st1 (rename_bexp b1) = false).
    rewrite <- H. symmetry. eapply beval_agree; eauto.
  exists st1; split.
  apply E_WhileEnd. auto.
  eapply agree_mon; eauto.

Case "while loop".
  destruct (live_while_charact b1 c1 L) as [P [Q R]].
  assert (beval st1 (rename_bexp b1) = true).
    rewrite <- H. symmetry. eapply beval_agree; eauto.
  destruct (IHceval1 (live (CWhile b1 c1) L) st1) as [st2 [E1 A1]].
    eapply agree_mon; eauto.
    auto.
  destruct (IHceval2 L st2) as [st3 [E2 A2]].
    auto. auto.
  exists st3; split.
  apply E_WhileLoop with st2; auto.
  auto.
Qed.

Theorem transf_correct_diverging:
  forall st c L st1,
  c / st ==> ∞ ->
  agree (live c L) st st1 ->
  correct_allocation c L = true ->
  transf_com c L / st1 ==> ∞.

Proof.

  cofix COINDHYP; intros until st1; intros DIV AG CORR; inv DIV.

Case "seq left".
  simpl in *. destruct (andb_prop _ _ CORR).
  apply Einf_Seq_1. apply COINDHYP with st. auto. auto. auto.

Case "seq right".
  simpl in *. destruct (andb_prop _ _ CORR).
  destruct (transf_correct_terminating _ _ _ H _ _ AG H1) as [st1' [E A]].
  apply Einf_Seq_2 with st1'. auto. apply COINDHYP with st2; auto.

Case "if true".
  simpl in *. destruct (andb_prop _ _ CORR).
  assert (beval st1 (rename_bexp b) = true).
    rewrite <- H. symmetry. eapply beval_agree; eauto. fsetdec.
  apply Einf_IfTrue. auto.
  apply COINDHYP with st. auto. eapply agree_mon; eauto. fsetdec. auto.

Case "if false".
  simpl in *. destruct (andb_prop _ _ CORR).
  assert (beval st1 (rename_bexp b) = false).
    rewrite <- H. symmetry. eapply beval_agree; eauto. fsetdec.
  apply Einf_IfFalse. auto.
  apply COINDHYP with st. auto. eapply agree_mon; eauto. fsetdec. auto.

Case "while body".
  destruct (live_while_charact b c1 L) as [P [Q R]].
  assert (beval st1 (rename_bexp b) = true).
    rewrite <- H. symmetry. eapply beval_agree; eauto.
  apply Einf_WhileBody. auto.
  apply COINDHYP with st. auto. eapply agree_mon; eauto. auto.

Case "while loop".
  destruct (live_while_charact b c1 L) as [P [Q R]].
  assert (beval st1 (rename_bexp b) = true).
    rewrite <- H. symmetry. eapply beval_agree; eauto.
  destruct (transf_correct_terminating _ _ _ H0 (live (CWhile b c1) L) st1) as [st2 [E1 A1]].
    eapply agree_mon; eauto. auto.
  apply Einf_WhileLoop with st2. auto. apply E1.
  apply (COINDHYP st' (WHILE b DO c1 END) L st2). auto. auto. auto.
Qed.

End RENAMING.

3. Examples

If we give it the identity renaming, our code transformation behaves exactly like dead code elimination.

Definition trivial_alloc := fun (x: id) => x.

Eval vm_compute in (correct_allocation trivial_alloc prog (VS.singleton vr)).
Eval vm_compute in (transf_com trivial_alloc prog (VS.singleton vr)).

Result is:

   r := a;                 ===>   r := a;
   q := 0;                        skip;
   while b <= r do                while b <= r do
     r := r - b;                    r := r - b;
     q := q + 1;                    skip;
   done                           done

Here is another allocation that collapses variables r and a.

Definition my_alloc := fun (x: id) => if beq_id x vr then va else x.

Eval vm_compute in (correct_allocation my_alloc prog (VS.singleton vr)).
Eval vm_compute in (transf_com my_alloc prog (VS.singleton vr)).

Here is the result. Note the removal of r := a.

   r := a;                 ===>   skip;
   q := 0;                        skip;
   while b <= r do                while b <= a do
     r := r - b;                    a := a - b;
     q := q + 1;                    skip;
   done                           done

Finally, if we try to collapse variable r and b, the validation function correct_allocation rejects the allocation.

Definition wrong_alloc := fun (x: id) => if beq_id x vr then vb else x.

Eval vm_compute in (correct_allocation wrong_alloc prog (VS.singleton vr)).

4. Putting it all together

We assume given a register allocation heuristic that computes a candidate allocation.

Parameter regalloc_heuristic: com -> VS.t -> (id -> id).

Typically, this heuristic will be written in Caml and linked with the Caml code extracted from the present Coq development.

We then define the register allocation pass. prog is the program to be transformed and obs is the set of variables we want to observe at the end of its execution. This compilation pass can fail at compile-time, in which case it returns None.

Definition regalloc (prog: com) (obs: VS.t) : option com :=
  let f := regalloc_heuristic prog obs in
  if correct_allocation f prog obs
  then Some(transf_com f prog obs)
  else None.

Module Regalloc

1. Renaming variables in a program

2. Semantic preservation

3. Examples

4. Putting it all together