val matrix_add_identity (#c #eq: _) (add: CE.cm c eq) (#m #n: pos) (mx: matrix c m n)
    : Lemma (matrix_eq_fun eq (matrix_add add (matrix_add_zero add m n) mx) mx)
let matrix_add_identity #c #eq (add: CE.cm c eq) (#m #n: pos) (mx: matrix c m n)
  : Lemma (matrix_add add (matrix_add_zero add m n) mx `matrix_eq_fun eq` mx) =
  matrix_equiv_from_proof eq (matrix_add add (matrix_add_zero add m n) mx) mx
                          (fun i j -> add.identity (ijth mx i j))
val matrix_add_identity (#c #eq: _) (add: CE.cm c eq) (#m #n: pos) (mx: matrix c m n)
    : Lemma (matrix_eq_fun eq (matrix_add add (matrix_add_zero add m n) mx) mx)
let matrix_add_identity #c #eq (add: CE.cm c eq) (#m: pos) (#n: pos) (mx: matrix c m n)
    : Lemma (matrix_eq_fun eq (matrix_add add (matrix_add_zero add m n) mx) mx) =
matrix_equiv_from_proof eq
  (matrix_add add (matrix_add_zero add m n) mx)
  mx
  (fun i j -> add.identity (ijth mx i j))
FStar.Matrix.fst.checked
prims.fst.checked FStar.Seq.Properties.fsti.checked FStar.Seq.Permutation.fsti.checked FStar.Seq.Equiv.fsti.checked FStar.Seq.Base.fsti.checked FStar.Pervasives.Native.fst.checked FStar.Pervasives.fsti.checked FStar.Mul.fst.checked FStar.Math.Lemmas.fst.checked FStar.IntegerIntervals.fst.checked FStar.Classical.fsti.checked FStar.Algebra.CommMonoid.Fold.fsti.checked FStar.Algebra.CommMonoid.Equiv.fst.checked
FStar.Matrix.fst
FStar.Algebra.CommMonoid.Equiv.equiv FStar.Algebra.CommMonoid.Equiv.cm Prims.pos FStar.Matrix.matrix FStar.Matrix.matrix_equiv_from_proof FStar.Matrix.matrix_add FStar.Matrix.matrix_add_zero FStar.IntegerIntervals.under FStar.Algebra.CommMonoid.Equiv.__proj__CM__item__identity FStar.Matrix.ijth Prims.unit Prims.l_True Prims.squash FStar.Matrix.matrix_eq_fun Prims.Nil FStar.Pervasives.pattern
(*
   Copyright 2022 Microsoft Research
   
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
   
       http://www.apache.org/licenses/LICENSE-2.0
       
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

   Author: A. Rozanov
*)

(*
   In this module we provide basic definitions to work with matrices via
   seqs, and define transpose transform together with theorems that assert
   matrix fold equality of original and transposed matrices.
*)


module FStar.Matrix

module CE = FStar.Algebra.CommMonoid.Equiv
module CF = FStar.Algebra.CommMonoid.Fold
module SP = FStar.Seq.Permutation
module SB = FStar.Seq.Base
module SProp = FStar.Seq.Properties
module ML = FStar.Math.Lemmas

open FStar.IntegerIntervals   
open FStar.Mul
open FStar.Seq.Equiv

(* 
   A little glossary that might help reading this file
   We don't list common terms like associativity and reflexivity.

   lhs, rhs               left hand side, right hand side
   liat                   subsequence of all elements except the last (tail read backwards)
   snoc                   construction of sequence from a pair (liat, last) (cons read backwards)
   un_snoc                decomposition of sequence into a pair (liat, last)
   foldm                  sum or product of all elements in a sequence using given CommMonoid
   foldm_snoc             recursively defined sum/product of a sequence, starting from the last element
   congruence             respect of equivalence ( = ) by a binary operation ( * ), a=b ==> a*x = b*x
   unit                   identity element (xu=x, ux=x) (not to be confused with invertible elements)
*)

type matrix c m n = z:SB.seq c { SB.length z = m*n }
 
let seq_of_matrix #c #m #n mx = mx

let ijth #c #m #n mx i j = SB.index mx (get_ij m n i j)

let ijth_lemma #c #m #n mx i j 
 : Lemma (ijth mx i j == SB.index (seq_of_matrix mx) (get_ij m n i j)) = ()

let matrix_of_seq #c m n s = s

let foldm #c #eq #m #n cm mx = SP.foldm_snoc cm mx
 
let matrix_fold_equals_fold_of_seq #c #eq #m #n cm mx
  : Lemma (ensures foldm cm mx `eq.eq` SP.foldm_snoc cm (seq_of_matrix mx)) [SMTPat(foldm cm mx)]
  = eq.reflexivity (foldm cm mx)

let matrix_fold_internal #c #eq #m #n (cm:CE.cm c eq) (mx: matrix c m n)
  : Lemma (ensures foldm cm mx == SP.foldm_snoc cm mx) = ()
     
(* A flattened matrix (seq) constructed from generator function
   Notice how the domains of both indices are strictly controlled. *)
let init #c (#m #n: pos) (generator: matrix_generator c m n)
  : matrix_of generator =  
  let mn = m * n in
  let generator_ij ij = generator (get_i m n ij) (get_j m n ij) in
  let flat_indices = indices_seq mn in  
  let result = SProp.map_seq generator_ij flat_indices in
  SProp.map_seq_len generator_ij flat_indices;
  assert (SB.length result == SB.length flat_indices);
  let aux (i: under m) (j: under n) 
    : Lemma (SB.index (SProp.map_seq generator_ij flat_indices) (get_ij m n i j) == generator i j) 
    = consistency_of_i_j m n i j;
      consistency_of_ij m n (get_ij m n i j);
      assert (generator_ij (get_ij m n i j) == generator i j);
      SProp.map_seq_index generator_ij flat_indices (get_ij m n i j) in
  let aux1 (ij: under mn) 
    : Lemma (SB.index (SProp.map_seq generator_ij flat_indices) ij == generator_ij ij) 
    = SProp.map_seq_index generator_ij flat_indices ij in 
  FStar.Classical.forall_intro aux1;
  FStar.Classical.forall_intro_2 aux;
  result
  
private let matrix_seq #c #m #n (gen: matrix_generator c m n) : (t:SB.seq c{ (SB.length t = (m*n)) /\
  (forall (i: under m) (j: under n). SB.index t (get_ij m n i j) == gen i j) /\ 
  (forall(ij: under (m*n)). SB.index t ij == gen (get_i m n ij) (get_j m n ij))
}) = init gen

