From 71c0a3413422e5ba146de4dbd9ef86d75baccb3d Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 26 Apr 2024 06:08:32 +0100
Subject: [PATCH 1/5] Bump toolchain to nightly-2024-04-25

---
 Arm/BitVec.lean         |  2 +-
 Arm/Map.lean            | 13 ++-----------
 Arm/SeparateProofs.lean |  3 +--
 lake-manifest.json      | 13 ++-----------
 lean-toolchain          |  2 +-
 5 files changed, 7 insertions(+), 26 deletions(-)

diff --git a/Arm/BitVec.lean b/Arm/BitVec.lean
index a7218c5b..76e5d389 100644
--- a/Arm/BitVec.lean
+++ b/Arm/BitVec.lean
@@ -241,7 +241,7 @@ protected theorem zero_le_sub (x y : BitVec n) :
 
 @[simp]
 protected theorem zero_or (x : BitVec n) : 0#n ||| x = x := by
-  unfold HOr.hOr instHOr OrOp.or instOrOpBitVec BitVec.or
+  unfold HOr.hOr instHOrOfOrOp BitVec.instOrOp BitVec.or
   simp only [toNat_ofNat, Nat.or_zero]
   congr
 
diff --git a/Arm/Map.lean b/Arm/Map.lean
index b6428f98..0ecd78ce 100644
--- a/Arm/Map.lean
+++ b/Arm/Map.lean
@@ -118,17 +118,8 @@ def Map.size (m : Map α β) : Nat :=
 @[simp] theorem Map.size_erase_le [DecidableEq α] (m : Map α β) (a : α) : (m.erase a).size ≤ m.size := by
   induction m <;> simp [erase, size] at *
   split
-  next =>
-    -- (FIXME) This could be discharged by omega in
-    -- leanprover/lean4:nightly-2024-02-24, but not in
-    -- leanprover/lean4:nightly-2024-03-01.
-    exact Nat.le_succ_of_le (by assumption)
-  next =>
-    simp;
-    -- (FIXME) This could be discharged by omega in
-    -- leanprover/lean4:nightly-2024-02-24, but not in
-    -- leanprover/lean4:nightly-2024-03-01.
-    exact Nat.succ_le_succ (by assumption)
+  next => omega
+  next => simp; omega
 
 @[simp] theorem Map.size_erase_eq [DecidableEq α] (m : Map α β) (a : α) : m.contains a = false → (m.erase a).size = m.size := by
   induction m <;> simp [erase, size] at *
diff --git a/Arm/SeparateProofs.lean b/Arm/SeparateProofs.lean
index 39cbb52f..3719e0a5 100644
--- a/Arm/SeparateProofs.lean
+++ b/Arm/SeparateProofs.lean
@@ -34,8 +34,7 @@ theorem n_minus_1_lt_2_64_1 (n : Nat)
   refine BitVec.val_bitvec_lt.mp ?a
   simp [BitVec.bitvec_to_nat_of_nat]
   have : n - 1 < 2 ^ 64 := by omega
-  simp_all [Nat.mod_eq_of_lt]
-  exact Nat.sub_lt_left_of_lt_add h1 h2
+  omega
 
 -- (FIXME) Prove for all bitvector widths.
 theorem BitVec.add_sub_self_left_64 (a m : BitVec 64) :
diff --git a/lake-manifest.json b/lake-manifest.json
index b5546a11..68788073 100644
--- a/lake-manifest.json
+++ b/lake-manifest.json
@@ -4,19 +4,10 @@
  [{"url": "https://github.com/leanprover/std4",
    "type": "git",
    "subDir": null,
-   "rev": "a178ab58c07c37d919079ac3a5e4514fd85b791b",
+   "rev": "406fe30baf9887a5cbc9cb3471d6d84fd29bfc52",
    "name": "std",
    "manifestFile": "lake-manifest.json",
-   "inputRev": "nightly-testing-2024-02-22",
-   "inherited": false,
-   "configFile": "lakefile.lean"},
-  {"url": "https://github.com/leanprover-community/lean-auto.git",
-   "type": "git",
-   "subDir": null,
-   "rev": "468c9144654768f46e7fce0e698792f28ce826d5",
-   "name": "auto",
-   "manifestFile": "lake-manifest.json",
-   "inputRev": "nightly-testing-2024-02-22",
+   "inputRev": "nightly-testing-2024-04-25",
    "inherited": false,
    "configFile": "lakefile.lean"}],
  "name": "lnsym",
