How to generalize dependently sized arrays to n dimensions?

Question

I have been toying with this for some time now but I haven't been able to convince GHC to make this work.

Basically it is quite easy to create dependently sized arrays in current versions of Haskell/GHC:

newtype Arr1 (w :: Nat) a = Arr1 (Int -> a)
newtype Arr2 (w :: Nat) (h :: Nat) a = Arr2 (Int -> a)

ix2 :: forall w h a. (KnownNat w) => Arr2 w h a -> Int -> Int -> a
ix2 (Arr2 f) x y = f ( y * w + x )
    where w = fromInteger $ natVal (Proxy :: Proxy w)

sub2 :: forall w h a. (KnownNat w) => Arr2 w h a -> Int -> Arr1 w a
sub2 (Arr2 f) y = Arr1 $ \x -> f (y * w + x)
    where w = fromInteger $ natVal (Proxy :: Proxy w)

mkArr2V :: forall w h a. (V.Unbox a, KnownNat w, KnownNat h) => V.Vector a -> Arr2 w h a
mkArr2V v = Arr2 $ (v V.!)

-- and so on ... errorchecking neglected

But current GHC versions give us much more expressibility. Basically it should be possible to create a single type for this:

newtype Mat (s :: [Nat]) a = Mat (Int -> a)

-- create array backed by vector
mkMatV :: forall s a. V.Vector a -> Mat s a
mkMatV v = Mat $ (v V.!)

This works in GHCi:

>>> let m = mkMatV (V.fromList [1,2,3,4]) :: Mat [2,2] Double
>>> :t m
m :: Mat '[2, 2] Double

But up to this point I am unsure on how to accomplish indexing into the arrays. A simple solution would be to use two different functions for nd and 1d indexing. Note this does not typecheck.

