Типобезопасные матрицы на Haskell

Зависимых типов в хаскелле конечно нет, но при большом желании (и наличии в хабах поста хаба ненормальное программирование) в по крайней мере некотрых случаях их можно эмулировать с помощью синглтонов. Например, в строке

transposeMatrix Nil = _

нам действительно нужно вернуть вектор Nilов длины n, и n действительно стирается в рантайме. Однако мы можем "удержать" ее, проиндексировав ею GADT и сохранив этот GADT в Matrix.

data Nat = Zero | Succ Nat 
data SNat s where
  SZero :: SNat 'Zero
  SSucc :: SNat n -> SNat ('Succ n)

-- Библиотека singletons умеет автоматически генерировать SNat из Nat с помощью template haskell, но для наглядности изобретаю велосипед

Здесь SNat — тип-синглтон (т.е для каждого n существует только одно значение SNat n), проиндексированный типом Nat. Теперь если мы паттерн-матчим значение n типа SNat n с SZero, мы убеждаем тайпчекер в том, что в этой ветке n ~ 'Zero. Аналогично если n является SSucc (n1 :: SNat n1), то тайпчекер выводит n ~ Succ n1. Проверить можно с помощью holes:

q :: SNat n -> ()
q SZero = _zero
q (SSucc n1) = _succ

Сокращенный вывод ghc:

Found hole: _zero :: ()
Constraints include
  n ~ 'Zero

Found hole: _succ :: ()
Relevant bindings include
  n1 :: SNat n1
Constraints include
  n ~ 'Succ n1 

Таким образом мы не только сохраняем тип в рантайм значение, но и можем "поднять" его обратно в тип.

Теперь мы можем сохранить n и m в виде SNat в Matrix:

data Matrix (m :: Nat) (n :: Nat) a = Matrix (SNat m) (SNat n) (Vector m (Vector n a))

Функция transposeMatrix тогда может выглядеть так:

transposeMatrix :: Matrix m n a -> Matrix n m a
transposeMatrix (Matrix m' n' vv) = Matrix n' m' (go m' n' vv)
  where
    go :: SNat m -> SNat n -> Vector m (Vector n a) -> Vector n (Vector m a)
    go SZero n Nil = replicateVec n Nil
    go _ SZero _ = Nil
    go m (SSucc n1) ((x :| xs) :| xss) = (x :| fmap headVec xss) :| go m n1 (xs :| fmap tailVec xss)

    replicateVec :: SNat n -> a -> Vector n a
    replicateVec SZero _ = Nil
    replicateVec (SSucc n) a = a :| replicateVec n a

Здесь replicateVec строит вектор длины n заданной SNat n, повторяя a. Паттерн матчингом по SNat n мы убеждаем тайпчекер в том, что мы действительно строим вектор длины n.

В функции go:

• в первом кейсе внешний вектор Nil, поэтому m ~ 'Zero, вызываем replicateVec с длиной n, чтобы создать Vector n (Vector 'Zero a),
• в втором кейсе n матчится с SZero, из этого n ~ Zero, достаточно вернуть Nil :: Vector 'Zero (Vector m a),
• в третьем кейсе матрица имеет тип Vector m (Vector ('Succ n1) a), паттерн матчингом n с (SSucc n1) мы получаем SNat n1, дальше единственное изменение — передача m и n1 в рекурсивный вызов go

Собственно на этом все.

Чтобы не прописывать при конструировании Matrix длины руками, можно сделать что-то вроде такого:

class Promote (n :: Nat) where
  promote :: Proxy n -> SNat n

instance Promote 'Zero where
  promote _ = SZero

instance Promote n => Promote ('Succ n) where
  promote _ = SSucc (promote $ Proxy @n)

matrix :: (Promote n, Promote m) => Vector n (Vector m a) -> Matrix n m a
matrix = Matrix (promote Proxy) (promote Proxy)

Пример:

> transposeMatrix $ matrix ( (1 :| 2 :| 3 :| Nil) :| (4 :| 5 :| 6 :| Nil) :| Nil )

Matrix (SSucc (SSucc (SSucc SZero))) (SSucc (SSucc SZero)) ((1 :| (4 :| Nil)) :| ((2 :| (5 :| Nil)) :| ((3 :| (6 :| Nil)) :| Nil)))

{-# LANGUAGE GADTs, DataKinds, KindSignatures, TypeApplications,
  DerivingStrategies, StandaloneDeriving, ScopedTypeVariables #-}

import Prelude (IO, Integer, print, Show, Functor(fmap), ($), (<$>))
import Data.Proxy (Proxy(Proxy))

data Nat = Zero | Succ Nat

data SNat s where
  SZero :: SNat 'Zero
  SSucc :: SNat n -> SNat ('Succ n)

deriving stock instance Show (SNat s)

infixr 3 :|
data Vector (n :: Nat) a where
  (:|) :: a -> Vector n a -> Vector ('Succ n) a
  Nil  :: Vector 'Zero a

deriving stock instance Show a => Show (Vector n a)

data Matrix (m :: Nat) (n :: Nat) a = Matrix (SNat m) (SNat n) (Vector m (Vector n a))
  deriving stock Show

instance Functor (Vector n) where
  fmap f (v :| vs) = f v :| fmap f vs
  fmap _ Nil = Nil

instance Functor (Matrix m n) where
  fmap f (Matrix m n vs) = Matrix m n $ fmap f <$> vs

headVec :: Vector ('Succ n) a -> a
headVec (a :| _) = a

tailVec :: Vector ('Succ n) a -> Vector n a
tailVec (_ :| a) = a

transposeMatrix :: Matrix m n a -> Matrix n m a
transposeMatrix (Matrix m' n' vv) = Matrix n' m' (go m' n' vv)
  where
    go :: SNat m -> SNat n -> Vector m (Vector n a) -> Vector n (Vector m a)
    go SZero n Nil = replicateVec n Nil
    go _ SZero _ = Nil
    go m (SSucc n1) ((x :| xs) :| xss) = (x :| fmap headVec xss) :| go m n1 (xs :| fmap tailVec xss)

    replicateVec :: SNat n -> a -> Vector n a
    replicateVec SZero _ = Nil
    replicateVec (SSucc n) a = a :| replicateVec n a

class Promote (n :: Nat) where
  promote :: Proxy n -> SNat n

instance Promote 'Zero where
  promote _ = SZero

instance Promote n => Promote ('Succ n) where
  promote _ = SSucc (promote $ Proxy @n)

matrix :: (Promote n, Promote m) => Vector n (Vector m a) -> Matrix n m a
matrix = Matrix (promote Proxy) (promote Proxy)

main :: IO ()
main = print $ transposeMatrix $ matrix ( ( (1 :: Integer) :| 2 :| 3 :| Nil) :| (4 :| 5 :| 6 :| Nil) :| Nil )