(* This auxiliary lemma establishes the decomposition of the seq-matrix 
   into the concatenation of its first (m-1) rows and its last row (thus snoc)  *)
let matrix_append_snoc_lemma #c (#m #n: pos) (generator: matrix_generator c m n)
  : Lemma (matrix_seq generator == (SB.slice (matrix_seq generator) 0 ((m-1)*n))
                                   `SB.append`
                                   (SB.slice (matrix_seq generator) ((m-1)*n) (m*n))) 
  = SB.lemma_eq_elim (matrix_seq generator) 
                  (SB.append (SB.slice (matrix_seq generator) 0 ((m-1)*n))
                          (SB.slice (matrix_seq generator) ((m-1)*n) (m*n)))

let matrix_seq_decomposition_lemma  #c (#m:greater_than 1) (#n: pos) (generator: matrix_generator c m n)
  : Lemma ((matrix_seq generator) == 
          SB.append (matrix_seq #c #(m-1) #n generator)
                    (SB.slice (matrix_seq generator) ((m-1)*n) (m*n)))
  = SB.lemma_eq_elim (matrix_seq generator)
                  ((matrix_seq #c #(m-1) #n generator) `SB.append` 
                  (SB.slice (matrix_seq generator) ((m-1)*n) (m*n))) 

(* This auxiliary lemma establishes the equality of the fold of the entire matrix
   to the op of folds of (the submatrix of the first (m-1) rows) and (the last row). *) 
let matrix_fold_snoc_lemma #c #eq 
                           (#m: not_less_than 2) 
                           (#n: pos) 
                           (cm: CE.cm c eq) 
                           (generator: matrix_generator c m n)
  : Lemma (assert ((m-1)*n < m*n);
            SP.foldm_snoc cm (matrix_seq generator) `eq.eq` 
    cm.mult (SP.foldm_snoc cm (matrix_seq #c #(m-1) #n generator))
            (SP.foldm_snoc cm (SB.slice (matrix_seq #c #m #n generator) ((m-1)*n) (m*n))))
  = SB.lemma_eq_elim (matrix_seq generator)
                  ((matrix_seq #c #(m-1) #n generator) `SB.append` 
                  (SB.slice (matrix_seq generator) ((m-1)*n) (m*n)));    
    SP.foldm_snoc_append cm (matrix_seq #c #(m-1) #n generator) 
                         (SB.slice (matrix_seq generator) ((m-1)*n) (m*n)) 

(* 
   There are many auxiliary lemmas like this that are extracted because
   lemma_eq_elim invocations often impact verification speed more than
   one might expect they would.
*)
let matrix_submatrix_lemma #c (#m: not_less_than 2) (#n: pos)  
                           (generator: matrix_generator c m n)
  : Lemma ((matrix_seq generator) == (matrix_seq (fun (i:under(m-1)) (j:under n) -> generator i j) 
                                     `SB.append` SB.init n (generator (m-1))))
  = SB.lemma_eq_elim (matrix_seq (fun (i:under (m-1)) (j:under n) -> generator i j)) 
                     (matrix_seq #c #(m-1) #n generator);
    SB.lemma_eq_elim (SB.slice (matrix_seq generator) ((m-1)*n) (m*n)) 
                     (SB.init n (generator (m-1)));
    matrix_seq_decomposition_lemma generator

let matrix_seq_of_one_row_matrix #c #m #n (generator : matrix_generator c m n) 
  : Lemma (requires m==1)
          (ensures matrix_seq generator == (SB.init n (generator 0))) = 
  SB.lemma_eq_elim (matrix_seq generator) (SB.init n (generator 0))

let one_row_matrix_fold_aux #c #eq #m #n (cm:CE.cm c eq) (generator : matrix_generator c m n) : Lemma 
  (requires m=1)
  (ensures foldm cm (init generator) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i)))) /\ 
           SP.foldm_snoc cm (seq_of_matrix (init generator)) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i))))) =
  let lhs_seq = matrix_seq generator in
  let rhs_seq = SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i))) in
  let lhs = SP.foldm_snoc cm (matrix_seq generator) in
  let rhs = SP.foldm_snoc cm rhs_seq in
  SP.foldm_snoc_singleton cm (SP.foldm_snoc cm (SB.init n (generator 0)));
  SB.lemma_eq_elim (SB.create 1 (SP.foldm_snoc cm (SB.init n (generator 0))))
    (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i))));
  matrix_seq_of_one_row_matrix generator;
  eq.symmetry rhs lhs


