Module CellEncoding

Require Import
  Utf8
  Coqlib Axioms
  AdomLib Util.

Require Import ArithLib Cells.

Set Implicit Arguments.

Local Open Scope N_scope.

Definition size (mn: N * N) : nat :=
  N.to_nat (fst mn + snd mn).

Function interleave (mn: N * N) {measure size mn} : N :=
  let '(m, n) := mn in
  match m, n with
  | N0, N0 => N0
  | N0, Npos 1 => Npos 2
  | N0, Npos (n'~0) => 4 * interleave (N0, Npos n')
  | N0, Npos (n'~1) => 2 + 4 * interleave (N0, Npos n')
  | Npos 1, N0 => Npos 1
  | Npos 1, Npos 1 => Npos 3
  | Npos 1, Npos (n'~0) => 1 + 4 * interleave (N0, Npos n')
  | Npos 1, Npos (n'~1) => 3 + 4 * interleave (N0, Npos n')
  | Npos (m'~0), N0 => 4 * interleave (Npos m', N0)
  | Npos (m'~0), Npos 1 => 2 + 4 * interleave (Npos m', N0)
  | Npos (m'~0), Npos (n'~0) => 4 * interleave (Npos m', Npos n')
  | Npos (m'~0), Npos (n'~1) => 2 + 4 * interleave (Npos m', Npos n')
  | Npos (m'~1), N0 => 1 + 4 * interleave (Npos m', N0)
  | Npos (m'~1), Npos 1 => 3 + 4 * interleave (Npos m', N0)
  | Npos (m'~1), Npos (n'~0) => 1 + 4 * interleave (Npos m', Npos n')
  | Npos (m'~1), Npos (n'~1) => 3 + 4 * interleave (Npos m', Npos n')
  end.

Function split (n: N) {measure N.to_nat n} : N * N :=
  match n with
  | 0 => (0, 0)
  | 1 => (1, 0)
  | 2 => (0, 1)
  | 3 => (1, 1)
  | Npos (n'~0~0) => let (x, y) := split (Npos n') in (2 * x, 2 * y)
  | Npos (n'~0~1) => let (x, y) := split (Npos n') in (1 + 2 * x, 2 * y)
  | Npos (n'~1~0) => let (x, y) := split (Npos n') in (2 * x, 1 + 2 * y)
  | Npos (n'~1~1) => let (x, y) := split (Npos n') in (1 + 2 * x, 1 + 2 * y)
  end.

Lemma split_4n n : split (4 * n) = let '(x, y) := split n in (2 * x, 2 * y).
Lemma split_4np1 n : split (1 + 4 * n) = let '(x, y) := split n in (1 + 2 * x, 2 * y).
Lemma split_4np2 n : split (2 + 4 * n) = let '(x, y) := split n in (2 * x, 1 + 2 * y).
Lemma split_4np3 n : split (3 + 4 * n) = let '(x, y) := split n in (1 + 2 * x, 1 + 2 * y).

Lemma split_interleave mn : split (interleave mn) = mn.

Lemma interleave_split n : interleave (split n) = n.

Unset Elimination Schemes.

