Type-level matrix

Matrix.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}

module Matrix where

data Nat = Z | S Nat

data Vec (n :: Nat) a where
  VNil  :: Vec 'Z a
  VCons :: a -> Vec n a -> Vec ('S n) a

deriving instance Show a => Show (Vec n a)

instance Functor (Vec n) where
  fmap _ VNil         = VNil
  fmap f (VCons x xs) = VCons (f x) (fmap f xs)

class KnownLen (n :: Nat) where
  vreplicate :: a -> Vec n a

instance KnownLen 'Z where
  vreplicate _ = VNil

instance KnownLen n => KnownLen ('S n) where
  vreplicate x = VCons x (vreplicate x)

newtype Matrix (w :: Nat) (h :: Nat) a =
  Matrix (Vec h (Vec w a))

deriving instance Show a => Show (Matrix w h a)

instance Functor (Matrix w h) where
  fmap f (Matrix rs) = Matrix (fmap (fmap f) rs)

data SomeMatrix a where
  SomeMatrix :: Matrix w h a -> SomeMatrix a

deriving instance Show a => Show (SomeMatrix a)

vsnoc :: Vec n a -> a -> Vec ('S n) a
vsnoc VNil         y = VCons y VNil
vsnoc (VCons x xs) y = VCons x (vsnoc xs y)

vzipWith :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
vzipWith _ VNil         VNil         = VNil
vzipWith f (VCons x xs) (VCons y ys) = VCons (f x y) (vzipWith f xs ys)

emptyMatrix :: Matrix 'Z 'Z a
emptyMatrix = Matrix VNil

-- | A 1×1 matrix holding a single cell.
--
-- >>> singleton 42
-- Matrix (VCons (VCons 42 VNil) VNil)
singleton :: a -> Matrix ('S 'Z) ('S 'Z) a
singleton x = Matrix (VCons (VCons x VNil) VNil)

infixr 5 >.
infixl 5 .>
infixr 5 >|
infixl 5 |>
infixl 6 |+|
infixl 7 |*|

-- | Prepend a row. The row's width must match the matrix's width —
-- checked at the type level, so a mismatched row is a compile-time error.
--
-- Prepend a single-cell row to a 'singleton':
--
-- >>> VCons 2 VNil >. singleton 1
-- Matrix (VCons (VCons 2 VNil) (VCons (VCons 1 VNil) VNil))
--
-- Grow the 'singleton' first with a column, then prepend a matching row:
--
-- >>> singleton 1 |> VCons 2 VNil
-- Matrix (VCons (VCons 1 (VCons 2 VNil)) VNil)
--
-- >>> VCons 3 (VCons 4 VNil) >. (singleton 1 |> VCons 2 VNil)
-- Matrix (VCons (VCons 3 (VCons 4 VNil)) (VCons (VCons 1 (VCons 2 VNil)) VNil))
(>.) :: Vec w a -> Matrix w h a -> Matrix w ('S h) a
row >. Matrix rs = Matrix (VCons row rs)

-- | Append a row. The row's width must match the matrix's width —
-- checked at the type level.
--
-- Append a single-cell row to a 'singleton':
--
-- >>> singleton 1 .> VCons 2 VNil
-- Matrix (VCons (VCons 1 VNil) (VCons (VCons 2 VNil) VNil))
--
-- Build a 2×2 matrix by first widening a 'singleton' and then appending
-- a row of matching width:
--
-- >>> singleton 1 |> VCons 2 VNil .> VCons 3 (VCons 4 VNil)
-- Matrix (VCons (VCons 1 (VCons 2 VNil)) (VCons (VCons 3 (VCons 4 VNil)) VNil))
(.>) :: Matrix w h a -> Vec w a -> Matrix w ('S h) a
Matrix rs .> row = Matrix (vsnoc rs row)

-- | Prepend a column. The column's height must match the matrix's
-- height — checked at the type level.
--
-- Prepend a single-cell column to a 'singleton':
--
-- >>> VCons 2 VNil >| singleton 1
-- Matrix (VCons (VCons 2 (VCons 1 VNil)) VNil)
--
-- Grow the 'singleton' first with a row, then prepend a matching column:
--
-- >>> singleton 1 .> VCons 2 VNil
-- Matrix (VCons (VCons 1 VNil) (VCons (VCons 2 VNil) VNil))
--
-- >>> VCons 3 (VCons 4 VNil) >| (singleton 1 .> VCons 2 VNil)
-- Matrix (VCons (VCons 3 (VCons 1 VNil)) (VCons (VCons 4 (VCons 2 VNil)) VNil))
(>|) :: Vec h a -> Matrix w h a -> Matrix ('S w) h a
col >| Matrix rs = Matrix (vzipWith VCons col rs)

-- | Append a column. The column's height must match the matrix's
-- height — checked at the type level.
--
-- Append a single-cell column to a 'singleton':
--
-- >>> singleton 1 |> VCons 2 VNil
-- Matrix (VCons (VCons 1 (VCons 2 VNil)) VNil)
--
-- Build a 2×2 matrix by first growing a 'singleton' downward and then
-- appending a column of matching height:
--
-- >>> singleton 1 .> VCons 2 VNil |> VCons 3 (VCons 4 VNil)
-- Matrix (VCons (VCons 1 (VCons 3 VNil)) (VCons (VCons 2 (VCons 4 VNil)) VNil))
(|>) :: Matrix w h a -> Vec h a -> Matrix ('S w) h a
Matrix rs |> col = Matrix (vzipWith vsnoc rs col)

-- | Add two matrices element-wise. The dimensions must match — checked
-- at the type level, so summing matrices of different shape is a
-- compile-time error.
--
-- >>> singleton 1 |+| singleton 2
-- Matrix (VCons (VCons 3 VNil) VNil)
--
-- A 2×2 example:
--
-- >>> let a = singleton 1 |> VCons 2 VNil .> VCons 3 (VCons 4 VNil)
-- >>> let b = singleton 10 |> VCons 20 VNil .> VCons 30 (VCons 40 VNil)
-- >>> a |+| b
-- Matrix (VCons (VCons 11 (VCons 22 VNil)) (VCons (VCons 33 (VCons 44 VNil)) VNil))
--
-- Adding the 0×0 matrix to itself (result typed as @Int@ via 'fmap'):
--
-- >>> emptyMatrix |+| fmap (const (0 :: Int)) emptyMatrix
-- Matrix VNil
(|+|) :: Num a => Matrix w h a -> Matrix w h a -> Matrix w h a
Matrix xs |+| Matrix ys = Matrix (vzipWith (vzipWith (+)) xs ys)

-- | Multiply two matrices. The left matrix's width must equal the right
-- matrix's height — checked at the type level. The result is @height-of-left
-- × width-of-right@.
--
-- >>> singleton 3 |*| singleton 4
-- Matrix (VCons (VCons 12 VNil) VNil)
--
-- A 2×2 times 2×2:
--
-- >>> let a = singleton 1 |> VCons 2 VNil .> VCons 3 (VCons 4 VNil)
-- >>> let b = singleton 5 |> VCons 6 VNil .> VCons 7 (VCons 8 VNil)
-- >>> a |*| b
-- Matrix (VCons (VCons 19 (VCons 22 VNil)) (VCons (VCons 43 (VCons 50 VNil)) VNil))
--
-- A 1×2 row vector times a 2×1 column vector gives a 1×1 matrix:
--
-- >>> let row = singleton 2 |> VCons 3 VNil
-- >>> let col = singleton 4 .> VCons 5 VNil
-- >>> row |*| col
-- Matrix (VCons (VCons 23 VNil) VNil)
(|*|) :: (Num a, KnownLen w') => Matrix w h a -> Matrix w' w a -> Matrix w' h a
Matrix xs |*| Matrix ys =
  let vdot :: Num a => Vec n a -> Vec n a -> a
      vdot VNil         VNil         = 0
      vdot (VCons x xs) (VCons y ys) = x * y + vdot xs ys
      transposeCols :: KnownLen w => Vec h (Vec w a) -> Vec w (Vec h a)
      transposeCols VNil         = vreplicate VNil
      transposeCols (VCons r rs) = vzipWith VCons r (transposeCols rs)
  in  Matrix (fmap (\r -> fmap (vdot r) (transposeCols ys)) xs)

-- | Render a matrix as an ASCII table. Each cell is 'show'n; every column
-- expands to fit its widest cell.
--
-- >>> putStr (prettyPrint (singleton 1 |> VCons 2 VNil .> VCons 3 (VCons 4 VNil)))
-- +---+---+
-- | 1 | 2 |
-- +---+---+
-- | 3 | 4 |
-- +---+---+
--
-- Column widths adapt to the widest cell:
--
-- >>> putStr (prettyPrint (singleton 100 |> VCons 2 VNil .> VCons 30 (VCons 4 VNil)))
-- +-----+---+
-- | 100 | 2 |
-- +-----+---+
-- | 30  | 4 |
-- +-----+---+
--
-- Non-numeric cells work too:
--
-- >>> putStr (prettyPrint (singleton "hi" |> VCons "there" VNil))
-- +------+---------+
-- | "hi" | "there" |
-- +------+---------+
prettyPrint :: Show a => Matrix w h a -> String
prettyPrint (Matrix rs) =
  let vecToList :: Vec n a -> [a]
      vecToList VNil         = []
      vecToList (VCons x xs) = x : vecToList xs
  in  case map (map show . vecToList) (vecToList rs) of
        []       -> ""
        cellRows ->
          let colCount  = length (head cellRows)
              colWidths = [ maximum (0 : [length (r !! i) | r <- cellRows])
                          | i <- [0 .. colCount - 1]
                          ]
              sep       = "+" ++ concatMap (\w -> replicate (w + 2) '-' ++ "+") colWidths ++ "\n"
              pad w s   = s ++ replicate (w - length s) ' '
              rowLine r = "|" ++ concat (zipWith (\w s -> " " ++ pad w s ++ " |") colWidths r) ++ "\n"
          in sep ++ concatMap (\r -> rowLine r ++ sep) cellRows

listToVec :: [a] -> (forall n. Vec n a -> r) -> r
listToVec []     k = k VNil
listToVec (x:xs) k = listToVec xs (\v -> k (VCons x v))

-- | Build a matrix from a list of rows. Returns 'Nothing' if the rows do
-- not all share the same width (i.e. the outer list is ragged). The result
-- is wrapped in 'SomeMatrix' because the dimensions are not known until
-- runtime.
--
-- >>> fromLists [[1,2,3],[4,5,6]]
-- Just (SomeMatrix (Matrix (VCons (VCons 1 (VCons 2 (VCons 3 VNil))) (VCons (VCons 4 (VCons 5 (VCons 6 VNil))) VNil))))
--
-- A ragged input is rejected:
--
-- >>> fromLists [[1,2],[3,4,5]]
-- Nothing
--
-- An empty list of rows is treated as a 0×0 matrix:
--
-- >>> fromLists ([] :: [[Int]])
-- Just (SomeMatrix (Matrix VNil))
fromLists :: forall a. [[a]] -> Maybe (SomeMatrix a)
fromLists []     = Just (SomeMatrix (emptyMatrix :: Matrix 'Z 'Z a))
fromLists (r:rs) =
  let matchVec :: Vec n a -> [b] -> Maybe (Vec n b)
      matchVec VNil         []     = Just VNil
      matchVec (VCons _ xs) (y:ys) = VCons y <$> matchVec xs ys
      matchVec _            _      = Nothing
  in  listToVec r $ \firstRow -> do
      fmap  (\otherRows -> listToVec (firstRow : otherRows)
            (\allRows -> SomeMatrix (Matrix allRows))) (traverse (matchVec firstRow) rs)