let fold_of_subgen_aux #c #eq (#m:pos{m>1}) #n (cm: CE.cm c eq) (gen: matrix_generator c m n) (subgen: matrix_generator c (m-1) n) : Lemma
  (requires subgen == (fun (i: under (m-1)) (j: under n) -> gen i j))
  (ensures forall (i: under (m-1)). SP.foldm_snoc cm (SB.init n (subgen i)) ==
                               SP.foldm_snoc cm (SB.init n (gen i))) =         
  let aux_pat (i: under (m-1)) : Lemma (SP.foldm_snoc cm (SB.init n (subgen i)) 
                                     == SP.foldm_snoc cm (SB.init n (gen i))) = 
  SB.lemma_eq_elim (SB.init n (subgen i)) (SB.init n (gen i)) in
  Classical.forall_intro aux_pat          

let arithm_aux (m: pos{m>1}) (n: pos) : Lemma ((m-1)*n < m*n) = ()

let terminal_case_aux #c #eq (#p:pos{p=1}) #n (cm:CE.cm c eq) (generator: matrix_generator c p n) (m: pos{m<=p}) : Lemma 
  (ensures SP.foldm_snoc cm (SB.slice (seq_of_matrix (init generator)) 0 (m*n)) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun (i:under m) -> SP.foldm_snoc cm (SB.init n (generator i)))))
  = one_row_matrix_fold_aux cm generator