-- slice from nd array
(!) :: forall s ss a. (KnownNat s) => Mat (s ': ss) a -> Int -> Mat ss a
(!) (Mat f) o = Mat $ \i -> f (o*s+i)
    where s = fromInteger $ natVal (Proxy :: Proxy (sum ss))

-- index into 1d array
(#) :: forall s ss a. (KnownNat s) => Mat (s ': '[]) a -> Int -> a
(#) (Mat f) o = Mat $ \i -> f o

Probably usable like this:

>>> :t m ! 0
Mat [2] Double
>>> m ! 0 # 0
1

Not that it is necessary to give indices in z,y,x order. My prefered solution would provide a single indexing function that changes it's return type based on the dimensionality of the array. To my knowledge this can somehow be achieved by using type classes, but I haven't figured that out yet. Bonus points if indices can be given in "natural" x,y,z order.

tl;dr: I am asking for a function that indexes n-dimensional arrays as defined above.

Answer 1

This can be indeed done with type classes. Some preliminaries:

{-# LANGUAGE
  UndecidableInstances, MultiParamTypeClasses, TypeFamilies,
  ScopedTypeVariables, FunctionalDependencies, TypeOperators,
  DataKinds, FlexibleInstances #-}

import qualified Data.Vector as V

import GHC.TypeLits
import Data.Proxy

newtype NVec (shape :: [Nat]) a = NVec {_data :: V.Vector a}

Before anything else, we should be able to tell the overall flat size of an n-dimensional vector. We'll use this to compute strides for indexing. We use a class to recurse on the type-level list.

class FlatSize (sh :: [Nat]) where
  flatSize :: Proxy sh -> Int

instance FlatSize '[] where
  flatSize _ = 1

instance (KnownNat s, FlatSize ss) => FlatSize (s ': ss) where
  flatSize _ = fromIntegral (natVal (Proxy :: Proxy s)) * flatSize (Proxy :: Proxy ss)       

We use a type class for indexing too. We provide different instances for the one-dimensional case (where we simply index into the underlying vector) and the higher-dimensional case (where we return a new NVec with reduced dimension). We use the same class for both cases, though.

infixl 5 !                                            
class Index (sh :: [Nat]) (a :: *) (b :: *) | sh a -> b where
  (!) :: NVec sh a -> Int -> b

instance Index '[s] a a where
  (NVec v) ! i = v V.! i         

instance (Index (s2 ': ss) a b, FlatSize (s2 ': ss), res ~ NVec (s2 ': ss) a) 
  => Index (s1 ': s2 ': ss) a res where
  (NVec v) ! i = NVec (V.slice (i * stride) stride v)
    where stride = flatSize (Proxy :: Proxy (s2 ': ss))

Indexing into a higher-dimensional vector is just taking a slice with the flat size of the resulting vector and the appropriate offset.

Some testing:

fromList :: forall a sh. FlatSize sh => [a] -> NVec sh a
fromList as | length as == flatSize (Proxy :: Proxy sh) = NVec (V.fromList as)
fromList _ = error "fromList: initializer list has wrong size"

v3 :: NVec [2, 2, 2] Int
v3 = fromList [
  2, 4,
  5, 6,

  10, 20,
  30, 0 ]

v2 :: NVec [2, 2] Int
v2 = v3 ! 0

vElem :: Int
vElem = v3 ! 0 ! 1 ! 1 -- 6 

---

As an extra, let me present a singletons solution too, because it's considerably more convenient. It lets us reuse more code (less custom typeclasses for a single function) and write in a more direct, functional style.

{-# LANGUAGE
  UndecidableInstances, MultiParamTypeClasses, TypeFamilies,
  ScopedTypeVariables, FunctionalDependencies, TypeOperators,
  DataKinds, FlexibleInstances, StandaloneDeriving, DeriveFoldable,
  GADTs, FlexibleContexts #-}

import qualified Data.Vector as V
import qualified Data.Foldable as F
import GHC.TypeLits
import Data.Singletons.Preludeimport 
import Data.Singletons.TypeLits

newtype NVec (shape :: [Nat]) a = NVec {_data :: V.Vector a}

flatSize becomes much simpler: we just lower the sh to the value level, and operate on it as usual:

flatSize :: Sing (sh :: [Nat]) -> Int
flatSize = fromIntegral . product . fromSing

We use a type family and a function for indexing. In the previous solution we used instances to dispatch on the dimensions; here we do the same with pattern matching:

type family Index (shape :: [Nat]) (a :: *) where
  Index (s  ': '[])       a = a
  Index (s1 ':  s2 ': ss) a = NVec (s2 ': ss) a

infixl 5 !
(!) :: forall a sh. SingI sh => NVec sh a -> Int -> Index sh a
(!) (NVec v) i = case (sing :: Sing sh) of
  SCons _ SNil       -> v V.! i
  SCons _ ss@SCons{} -> NVec (V.slice (i * stride) stride v) where
    stride = flatSize ss

We can also use the Nat singletons for safe indexing and initialization (i. e. with statically checked bounds and sizes). For initialization we define a list type with static size (Vec).

safeIx ::
  forall a s sh i. (SingI (s ': sh), (i + 1) <= s) =>
  NVec (s ': sh) a -> Sing i -> Index (s ': sh) a
safeIx v si = v ! (fromIntegral $ fromSing si)                    

data Vec n a where
  VNil :: Vec 0 a
  (:>) :: a -> Vec (n - 1) a -> Vec n a
infixr 5 :>
deriving instance F.Foldable (Vec n)

fromVec :: forall a sh. SingI sh => Vec (Foldr (:*$) 1 sh) a -> NVec sh a
fromVec = fromList . F.toList

Some examples for the safe functions:

-- Other than 8 elements in the Vec would be a type error
v3 :: NVec [2, 2, 2] Int
v3 = fromVec
     (2 :> 4  :>
      5 :> 6  :>

      10 :> 20 :>
      30 :> 0  :> VNil)

vElem :: Int
vElem = v3
  `safeIx` (sing :: Sing 0)
  `safeIx` (sing :: Sing 1)
  `safeIx` (sing :: Sing 1) -- 6