如何将依赖大小的数组推广到n维？

Question

我已经玩了一段时间了,但是我无法说服GHC做这项工作.

基本上,在当前版本的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

但目前的GHC版本为我们提供了更多的可表达性.基本上应该可以为此创建一个类型:

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.!)

这适用于GHCi:

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

但到目前为止,我不确定如何完成对数组的索引.一个简单的解决方案是使用两个不同的函数进行nd和1d索引.请注意,这不是类型检查.

-- 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

可能像这样使用:

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

并不是必须以z,y,x顺序给出索引.我首选的解决方案将提供单个索引函数,该函数根据数组的维度更改其返回类型.据我所知,这可以通过使用类型类来实现,但我还没想到.如果指数可以"自然"x,y,z顺序给出,则可获得奖励积分.

tl; dr:我要求一个函数来索引上面定义的n维数组.

Answer 1

这确实可以通过类型类来完成.一些预赛:

{-# 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}

在此之前,我们应该能够分辨出n维向量的整体平面大小.我们将使用它来计算索引的步幅.我们使用类来递归类型级别列表.

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)       

我们也使用类型类进行索引.我们为一维情况(我们简单地索引到底层向量)和高维情况(我们返回一个NVec具有减少维度的新情况)提供不同的实例.但是,我们对这两种情况使用相同的类.

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))

索引到更高维向量只是得到具有所得向量的平坦大小和适当偏移的切片.

一些测试:

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 

另外,让我提出一个singletons解决方案,因为它更方便.它允许我们重用更多代码(更少的单个函数的自定义类型类),并以更直接,更实用的方式编写.

{-# 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变得简单得多:我们只是sh将值降低到值,并像往常一样对其进行操作:

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

我们使用类型族和函数进行索引.在之前的解决方案中,我们使用实例来分析维度; 这里我们对模式匹配做同样的事情:

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

我们还可以使用Nat单例进行安全索引和初始化(即使用静态检查的边界和大小).对于初始化,我们定义一个静态大小(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

安全功能的一些示例:

-- 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