Section ISO.

  Class injection (A B: Type) := inj : A → B.

  Class bijection A B `(B_of_A: injection A B) `(A_of_B: injection B A) : Prop :=
    Bijection
      { forward : ∀ a : A, inj (inj a) = a ;
        backward: ∀ b : B, inj (inj b) = b }.

  Definition inverse_bijection A B `(bijection A B) : @bijection B A _ _ :=
    Bijection _ _ backward forward.

End ISO.

Definition InjectionTrans B A C `(injection A B) `(injection B C) : injection A C
  := inj ∘ inj.
Arguments InjectionTrans B A C {_} {_} _.

Instance BijectionTrans B A C `(bijection A B) `(bijection B C)
  : bijection (InjectionTrans B A C) (InjectionTrans B C A).

Instance NOfPos : injection positive N := Pos.pred_N.
Instance PosOfN : injection N positive := N.succ_pos.

Instance PosAsN : bijection PosOfN NOfPos.

Instance PosOfZ : injection Z positive := λ z,
  match z with
  | Z0 => xH
  | Zpos p => xO p
  | Zneg p => xI p
  end.

Instance ZOfPos : injection positive Z := λ p,
  match p with
  | xH => Z0
  | xO p => Zpos p
  | xI p => Zneg p
  end.

Instance ZAsPos : bijection PosOfZ ZOfPos.

Instance NOfZ : injection Z N := λ z,
  match z with
  | Z0 => N0
  | Zpos p => Npos (xO p)
  | Zneg p => Npos (2 * p - 1)
  end.

Instance ZOfN : injection N Z := λ p,
  match p with
  | N0 => Z0
  | Npos xH => Zneg xH
  | Npos (xO p) => Zpos p
  | Npos (xI p) => Zneg (p+1)
  end.

Instance ZAsN : bijection NOfZ ZOfN.

Instance NOfPair A B `(injection A N) `(injection B N) : injection (A * B) N :=
  λ ab, let '(a, b) := ab in interleave (inj a, inj b).

Instance PairOfN A B `(injection N A) `(injection N B) : injection N (A * B) :=
  λ n, let '(a, b) := split n in (inj a, inj b).

Arguments NOfPair A B {_} {_} _.
Arguments PairOfN A B {_} {_} _.

Instance PairBijection A B `(bijection A N) `(bijection B N) : bijection (NOfPair A B) (PairOfN A B).

Definition NOfMemoryChunk : MemChunkTree.typed_memory_chunk -> N := λ c,
  match c with
    | exist AST.Mint8signed _ => 0
    | exist AST.Mint8unsigned _ => 1
    | exist AST.Mint16signed _ => 2
    | exist AST.Mint16unsigned _ => 3
    | exist AST.Mint32 _ => 4
    | exist AST.Mint64 _ => 5
    | exist AST.Mfloat32 _ => 6
    | exist AST.Mfloat64 _ => 7
    | exist _ H => match H with end
  end%N.

Definition MemoryChunkOfN : N -> MemChunkTree.typed_memory_chunk := λ n,
  match n with
    | 0 => exist _ AST.Mint8signed I
    | 1 => exist _ AST.Mint8unsigned I
    | 2 => exist _ AST.Mint16signed I
    | 3 => exist _ AST.Mint16unsigned I
    | 4 => exist _ AST.Mint32 I
    | 5 => exist _ AST.Mint64 I
    | 6 => exist _ AST.Mfloat32 I
    | _ => exist _ AST.Mfloat64 I
  end%N.

Instance NOfCell : injection cell N := λ ab,
  (match ab with
  | AClocal i i' o c => 24 * inj (i, i', o) + 3 * NOfMemoryChunk c
  | ACtemp i r => 3 * inj (i, r) + 1
  | ACglobal b o c => 24 * inj (b, o) + 2 + 3 * NOfMemoryChunk c
  end)%N.

Instance CellOfN : injection N cell := λ n,
  let '(q, r) := N.div_eucl n 3 in
  (match r with
  | 0 =>
    let '(q', r') := N.div_eucl q 8 in
    let c := MemoryChunkOfN r' in
    let '(i, i', o) := inj q' in AClocal i i' o c
  | 1 => let '(i, r) := inj q in ACtemp i r
  | _ =>
    let '(q', r') := N.div_eucl q 8 in
    let c := MemoryChunkOfN r' in
    let '(b, o) := inj q' in ACglobal b o c
  end)%N.

Instance PosOfCell : injection cell positive := InjectionTrans N cell positive.
Instance CellOfPos : injection positive cell := InjectionTrans N positive cell.

Instance CellAsPos : bijection PosOfCell CellOfPos.