val 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 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
val 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 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
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 Prims.pos Prims.b2t Prims.op_GreaterThan FStar.Algebra.CommMonoid.Equiv.cm FStar.Matrix.matrix_generator Prims.op_Subtraction FStar.Classical.forall_intro FStar.IntegerIntervals.under Prims.eq2 FStar.Seq.Permutation.foldm_snoc FStar.Seq.Base.init Prims.unit Prims.l_True Prims.squash Prims.Nil FStar.Pervasives.pattern FStar.Seq.Base.lemma_eq_elim Prims.l_Forall
(*
   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))
val 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)))

    cm: FStar.Algebra.CommMonoid.Equiv.cm c eq ->
    gen: FStar.Matrix.matrix_generator c m n ->
    subgen: FStar.Matrix.matrix_generator c (m - 1) n
  -> FStar.Pervasives.Lemma (requires subgen == (fun i j -> gen i j))
      (ensures
        forall (i: FStar.IntegerIntervals.under (m - 1)).
          FStar.Seq.Permutation.foldm_snoc cm (FStar.Seq.Base.init n (subgen i)) ==
          FStar.Seq.Permutation.foldm_snoc cm (FStar.Seq.Base.init n (gen i)))