#push-options "--ifuel 0 --fuel 1 --z3rlimit 10"
let terminal_case_two_aux #c #eq (#p:pos) #n (cm:CE.cm c eq) (generator: matrix_generator c p n) (m: pos{m=1}) : Lemma 
  (ensures SP.foldm_snoc cm (SB.slice (seq_of_matrix (init generator)) 0 (m*n)) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun (i:under m) -> SP.foldm_snoc cm (SB.init n (generator i)))))
  = 
   SP.foldm_snoc_singleton cm (SP.foldm_snoc cm (SB.init n (generator 0)));
   assert (SP.foldm_snoc cm (SB.init m (fun (i:under m) -> SP.foldm_snoc cm (SB.init n (generator i)))) `eq.eq`
           SP.foldm_snoc cm (SB.init n (generator 0)));
   let line = SB.init n (generator 0) in
   let slice = SB.slice (matrix_seq generator) 0 n in
   let aux (ij: under n) : Lemma (SB.index slice ij == SB.index line ij) = 
     Math.Lemmas.small_div ij n;
     Math.Lemmas.small_mod ij n 
   in Classical.forall_intro aux;
   SB.lemma_eq_elim line slice;        
   eq.symmetry (SP.foldm_snoc cm (SB.init m (fun (i:under m) -> SP.foldm_snoc cm (SB.init n (generator i)))))
               (SP.foldm_snoc cm line) 
#pop-options

let liat_equals_init #c (m:pos) (gen: under m -> c)
  : Lemma (fst (SProp.un_snoc (SB.init m gen)) == SB.init (m-1) gen) = 
  SB.lemma_eq_elim (fst (SProp.un_snoc (SB.init m gen))) (SB.init (m-1) gen)

let math_aux (m n: pos) (j: under n) : Lemma (j+((m-1)*n) < m*n) = ()

let math_aux_2 (m n: pos) (j: under n) : Lemma (get_j m n (j+(m-1)*n) == j) 
  = 
  Math.Lemmas.modulo_addition_lemma j n (m-1);
  Math.Lemmas.small_mod j n 
  
let math_aux_3 (m n: pos) (j: under n) : Lemma (get_i m n (j+(m-1)*n) == (m-1)) 
  = 
  Math.Lemmas.division_addition_lemma j n (m-1);
  Math.Lemmas.small_div j n 

let math_aux_4 (m n: pos) (j: under n) : Lemma ((j+((m-1)*n)) - ((m-1)*n) == j) = ()

let seq_eq_from_member_eq #c (n: pos) (p q: (z:SB.seq c{SB.length z=n}))
                          (proof: (i: under n) -> Lemma (SB.index p i == SB.index q i))
  : Lemma (p == q) = 
  Classical.forall_intro proof;
  SB.lemma_eq_elim p q

let math_wut_lemma (x: pos) : Lemma (requires x>1) (ensures x-1 > 0) = ()

(* This proof used to be very unstable, so I rewrote it with as much precision
   and control over lambdas as possible. 

   I also left intact some trivial auxiliaries and the quake option 
   in order to catch regressions the moment they happen instead of several 
   releases later -- Alex *)