diff --git a/lean-toolchain b/lean-toolchain
index 2f334869..f4646857 100644
--- a/lean-toolchain
+++ b/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:nightly-2024-03-28
+leanprover/lean4:nightly-2024-04-24

From fbb95d4a8fd2d10b3df63efacffb48c1554ac300 Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 26 Apr 2024 06:09:02 +0100
Subject: [PATCH 2/5] Add basic `elim` tactic

This is a minimal attempt at providing a elimination tactic `elim`.

Elimination rules are a unified and principled way to destruct various
kinds of premises, and ubiquituous in Isabelle.

While std and mathlib have custom destruction tactics, we try to avoid
external dependencies as much as possible here.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 Tactics/Elim.lean | 159 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 159 insertions(+)
 create mode 100644 Tactics/Elim.lean

diff --git a/Tactics/Elim.lean b/Tactics/Elim.lean
new file mode 100644
index 00000000..c14e48fc
--- /dev/null
+++ b/Tactics/Elim.lean
@@ -0,0 +1,159 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Hanno Becker
+-/
+
+/- This is a minimal attempt and providing a elimination tactic `elim`
+   Elimination rules are a unified and principled way to destruct various
+   kinds of premises, and ubiquituous in Isabelle. See below for some examples.
+
+   While std and mathlib have custom destruction tactics, we try to avoid
+   external dependencies as much as possible here. -/
+
+import Lean
+open Lean Meta Elab.Tactic
+
+--
+-- Almost a copy from stdlib, but remembering the fvarid
+--
+
+/-- Try to close goal by assumption. Upon succes, return fvar id of
+  matching assumption. Otherwise, return none. --/
+def assumptionCore' (mvarId : MVarId) : MetaM (Option FVarId) :=
+  mvarId.withContext do
+    mvarId.checkNotAssigned `assumption
+    match (← findLocalDeclWithType? (← mvarId.getType)) with
+    | none => return none
+    | some fvarId => mvarId.assign (mkFVar fvarId); return (some fvarId)
+
+/-- Close goal `mvarId` using an assumption. Throw error message if failed. -/
+def assumption' (mvarId : MVarId) : MetaM FVarId := do
+  let res ← assumptionCore' mvarId
+  match res with
+  | none => throwTacticEx `assumption mvarId "assumption' tactic failed"
+  | some fvarid => return fvarid
+
+--
+--
+--
+
+def erule (e : Expr) (mvid : MVarId) (with_intro : Bool := true) : MetaM (List MVarId) := do
+  let s ← saveState
+  try
+    let mvids ← mvid.apply e
+    match mvids with
+    | [] => throwError "ill-formed elimination rule {e}"
+    | main :: other =>
+      -- Try to solve main goal by assumption, remember fvar of hypothesis
+      let fvid ← assumption' main
+      -- Remove hypothesis from all other goals
+      let other' ← other.mapM (fun mvid => do
+        if (← mvid.isAssignedOrDelayedAssigned) then
+          return mvid
+        let mvid ← mvid.clear fvid
+        if with_intro then
+          let (_, mvid) ← mvid.intros
+          return mvid
+        else
+          return mvid
+      )
+      return other'
+  catch _ =>
+    restoreState s
+    throwError "erule_tac failed"
+
+-- Run erule repeatedly
+def elim (e : Expr) (mvid : MVarId) : MetaM (List MVarId) := do
+   Meta.repeat1' (erule e) [mvid]
+
+elab "erule" e:term : tactic => do
+   let e ← elabTermForApply e
+   Elab.Tactic.liftMetaTactic (erule e)
+
+elab "elim" e:term : tactic => do
+   let e ← elabTermForApply e
+   Elab.Tactic.liftMetaTactic (elim e)
+
+-- Some examples of elimination rules
+--
+-- Elimination rules abound, but here are some simple examples.
+-- While those can be handled by existing tactics as well, the point
+-- is to show that `elim` unifies various types of deconstruction
+-- in a single tactic, parametrized by a suitable elimination rule.
+
+-- 1/ Disjunction elimination
+example {A B C D E : Prop} (h : E): A ∨ B → C ∨ D → E := by
+  intros
+  elim Or.elim /- Now we have 4 goals -/
+  <;> exact h
+  done
+
+-- 2/ Conjunction elimination
+theorem conjE {P Q R: Prop}: P ∧ Q → (P → Q → R) → R := by
+  intros h g; cases h; apply g <;> assumption
+
+example {A B C D E} (h : E) (t : A ∧ B) (t' : C ∧ D) : E := by
+  intros
+  elim conjE -- Separate hypothesis A, B, C, D
+  exact h
+  done
+
+-- 3/ Destructing existentials
+
+theorem exE' {R : Prop} {α : Type} {P : α → Prop}:
+   (∃x:α, P x) → (∀x:α, (P x → R)) → R := by
+   intro f g
+   cases f
+   apply g
+   assumption
+
+example {R : Prop} (h : R) (a: ∃x:Nat, x > 42 ∧ x < 100) (b: ∃y:Nat, y > 100 ∧ y < 128) : R := by
+  elim exE'
+  exact h
+  done
+
+-- 3/ Destructing compound definitions
+
+structure dummyStruct :=
+  mkDummy :: (f0: Nat) (f1: Nat)
+
+def isValidDummy (x : dummyStruct) : Prop :=
+   (x.f0 * x.f1) > 42 ∧
+   (∃d:Nat, d > 5 ∧ d ∣ x.f0 ∧ d ∣ x.f1)
+
+def isValidDummyE {x : dummyStruct} (v: isValidDummy x):
+   (∀d:Nat, d > 5 → d ∣ x.f0 → d ∣ x.f1 → x.f0 * x.f1 > 42 → R) → R := by
+  intro h
+  simp [isValidDummy] at v
+  elim conjE
+  elim exE'
+  elim conjE
+  apply h <;> assumption
+  done
+
+example {x y : dummyStruct} (vx : isValidDummy x) (vy : isValidDummy y) : true := by
+  elim isValidDummyE -- Breaks up both validity assumptions
+  simp
+  done
+
+-- 4/ Case-splitting on natural numbers
+
+theorem bound_nat_succ {i n : Nat} (lt: i < n) : i = n-1 ∨ i < n-1 := by omega
+theorem bound_nat_succE {R : Prop} {i n : Nat} (lt: i < n) (t0 : i = n-1 → R) (t1 : i < n-1 → R) : R :=
+  by
+     cases (bound_nat_succ lt)
+     apply t0
+     assumption
+     apply t1
+     assumption
+
+-- Syntax directed version of bound_nat_succE
+theorem bound_nat_succE' {R : Prop} {i n : Nat} (lt: i < Nat.succ n) (t0 : i = n → R) (t1 : i < n → R) : R := by
+  apply (bound_nat_succE lt) <;> simp <;> assumption
+
+example {P : Nat → Prop} {i : Nat} (h : P i /- just to avoid sorry-/)
+  (lt: i < Nat.succ (Nat.succ (Nat.succ (Nat.succ 0)))) : P i := by
+  elim bound_nat_succE' /- Now we have 4 goals for i=0,1,2,3 -/
+  <;> exact h
+  done

From cc53bcd9e92b3be45edd114794b628bdcbc76bf4 Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 26 Apr 2024 06:12:35 +0100
Subject: [PATCH 3/5] Add basic material on fixed-length lists

This commit adds definition and basic constructions and lemmas
for a type `Vec t n` of length-n lists of type `t`.

This exists in Mathlib, but in the interest of keeping dependencies
down (and as a useful learning exercise), we re-implement the required
material here.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 Arm/Vec.lean | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 Arm/Vec.lean

diff --git a/Arm/Vec.lean b/Arm/Vec.lean
new file mode 100644
index 00000000..58f8f056
--- /dev/null
+++ b/Arm/Vec.lean
@@ -0,0 +1,60 @@
+import Lean
+import Tactics.Elim
+import Init.Data.List.Lemmas
+
+open Lean Meta Elab.Tactic
+
+/- Ad-hoc definitions and lemmas about fixed-length lists
+
+MathLib has `Vec` which could be used instead; for now, we use our own
+version to keep dependencies down.
+
+-/
+
+abbrev Vec α (n : Nat) := { v : List α // v.length = n }
+
+@[simp]
+theorem Vec_eq_transport {k l : Nat} (h : k = l) (x : Vec α k) : (h▸x).val = x.val := by cases h; simp
+
+@[simp] def Vec.inBounds (_ : Vec α n) (i : Nat) : Prop := i < n
+
+@[simp] def Vec.get {n : Nat} (v : Vec α n) (i : Nat) (ok : v.inBounds i) : α :=
+  let hi: i < v.val.length := by simpa [v.property] using ok
+  v.val[i]'hi
+
+def Vec.empty : Vec α 0 := ⟨[], by simp⟩
+
+def Vec.cons {n : Nat} (x : α) (v : Vec α n) : Vec α (n+1) :=
+  ⟨List.cons x v.val, by simp [v.property]⟩
+
+theorem Vec.ext'' (x y : Vec α n) (h: x.val = y.val) : x = y := by
+   cases x <;> cases y; simp_all
+
+@[simp] def Vec.append {n m : Nat} (v : Vec α n) (w : Vec α m) : Vec α (n + m) :=
+  ⟨v.val ++ w.val, by simp [v.property, w.property]⟩
+
+instance : HAppend (Vec α n) (Vec α m) (Vec α (n+m)) where
+   hAppend xs ys := Vec.append xs ys
+
+@[simp] def Vec.push {n : Nat} (v : Vec α n) (x : α): Vec α (n+1) :=
+  ⟨v.val ++ [x], by simp [v.property]⟩
+
+-- Support array-like access st[i]
+@[simp] instance GetElem_Vec : GetElem (Vec α n) Nat α Vec.inBounds where
+  getElem := Vec.get
+
+-- Extensionality for fixed-length lists
+@[ext] def Vec.ext {n : Nat} (v w : Vec α n) (h: ∀(i : Nat), (h : i < n) → v[i] = w[i]) : v = w := by
+  apply Subtype.eq
+  apply List.ext_get <;> simp [v.property, w.property]
+  intros
+  simp_all [getElem, GetElem_Vec, Vec.get]
+  done
+
+def Vec.ext' {n : Nat} (v w : Vec α n) : (v=w) <-> (∀(i : Nat), (h : i < n) → v[i] = w[i]) := by
+  apply Iff.intro
+  · simp_all
+  · intros
+    ext
+    simp_all
+  done

From bce29291bbccf95a521a97ca58e57f68a95f8832 Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 26 Apr 2024 06:13:52 +0100
Subject: [PATCH 4/5] Add some helper lemmas for BitVec

`simp` rules for `extractLsb'` and `cast` / `append`

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 Arm/BitVec.lean | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/Arm/BitVec.lean b/Arm/BitVec.lean
index 76e5d389..bb4d1f67 100644
--- a/Arm/BitVec.lean
+++ b/Arm/BitVec.lean
@@ -391,6 +391,26 @@ protected theorem truncate_to_lsb_of_append (m n : Nat) (x : BitVec m) (y : BitV
   truncate n (x ++ y) = y := by
   simp only [truncate_append, Nat.le_refl, ↓reduceDite, zeroExtend_eq]
 
+@[simp] theorem extractLsb'_cast {k l: Nat} (e: k=l) (lo : Nat) (x : BitVec n):
+   BitVec.cast e (extractLsb' lo k x) = extractLsb' lo l x := by cases e; simp
+
+@[simp] theorem extractLsb'_cast' {k: Nat} (e: n=m) (lo : Nat) (x : BitVec n):
+   (extractLsb' lo k (BitVec.cast e x)) = extractLsb' lo k x := by cases e; simp
+
+@[simp] theorem getLsb'_extract (lo k : Nat) (x : BitVec n) (i : Nat) :
+    getLsb (extractLsb' lo k x) i = (i < k && getLsb x (lo+i)) := by
+  unfold getLsb
+  simp [Nat.lt_succ]
+
+@[simp] theorem extractLsb'_of_append {x : BitVec w} {y : BitVec v} :
+    (x ++ y).extractLsb' v w = x := by
+  apply eq_of_getLsb_eq
+  intro i
+  have k: ¬(v + i) < v := by omega
+  have k': v + i - v = i := by omega
+  simp [getLsb'_extract, k, k']
+  done
+
 ----------------------------------------------------------------------
 
 /- Bitvector pattern component syntax category, originally written by

From e262ae63d5d53feee665b016b55ea6a026df14bb Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Fri, 26 Apr 2024 06:17:29 +0100
Subject: [PATCH 5/5] Start FIPS-aligned model of AES specification

This commit adds `AESSpec.lean` which aims to model the AES
specification in a way that's as close to the FIPS specification
as possible (the existing specification in `AESArm.lean` is more
aligned to the Arm ASL).

This is work in progress. For the moment, we only have definitions
of the AES state as a bitvector, byte vector, or byte array, as well as
conversions between them.

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 Specs/AESSpec.lean | 182 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 182 insertions(+)
 create mode 100644 Specs/AESSpec.lean

diff --git a/Specs/AESSpec.lean b/Specs/AESSpec.lean
new file mode 100644
index 00000000..7d27a0c1
--- /dev/null
+++ b/Specs/AESSpec.lean
@@ -0,0 +1,182 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Hanno Becker
+-/
+
+import Arm.BitVec
+import Arm.Vec
+
+import Lean
+open Lean Meta
+
+open BitVec
+
+/-
+  The goal of this file is to provide a model of the AES specification
+  that is as close as possible to the original FIPS document.
+
+  TODO: This is work in progress. For the moment, we only have definitions
+    of the AES state as a bitvector, byte vector, or byte array, as well as
+    conversions between them.
+-/
+
+/-
+   3.3 Indexing of Byte Sequences
+-/
+
+abbrev Byte : Type := BitVec 8
+
+def bitvec_split_byte {k : Nat} (invec : BitVec (8*k)) : Byte × (BitVec (8*(k-1))) :=
+  let byte : BitVec 8 := invec.truncate 8
+  let remainder := BitVec.extractLsb' 8 (8*(k-1)) invec
+  (byte, remainder)
+
+theorem bitvec_split_byte_append {k : Nat} (v : BitVec (8*(k-1))) (b : Byte) (h : 8 * (k-1) + 8 = 8 * k) :
+    bitvec_split_byte (BitVec.cast h (v ++ b)) = (b, v) := by simp [bitvec_split_byte]
+
+theorem bitvec_split_byte_append' {k : Nat} (v : BitVec (8*(k+1))):
+   let (b, v') := bitvec_split_byte v
+   v' ++ b = v := by
+  apply eq_of_getLsb_eq
+  intro i
+  simp [getLsb_append]
+  by_cases h: ((i : Nat) < 8)
+  · simp [h]
+  · simp [h]
+    have lt: (i : Nat) - 8 < 8*k := by simp_all; omega
+    have e : (8 + ((i : Nat) - 8)) = i := by simp_all
+    simp [e, lt]
+    done
+  done
+
+-- Splitting a bitvector of 8*k entries into bytes, little endian
+def bitvec_to_byte_seq (k : Nat) (invec: BitVec (8*k)) : Vec Byte k :=
+  if k_gt_0: (k > 0) then
+    let (byte, remainder) := bitvec_split_byte invec
+    have h: k - 1 + 1 = k := by omega
+    h ▸ Vec.cons byte (bitvec_to_byte_seq (k-1) remainder)
+ else
+    have h: k = 0 := by omega
+    h ▸ Vec.empty
+
+-- Concatenating a little endian byte sequence into a bitvector
+def byte_seq_to_bitvec: Vec Byte k → BitVec (8*k)
+  | ⟨[], h⟩ =>
+      let h' : 0 = 8*k := by simp_all; omega
+      BitVec.cast h' (BitVec.ofNat 0 0)
+  | ⟨a :: as, h⟩ =>
+      let g : 8 * (k - 1) + 8 = 8 * k := by simp [←h]; omega
+      BitVec.cast g (byte_seq_to_bitvec ⟨as, by simp[←h]⟩ ++ a)
+
+@[simp]
+theorem byte_seq_to_bitvec_cons:
+  byte_seq_to_bitvec (Vec.cons byte v) = (byte_seq_to_bitvec v) ++ byte := by
+    cases v; simp only [Vec.cons, byte_seq_to_bitvec]; rfl
+
+-- Example
+def example_bitvec : BitVec 128 := 0x0102030405060708090a0b0c0d0e0f#128
+def example_byteseq  := bitvec_to_byte_seq 16 example_bitvec
+def example_bitvec'  := byte_seq_to_bitvec example_byteseq
+def example_byteseq' := bitvec_to_byte_seq 16 example_bitvec'
+
+#eval example_bitvec
+#eval example_bitvec'
+#eval example_byteseq
+#eval example_byteseq'
+
+@[simp]
+def Vec_0_singleton {k : Nat}: (x y : Vec α k) → k = 0 → x = y
+  | ⟨[], _⟩, ⟨[], _⟩, _ => by simp
+  | ⟨_ :: _, hx⟩, _, h => by simp at hx; omega
+  | _, ⟨_ :: _, hy⟩, h => by simp at hy; omega
+
+theorem bitvec_to_byte_seq_invA: ∀(x : Vec Byte k), bitvec_to_byte_seq k (byte_seq_to_bitvec x) = x
+ | ⟨[], h⟩ => by
+      have h: k = 0 := by simp_all
+      rw [byte_seq_to_bitvec]
+      rw [bitvec_to_byte_seq]
+      simp_all [Vec.empty]
+      apply Vec_0_singleton
+      assumption
+ | ⟨x :: xs, h⟩ => by
+    have k_gt_0: k > 0 := by simp_all; omega
+    let h' : xs.length = k - 1 := by simp_all; omega
+    let x' : Vec Byte (k-1) := ⟨xs, h'⟩
+    let h : bitvec_to_byte_seq (k - 1) (byte_seq_to_bitvec x') = x' :=
+      bitvec_to_byte_seq_invA x'
+    rw [byte_seq_to_bitvec, bitvec_to_byte_seq]
+    simp_all [bitvec_split_byte_append, Vec.cons]
+    apply Vec.ext''; simp
+    done
+
+theorem bitvec_to_byte_seq_invB: ∀(k : Nat) (x : BitVec (8*k)), byte_seq_to_bitvec (bitvec_to_byte_seq k x) = x
+ | 0, x => by simp
+ | k+1, x => by
+    rw [bitvec_to_byte_seq]
+    have kp1_gt_0: k + 1 > 0 := by simp
+    simp [kp1_gt_0]
+    let x' := (bitvec_split_byte x).snd
+    have t' : byte_seq_to_bitvec (bitvec_to_byte_seq (k + 1 - 1) x') = x' := by
+      exact (bitvec_to_byte_seq_invB k _)
+    simp only [t', bitvec_split_byte_append' x]
+
+/-
+  3.4 The state
+-/
+
+/- We deal with three different presentations of the AES state:
+   1/ As a bitvector of length 128
+   2/ As a byte vector length 16
+   3/ As a 4x4 grid of bytes
+-/
+
+/-- Length 16 vectors <-> 4x4 arrays -/
+
+abbrev AESStateBitVec := BitVec 128
+abbrev AESStateVec := Vec Byte 16
+abbrev AESStateArr := Vec (Vec Byte 4) 4
+
+def AES_State_BitVec_to_Vec (x: AESStateBitVec): AESStateVec := bitvec_to_byte_seq 16 x
+
+def AESStateVec_to_Arr (x : AESStateVec): AESStateArr :=
+  ⟨[⟨[x[0], x[4], x[8] , x[12]], by simp⟩,
+    ⟨[x[1], x[5], x[9] , x[13]], by simp⟩,
+    ⟨[x[2], x[6], x[10], x[14]], by simp⟩,
+    ⟨[x[3], x[7], x[11], x[15]], by simp⟩], by simp⟩
+
+def AESStateArr_to_Vec (x : AESStateArr) : AESStateVec :=
+  ⟨[ x[0][0], x[1][0], x[2][0], x[3][0],
+     x[0][1], x[1][1], x[2][1], x[3][1],
+     x[0][2], x[1][2], x[2][2], x[3][2],
+     x[0][3], x[1][3], x[2][3], x[3][3] ], by simp⟩
+
+theorem lt_succ_iff_lt_or_eq {n m : Nat} (h : 0 < m) : (n < m) <-> (n < m-1 ∨ n = m-1) := by
+  have hs: m = Nat.succ (m - 1) := by omega
+  rw [hs, Nat.lt_succ_iff_lt_or_eq]; simp
+  done
+
+theorem AESStateVec_to_Arr' (x : AESStateVec) (i j : Nat) (hi: i < 4) (hj: j < 4):
+  have h: 4*j + i < 16 := by omega
+  (AESStateVec_to_Arr x)[i][j] = x[4*j + i] := by
+     simp [lt_succ_iff_lt_or_eq, Nat.not_lt_zero] at hi
+     simp [lt_succ_iff_lt_or_eq, Nat.not_lt_zero] at hj
+     elim Or.elim <;> try simp_all <;> simp [AESStateVec_to_Arr, Vec.get, GetElem.getElem, GetElem_Vec, List.get]
+     done
+
+def AESStateArr_to_Vec_invA (x : AESStateVec): (AESStateArr_to_Vec (AESStateVec_to_Arr x)) = x := by
+  simp [AESStateArr_to_Vec, AESStateVec_to_Arr']
+  ext
+  rename_i i hi
+  simp [lt_succ_iff_lt_or_eq, Nat.not_lt_zero] at hi
+  elim Or.elim <;> try simp_all <;> simp [AESStateVec_to_Arr, Vec.get, GetElem.getElem, GetElem_Vec, List.get]
+  done
+
+def AESStateArr_to_Vec_invB (x : AESStateArr): (AESStateVec_to_Arr (AESStateArr_to_Vec x)) = x := by
+  simp [AESStateArr_to_Vec, AESStateVec_to_Arr']
+  ext
+  rename_i i hi j hj
+  simp [lt_succ_iff_lt_or_eq, Nat.not_lt_zero] at hi
+  simp [lt_succ_iff_lt_or_eq, Nat.not_lt_zero] at hj
+  elim Or.elim <;> try simp_all <;> simp [AESStateVec_to_Arr, Vec.get, GetElem.getElem, GetElem_Vec, List.get]
+  done