A matrix indexed by its diagonal

MatrixDiagonal.idr

module Main

import Data.Vect

namespace row
  -- (VectD n a i x) is like (xs : Vect n a),
  -- with an extra (x = index' m (toList xs)) index
  data VectD : Nat -> (a : Type) -> Nat -> Maybe a -> Type where
    Nil  : VectD Z a i Nothing
    (::) : (x : a)
        -> VectD n a (pred i) x'
        -> VectD (S n) a i (if isZero i then Just x else x')

  example1 : VectD 3 Nat 0 (Just 4)
  example1 = [4, 5, 6]

  example2 : VectD 3 Nat 1 (Just 5)
  example2 = [4, 5, 6]

  example3 : VectD 3 Nat 2 (Just 6)
  example3 = [4, 5, 6]

  example4 : VectD 3 Nat 3 Nothing
  example4 = [4, 5, 6]

namespace diagonal
  maybeCons : Maybe a -> List a -> List a
  maybeCons Nothing  xs = xs
  maybeCons (Just x) xs = x :: xs
  
  -- a diagonal which starts at an horizontal offset
  diagonalAt : Nat -> List (List a) -> List a
  diagonalAt i Nil = Nil
  diagonalAt i (xs :: xss) = maybeCons (index' i xs) (diagonalAt (S i) xss)
  
  example1 : diagonalAt 0 [[1, 2, 3]
                          ,[4, 5, 6]
                          ]
           = the (List Nat) [1, 5]
  example1 = Refl
  
  example2 : diagonalAt 1 [[1, 2, 3]
                          ,[4, 5, 6]
                          ]
           = the (List Nat) [2, 6]
  example2 = Refl
  
  example3 : diagonalAt 2 [[1, 2, 3]
                          ,[4, 5, 6]
                          ]
           = the (List Nat) [3]
  example3 = Refl
  
  example4 : diagonalAt 3 [[1, 2, 3]
                          ,[4, 5, 6]
                          ]
           = the (List Nat) []
  example4 = Refl
  
namespace matrix
  -- (MatrixD m n a i xs) is like (xss : Vect m (Vect n a)),
  -- with an extra (xs = diagonalAt i (toList (map toList xss))) index
  data MatrixD : Nat -> Nat -> (a : Type) -> Nat -> List a -> Type where
    Nil  : MatrixD Z m a i Nil
    (::) : VectD n a i x
        -> MatrixD m n a (S i) xs
        -> MatrixD (S m) n a i (maybeCons x xs)

example1 : MatrixD 2 2 Nat 0 [1, 4]
example1 = [[1, 2]
           ,[3, 4]
           ]

example2 : MatrixD 2 3 Nat 0 [1, 5]
example2 = [[1, 2, 3]
           ,[4, 5, 6]
           ]

example3 : MatrixD 2 3 Nat 1 [2, 6]
example3 = [[1, 2, 3]
           ,[4, 5, 6]
           ]

example4 : MatrixD 2 3 Nat 2 [3]
example4 = [[1, 2, 3]
           ,[4, 5, 6]
           ]

example5 : MatrixD 3 2 Nat 0 [1, 4]
example5 = [[1, 2]
           ,[3, 4]
           ,[5, 6]
           ]


main : IO ()
main = do
  putStrLn "typechecks."