#push-options "--ifuel 0 --fuel 0 --z3rlimit 15"
#restart-solver
let rec matrix_fold_equals_double_fold #c #eq (#p:pos) #n (cm:CE.cm c eq) 
                                       (generator: matrix_generator c p n) (m: pos{m<=p}) 
  : Lemma (ensures SP.foldm_snoc cm (SB.slice (seq_of_matrix (init generator)) 0 (m*n)) `eq.eq`
                   SP.foldm_snoc cm (SB.init m (fun (i: under m) -> SP.foldm_snoc cm (SB.init n (generator i)))))
          (decreases m) =
  if p=1 then terminal_case_aux cm generator m 
  else if m=1 then terminal_case_two_aux cm generator m
  else 
    let lhs_seq = (SB.slice (matrix_seq generator) 0 (m*n)) in
    let rhs_seq_gen = fun (i: under m) -> SP.foldm_snoc cm (SB.init n (generator i)) in
    let rhs_seq_subgen = fun (i: under (m-1)) -> SP.foldm_snoc cm (SB.init n (generator i)) in
    let rhs_seq = SB.init m rhs_seq_gen in
    let lhs = SP.foldm_snoc cm lhs_seq in
    let rhs = SP.foldm_snoc cm rhs_seq in
    let matrix = lhs_seq in 
    let submatrix = SB.slice (matrix_seq generator) 0 ((m-1)*n) in
    let last_row = SB.slice (matrix_seq generator) ((m-1)*n) (m*n) in  
    SB.lemma_len_slice (matrix_seq generator) ((m-1)*n) (m*n);
    assert (SB.length last_row = n);
    SB.lemma_eq_elim matrix (SB.append submatrix last_row); 
    SP.foldm_snoc_append cm submatrix last_row; 
    matrix_fold_equals_double_fold #c #eq #p #n cm generator (m-1);  
    SB.lemma_eq_elim (SB.init (m-1) rhs_seq_gen)
                     (SB.init (m-1) rhs_seq_subgen); 
    let aux (j: under n) : Lemma (SB.index last_row j == generator (m-1) j) =
      SB.lemma_index_app2 submatrix last_row (j+((m-1)*n));
      math_aux_2 m n j;
      math_aux_3 m n j;
      math_aux_4 m n j;
      () in Classical.forall_intro aux; 
    let rhs_liat, rhs_last = SProp.un_snoc rhs_seq in 
    let rhs_last_seq = SB.init n (generator (m-1)) in
    liat_equals_init m rhs_seq_gen;
    SP.foldm_snoc_decomposition cm rhs_seq;    
    let aux_2 (j: under n) : Lemma (SB.index last_row j == SB.index rhs_last_seq j) = () in
    seq_eq_from_member_eq n last_row rhs_last_seq aux_2;
    SB.lemma_eq_elim rhs_liat (SB.init (m-1) rhs_seq_gen);
    cm.commutativity (SP.foldm_snoc cm submatrix) (SP.foldm_snoc cm last_row);
    eq.transitivity lhs (SP.foldm_snoc cm submatrix `cm.mult` SP.foldm_snoc cm last_row)
                        (SP.foldm_snoc cm last_row `cm.mult` SP.foldm_snoc cm submatrix);
    eq.reflexivity (SP.foldm_snoc cm last_row);
    cm.congruence (SP.foldm_snoc cm last_row) (SP.foldm_snoc cm submatrix)
                  (SP.foldm_snoc cm last_row) (SP.foldm_snoc cm (SB.init (m-1) rhs_seq_subgen));
    eq.transitivity lhs (SP.foldm_snoc cm last_row `cm.mult` SP.foldm_snoc cm submatrix) rhs 
#pop-options

let matrix_fold_equals_fold_of_seq_folds #c #eq #m #n cm generator : Lemma 
  (ensures foldm cm (init generator) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i)))) /\ 
           SP.foldm_snoc cm (seq_of_matrix (init generator)) `eq.eq`
           SP.foldm_snoc cm (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i))))) =  
  matrix_fold_equals_double_fold cm generator m;
  assert ((SB.slice (seq_of_matrix (init generator)) 0 (m*n)) == seq_of_matrix (init generator));
  SB.lemma_eq_elim (SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i))))
                   (SB.init m (fun (i: under m) -> SP.foldm_snoc cm (SB.init n (generator i))));
  assert ((SB.init m (fun i -> SP.foldm_snoc cm (SB.init n (generator i)))) ==
          (SB.init m (fun (i: under m) -> SP.foldm_snoc cm (SB.init n (generator i)))));
()

(* This auxiliary lemma shows that the fold of the last line of a matrix
   is equal to the corresponding fold of the generator function *) 
let matrix_last_line_equals_gen_fold #c #eq 
                                        (#m #n: pos)  
                                        (cm: CE.cm c eq) 
                                        (generator: matrix_generator c m n) 
  : Lemma (SP.foldm_snoc cm (SB.slice (matrix_seq generator) ((m-1)*n) (m*n))
           `eq.eq` CF.fold cm 0 (n-1) (generator (m-1))) = 
  let slice = SB.slice #c in
  let foldm_snoc = SP.foldm_snoc #c #eq in
  assert (matrix_seq generator == seq_of_matrix (init generator));
  let init = SB.init #c in
  let lemma_eq_elim = SB.lemma_eq_elim #c in 
  lemma_eq_elim (slice (matrix_seq generator) ((m-1)*n) (m*n)) 
                (init n (generator (m-1))); 
  let g : ifrom_ito 0 (n-1) -> c = generator (m-1) in 
  CF.fold_equals_seq_foldm cm 0 (n-1) g;
  let gen = CF.init_func_from_expr g 0 (n-1) in
  eq.reflexivity (foldm_snoc cm (init (closed_interval_size 0 (n-1)) gen));  
  lemma_eq_elim (slice (matrix_seq generator) ((m-1)*n) (m*n))    
                (init (closed_interval_size 0 (n-1)) gen); 
  eq.symmetry (CF.fold cm 0 (n-1) (generator (m-1)))
              (foldm_snoc cm (init (closed_interval_size 0 (n-1)) gen)); 
  eq.transitivity (foldm_snoc cm (slice (matrix_seq generator) ((m-1)*n) (m*n)))
                  (foldm_snoc cm (init (closed_interval_size 0 (n-1)) gen))
                  (CF.fold cm 0 (n-1) (generator (m-1))) 
 
(* This lemma proves that a matrix fold is the same thing as double-fold of
   its generator function against full indices ranges *)
#push-options "--ifuel 0 --fuel 0"
let rec matrix_fold_aux #c #eq // lemma needed for precise generator domain control
                           (#gen_m #gen_n: pos) // full generator domain
                           (cm: CE.cm c eq) 
                           (m: ifrom_ito 1 gen_m) (n: ifrom_ito 1 gen_n) //subdomain
                           (generator: matrix_generator c gen_m gen_n)
  : Lemma (ensures SP.foldm_snoc cm (matrix_seq #c #m #n generator) `eq.eq` 
                   CF.fold cm 0 (m-1) (fun (i: under m) -> CF.fold cm 0 (n-1) (generator i)))
          (decreases m) = 
  Classical.forall_intro_2 (ijth_lemma (init generator));
  let slice = SB.slice #c in
  let foldm_snoc = SP.foldm_snoc #c #eq in
  let lemma_eq_elim = SB.lemma_eq_elim #c in
  if m = 1 then begin
    matrix_fold_equals_fold_of_seq cm (init generator);
    matrix_last_line_equals_gen_fold #c #eq #m #n cm generator;  
    CF.fold_singleton_lemma cm 0 (fun (i:under m) -> CF.fold cm 0 (n-1) (generator i));
    assert (CF.fold cm 0 (m-1) (fun (i: under m) -> CF.fold cm 0 (n-1) (generator i)) 
            == CF.fold cm 0 (n-1) (generator 0)) 
  end else begin     
    Classical.forall_intro_3 (Classical.move_requires_3 eq.transitivity);  
    matrix_fold_aux cm (m-1) n generator;    
    let outer_func (i: under m) = CF.fold cm 0 (n-1) (generator i) in
    let outer_func_on_subdomain (i: under (m-1)) = CF.fold cm 0 (n-1) (generator i) in
    CF.fold_equality cm 0 (m-2) outer_func_on_subdomain outer_func;
    CF.fold_snoc_decomposition cm 0 (m-1) outer_func;  
    matrix_fold_snoc_lemma #c #eq #m #n cm generator;
    matrix_last_line_equals_gen_fold #c #eq #m #n cm generator;
    cm.congruence (foldm_snoc cm (matrix_seq #c #(m-1) #n generator))
                  (foldm_snoc cm (slice (matrix_seq #c #m #n generator) ((m-1)*n) (m*n)))
                  (CF.fold cm 0 (m-2) outer_func)
                  (CF.fold cm 0 (n-1) (generator (m-1))) 
  end  
#pop-options

(* This lemma establishes that the fold of a matrix is equal to
   nested Algebra.CommMonoid.Fold.fold over the matrix generator *)
let matrix_fold_equals_func_double_fold #c #eq #m #n cm generator
  : Lemma (foldm cm (init generator) `eq.eq` 
           CF.fold cm 0 (m-1) (fun (i:under m) -> CF.fold cm 0 (n-1) (generator i))) 
  = matrix_fold_aux cm m n generator

(* This function provides the transposed matrix generator, with indices swapped
   Notice how the forall property of the result function is happily proved 
   automatically by z3 :) *)
let transposed_matrix_gen #c #m #n (generator: matrix_generator c m n) 
  : (f: matrix_generator c n m { forall i j. f j i == generator i j }) 
  = fun j i -> generator i j

(* This lemma shows that the transposed matrix is 
   a permutation of the original one *)
let matrix_transpose_is_permutation #c #m #n generator
  : Lemma (SP.is_permutation (seq_of_matrix (init generator))
                          (seq_of_matrix (init (transposed_matrix_gen generator)))
                          (transpose_ji m n)) = 
  let matrix_transposed_eq_lemma #c (#m #n: pos) 
                                        (gen: matrix_generator c m n) 
                                        (ij: under (m*n)) 
    : Lemma (SB.index (seq_of_matrix (init gen)) ij ==
             SB.index (seq_of_matrix (init (transposed_matrix_gen gen))) (transpose_ji m n ij)) 
    = 
      ijth_lemma (init gen) (get_i m n ij) (get_j m n ij);
      ijth_lemma (init (transposed_matrix_gen gen)) 
                 (get_i n m (transpose_ji m n ij)) 
                 (get_j n m (transpose_ji m n ij));
    () in 
  let transpose_inequality_lemma (m n: pos) (ij: under (m*n)) (kl: under (n*m)) 
    : Lemma (requires kl <> ij) (ensures transpose_ji m n ij <> transpose_ji m n kl) = 
      dual_indices m n ij;
      dual_indices m n kl in
  Classical.forall_intro (matrix_transposed_eq_lemma generator); 
  Classical.forall_intro_2 (Classical.move_requires_2 
                           (transpose_inequality_lemma m n));
  SP.reveal_is_permutation (seq_of_matrix (init generator))
                           (seq_of_matrix (init (transposed_matrix_gen generator)))
                           (transpose_ji m n) 

(* Fold over matrix equals fold over transposed matrix *)
let matrix_fold_equals_fold_of_transpose #c #eq #m #n
                                         (cm: CE.cm c eq) 
                                         (gen: matrix_generator c m n)
  : Lemma (foldm cm (init gen) `eq.eq`
           foldm cm (init (transposed_matrix_gen gen))) = 
  let matrix_seq #c #m #n (g: matrix_generator c m n) = (seq_of_matrix (init g)) in
  let matrix_mn = matrix_seq gen in
  let matrix_nm = matrix_seq (transposed_matrix_gen gen) in
  matrix_transpose_is_permutation gen;
  SP.foldm_snoc_perm cm (matrix_seq gen)
                     (matrix_seq (transposed_matrix_gen gen))
                     (transpose_ji m n);
  matrix_fold_equals_fold_of_seq cm (init gen);
  matrix_fold_equals_fold_of_seq cm (init (transposed_matrix_gen gen));
  eq.symmetry (foldm cm (init (transposed_matrix_gen gen))) 
              (SP.foldm_snoc cm (matrix_seq (transposed_matrix_gen gen)));
  eq.transitivity (foldm cm (init gen)) (SP.foldm_snoc cm (matrix_seq gen))
                                        (SP.foldm_snoc cm (matrix_seq (transposed_matrix_gen gen)));
  eq.transitivity (foldm cm (init gen)) (SP.foldm_snoc cm (matrix_seq (transposed_matrix_gen gen)))
                  (foldm cm (init (transposed_matrix_gen gen))) 

let matrix_eq_fun #c (#m #n: pos) (eq: CE.equiv c) (ma mb: matrix c m n) =   
  eq_of_seq eq (seq_of_matrix ma) (seq_of_matrix mb)

(*
   Matrix equivalence, defined as element-wise equivalence of its underlying 
   flattened sequence, is constructed trivially from the element equivalence
   and the lemmas defined above.
*)
let matrix_equiv #c (eq: CE.equiv c) (m n: pos) : CE.equiv (matrix c m n) = 
  CE.EQ (matrix_eq_fun eq)
        (fun m -> eq_of_seq_reflexivity eq (seq_of_matrix m))
        (fun ma mb -> eq_of_seq_symmetry eq (seq_of_matrix ma) (seq_of_matrix mb))
        (fun ma mb mc -> eq_of_seq_transitivity eq (seq_of_matrix ma) (seq_of_matrix mb) (seq_of_matrix mc))

(* Equivalence of matrices means equivalence of all corresponding elements *)
let matrix_equiv_ijth #c (#m #n: pos) (eq: CE.equiv c) (ma mb: matrix c m n) (i: under m) (j: under n)
  : Lemma (requires (matrix_equiv eq m n).eq ma mb) (ensures ijth ma i j `eq.eq` ijth mb i j) = 
  eq_of_seq_element_equality eq (seq_of_matrix ma) (seq_of_matrix mb)

(* Equivalence of all corresponding elements means equivalence of matrices *)
let matrix_equiv_from_element_eq #c (#m #n: pos) (eq: CE.equiv c) (ma mb: matrix c m n)
  : Lemma (requires (forall (i: under m) (j: under n). ijth ma i j `eq.eq` ijth mb i j))
          (ensures matrix_eq_fun eq ma mb) = 
  assert (SB.length (seq_of_matrix ma) = SB.length (seq_of_matrix mb));
  let s1 = seq_of_matrix ma in
  let s2 = seq_of_matrix mb in
  assert (forall (ij: under (m*n)). SB.index s1 ij == ijth ma (get_i m n ij) (get_j m n ij));
  assert (forall (ij: under (m*n)). SB.index s2 ij == ijth mb (get_i m n ij) (get_j m n ij));
  assert (forall (ij: under (m*n)). SB.index s1 ij `eq.eq` SB.index s2 ij);  
  eq_of_seq_from_element_equality eq (seq_of_matrix ma) (seq_of_matrix mb)

(* We construct addition CommMonoid from the following definitions *)
let matrix_add_is_associative #c #eq #m #n (add: CE.cm c eq) (ma mb mc: matrix c m n)
  : Lemma (matrix_add add (matrix_add add ma mb) mc `(matrix_equiv eq m n).eq` 
           matrix_add add ma (matrix_add add mb mc)) =  
  matrix_equiv_from_proof eq 
    (matrix_add add (matrix_add add ma mb) mc)
    (matrix_add add ma (matrix_add add mb mc))  
    (fun i j -> add.associativity (ijth ma i j) (ijth mb i j) (ijth mc i j))

let matrix_add_is_commutative #c #eq (#m #n: pos) (add: CE.cm c eq) (ma mb: matrix c m n)
  : Lemma (matrix_add add ma mb `(matrix_equiv eq m n).eq` matrix_add add mb ma) = 
  matrix_equiv_from_proof eq (matrix_add add ma mb) (matrix_add add mb ma)
    (fun i j -> add.commutativity (ijth ma i j) (ijth mb i j)) 

let matrix_add_congruence #c #eq (#m #n: pos) (add: CE.cm c eq) (ma mb mc md: matrix c m n)
  : Lemma (requires matrix_eq_fun eq ma mc /\ matrix_eq_fun eq mb md)
          (ensures matrix_add add ma mb `matrix_eq_fun eq` matrix_add add mc md) =
  matrix_equiv_from_proof eq (matrix_add add ma mb) (matrix_add add mc md)
                          (fun i j -> matrix_equiv_ijth eq ma mc i j;
                                   matrix_equiv_ijth eq mb md i j; 
                                   add.congruence (ijth ma i j) (ijth mb i j) 
                                                  (ijth mc i j) (ijth md i j)) 
 
let matrix_add_zero #c #eq (add: CE.cm c eq) (m n: pos) 
  : (z: matrix c m n { forall (i: under m) (j: under n). ijth z i j == add.unit }) 
  = matrix_of_seq m n (SB.create (m*n) add.unit)

let matrix_add_identity #c #eq (add: CE.cm c eq) (#m #n: pos) (mx: matrix c m n)
val matrix_add_identity (#c #eq: _) (add: CE.cm c eq) (#m #n: pos) (mx: matrix c m n)
    : Lemma (matrix_eq_fun eq (matrix_add add (matrix_add_zero add m n) mx) mx)
add: FStar.Algebra.CommMonoid.Equiv.cm c eq -> mx: FStar.Matrix.matrix c m n
  -> FStar.Pervasives.Lemma
    (ensures
      FStar.Matrix.matrix_eq_fun eq
        (FStar.Matrix.matrix_add add (FStar.Matrix.matrix_add_zero add m n) mx)
        mx)