module Data.Index
(
Dim( each
, size
, rank
, correct
, reflect
, toIndex
, fromIndex
, next
, prev)
, zero
, Rank(..)
, (:.)(..), Z(..)
, Mode(..), roll, unroll, modeProxy
, foldlRange
, foldrRange
, withRange
, foldlRangeIndices
, foldrRangeIndices
, withRangeIndices
, Ranged
, InRange
, Range()
, Peano(..)
, ToPeano
, Size
, bounds
, range
, dimHead, dimTail
, pdimHead, pdimTail
, cnat
, And
, dim, dimu, dimr
, module Data.Proxy
) where
import GHC.Generics
import GHC.TypeLits
import Data.Proxy
import Data.Tagged
import qualified Data.Ix as Ix
import Data.Monoid (Monoid(..))
import Data.Data (Typeable, Data)
import Data.Foldable (Foldable)
import Data.Traversable (Traversable)
import Control.Applicative
import Language.Haskell.TH hiding (Range(..))
import Language.Haskell.TH.Quote
data a :. b = !Int :. !b
deriving (Show,Read,Eq,Ord,Generic,Functor,Foldable,Traversable)
data Z = Z
deriving (Show,Read,Eq,Ord,Typeable,Data,Generic)
infixr 9 :.
class Rank a b where
setBound :: a -> b
instance Rank Z Z where
setBound _ = Z
instance Rank xs ys => Rank (x:.xs) (y:.ys) where
setBound (x:.xs) = x:.setBound xs
zero :: Dim a => a
zero = each 0
class Ord n => Dim n where
each :: Int -> n
rank :: n -> Int
size :: proxy n -> Int
reflect' :: proxy n -> n
reflect :: n
reflect = reflect' Proxy
next :: n -> n
prev :: n -> n
toIndex :: n -> Int
fromIndex' :: proxy n -> Int -> n
fromIndex :: Int -> n
fromIndex = fromIndex' Proxy
correct :: n -> n
correctOnce :: n -> n
lastDim :: proxy n -> n
zipDims :: (Int -> Int -> Int -> Int) -> n -> n -> n
unsafeZipDims :: (Int -> Int -> Int -> Int) -> n -> n -> n
mapDim :: (Int -> Int -> Int) -> n -> n
unsafeMapDim :: (Int -> Int -> Int) -> n -> n
instance Dim Z where
each _ = Z
rank _ = 0
size _ = 1
reflect' _ = Z
next _ = error "next: Z"
prev _ = error "prev: Z"
toIndex _ = 0
fromIndex' _ i | i > 0 = error "fromIndex: index too large"
| otherwise = Z
correct _ = Z
correctOnce _ = Z
lastDim _ = Z
zipDims _ _ _ = Z
unsafeZipDims _ _ _ = Z
mapDim _ _ = Z
unsafeMapDim _ _ = Z
instance (KnownNat x, Dim xs) => Dim (x:.xs) where
each a
| 0 <= a && a < dimHead d = d
| otherwise = error "each: out of range"
where
d = a:.each a
rank (_:.xs) = 1 + rank xs
size d = pdimHead d * size (pdimTail d)
reflect' d = pdimHead d :. reflect' (pdimTail d)
next d@(x:.xs)
| x < dimHead d 1 = (x+1) :. xs
| otherwise = 0 :. next xs
prev (x:.xs)
| x > 0 = (x1) :. xs
| xs > zero = 0 :. prev xs
| otherwise = maxBound
toIndex d@(x:.xs) = x + dimHead d * toIndex xs
fromIndex' d = \ix -> (ix `mod` h) :. fromIndex' (pdimTail d) (ix `div` h)
where h = pdimHead d
correct d@(x:.xs)
| 0 <= x && x < dimHead d = x:.correct xs
| otherwise = error "correct: index not in range"
correctOnce d@(x:.xs)
| 0 <= x && x < dimHead d = x:.xs
| otherwise = error "correctOnce: index not in range"
lastDim d = (pdimHead d 1) :. lastDim (pdimTail d)
zipDims f d@(x:.xs) (y:.ys)
| 0 <= r && r < dimHead d = r :. zipDims f xs ys
| otherwise = error "zipDims: index not in range"
where r = f (dimHead d) x y
unsafeZipDims f d@(x:.xs) (y:.ys) = f (dimHead d) x y :. unsafeZipDims f xs ys
mapDim f d@(x:.xs)
| 0 <= r && r < dimHead d = r :. mapDim f xs
| otherwise = error "mapDims: index not in range"
where r = f (dimHead d) x
unsafeMapDim f d@(x:.xs) = f (dimHead d) x :. unsafeMapDim f xs
data Mode :: * -> * where
Unroll :: Ranged i => Mode i
Roll :: Dim i => Mode i
unroll :: Ranged a => Proxy a -> Mode a
unroll _ = Unroll
roll :: Dim a => Proxy a -> Mode a
roll _ = Roll
modeProxy :: Mode a -> Proxy a
modeProxy _ = Proxy
range :: Mode n -> [n]
range mode = foldrRange mode (:) []
proxyOf :: a -> Proxy a
proxyOf = const Proxy
foldrRange :: Mode n -> (n -> b -> b) -> b -> b
foldrRange Roll cons nil = go zero
where
top = lastDim Proxy
go !i
| i < top = i `cons` go (next i)
| otherwise = cons top nil
foldrRange Unroll cons nil = sfoldrRange_ (tagPeano zero) cons nil
foldlRange :: Mode n -> (b -> n -> b) -> b -> b
foldlRange Roll f = go zero
where
top = lastDim Proxy
go !i !acc
| i < top = go (next i) (f acc i)
| otherwise = f acc i
foldlRange Unroll f = sfoldlRange_ (tagPeano zero) f
withRange :: Applicative m => Mode a -> (a -> m ()) -> m ()
withRange m f = foldrRange m (\d acc -> acc *> f d) (pure ())
foldlRangeIndices :: Mode n -> (b -> Int -> b) -> b -> b
foldlRangeIndices m@Roll cons = go 0
where
s = size m
go !i !acc
| i < s = go (i+1) (cons acc i)
| otherwise = acc
foldlRangeIndices m@Unroll cons = sfoldlRangeIndices_ (tagPeanoI m) cons
foldrRangeIndices :: Mode n -> (Int -> b -> b) -> b -> b
foldrRangeIndices m@Roll cons nil = go 0
where
s = size m
go !i
| i < s = i `cons` go (i+1)
| otherwise = nil
foldrRangeIndices m@Unroll cons nil = sfoldrRangeIndices_ (tagPeanoI m) cons nil
withRangeIndices :: Applicative m => Mode n -> (Int -> m ()) -> m ()
withRangeIndices m f = foldrRangeIndices m (\d acc -> acc *> f d) (pure ())
zipMod :: Dim n => (Int -> Int -> Int) -> n -> n -> n
zipMod f = unsafeZipDims (\h a b -> f a b `mod` h)
unsafeZipDims':: Dim n => (Int -> Int -> Int) -> n -> n -> n
unsafeZipDims' f = unsafeZipDims (\_ a b -> f a b)
mapMod :: Dim n => (Int -> Int) -> n -> n
mapMod f = unsafeMapDim (\h a -> f a `mod` h)
instance Num Z where
_ + _ = Z
_ _ = Z
_ * _ = Z
negate _ = Z
signum _ = Z
abs _ = Z
fromInteger _ = Z
instance Dim (x:.xs) => Num (x:.xs) where
(+) = zipMod (+)
() = zipMod ()
(*) = zipMod (*)
negate = mapMod negate
signum = mapMod signum
abs = mapMod abs
fromInteger = fromIndex . fromIntegral
instance Real Z where
toRational _ = 0
instance (Num (x:.xs), Dim (x:.xs)) => Real (x:.xs) where
toRational = toRational . toIndex
instance Enum Z where
toEnum 0 = Z
toEnum _ = error "toEnum: must be 0"
fromEnum _ = 0
succ _ = error "succ: Z"
pred _ = error "pred: Z"
enumFrom _ = [Z]
enumFromThen _ _ = [Z]
enumFromTo _ _ = [Z]
enumFromThenTo _ _ _ = [Z]
instance Dim (x:.xs) => Enum (x:.xs) where
toEnum = fromIndex
fromEnum = toIndex
succ d
| d < lastDim Proxy = next d
| otherwise = error "succ: index within bound"
pred d
| d > zero = prev d
| otherwise = error "succ: index within bound"
enumFrom !l
| l < maxBound = l : enumFrom (next l)
| otherwise = [l]
enumFromThen !l !j
| l < maxBound = l : enumFromThen (zipMod (+) l j) j
| otherwise = [l]
enumFromTo !l !h
| l < h = l : enumFromTo (next l) h
| otherwise = [h]
instance Integral Z where
div _ _ = Z
mod _ _ = Z
quot _ _ = Z
rem _ _ = Z
quotRem _ _ = (Z, Z)
divMod _ _ = (Z, Z)
toInteger _ = 0
instance (Integral xs, Dim (x:.xs), Enum (x:.xs)) => Integral (x:.xs) where
div = unsafeZipDims' div
mod = unsafeZipDims' mod
quot = unsafeZipDims' quot
rem = unsafeZipDims' rem
quotRem xs ys = (quot xs ys, rem xs ys)
divMod xs ys = (div xs ys, mod xs ys)
toInteger = toInteger . toIndex
instance Bounded Z where
minBound = Z
maxBound = Z
instance Dim (x:.xs) => Bounded (x:.xs) where
minBound = zero
maxBound = lastDim Proxy
instance Monoid Z where
mempty = Z
mappend _ _ = Z
instance (Dim (x:.xs), Monoid xs) => Monoid (x:.xs) where
mempty = zero
mappend = (+)
instance KnownNat s => Applicative ((:.) s) where
pure x = 1:.x
d@(ix0 :. f) <*> ix1 :. x = ((ix0*ix1) `mod` dimHead d) :. f x
(*>) = (>>)
(<*) = flip (>>)
instance KnownNat s => Monad ((:.) s) where
return x = 1:.x
d@(ix0:.a) >>= f = case f a of
ix1 :. b -> ((ix0*ix1) `mod` dimHead d) :. b
instance Ix.Ix Z where
range _ = [Z]
index _ _ = 0
inRange _ _ = True
rangeSize _ = 0
instance (Ranged (x:.xs), Num xs) => Ix.Ix (x:.xs) where
range _ = range Unroll
index _ c = toIndex c
inRange _ c = minBound <= c && c <= maxBound
rangeSize (_,b) = size (proxyOf b)
dimQQ :: ExpQ -> TypeQ -> QuasiQuoter
dimQQ val ty = QuasiQuoter
{ quoteExp = \s -> [| $val :: $(quoteType (dimQQ val ty) s) |]
, quoteType = \s ->
let cons a b = [t| $(litT $ numTyLit a) :. $b |]
nil = [t| Z |]
in [t| $ty $(foldr cons nil . map read . words $ s) |]
, quotePat = error "dim in pattern context"
, quoteDec = error "dim in declaration context"
}
dim :: QuasiQuoter
dim = dimQQ [| Proxy |] [t| Proxy |]
dimu :: QuasiQuoter
dimu = dimQQ [| Unroll |] [t| Mode |]
dimr :: QuasiQuoter
dimr = dimQQ [| Roll |] [t| Mode |]
dimHead :: KnownNat x => x:.xs -> Int
dimHead = cnat . proxyHead
where proxyHead :: x:.xs -> Proxy x
proxyHead _ = Proxy
dimTail :: x:.xs -> xs
dimTail (_:.xs) = xs
pdimHead :: KnownNat x => proxy (x:.xs) -> Int
pdimHead = cnat . proxyHead
where proxyHead :: proxy (x:.xs) -> Proxy x
proxyHead _ = Proxy
pdimTail :: proxy (x:.xs) -> Proxy xs
pdimTail _ = Proxy
class (Dim n, Range (ToPeano (Size n))) => Ranged n
instance (Dim n, Range (ToPeano (Size n))) => Ranged n
data Peano = Zero | Succ Peano
type family ToPeano (n :: Nat) :: Peano where
ToPeano 0 = Zero
ToPeano n = Succ (ToPeano (n1))
type family Size (dim :: *) :: Nat where
Size (x:.Z) = x
Size (x:.xs) = x * Size xs
type family InRange (a :: *) (b :: *) :: Bool where
InRange Z Z = True
InRange (x:.xs) (y:.ys) = And (x <=? y) (InRange xs ys)
type family And (a :: Bool) (b :: Bool) :: Bool where
And True True = True
And a b = False
class Range (n :: Peano) where
swithRange_ :: (Dim o, Applicative m) => Tagged n o -> (o -> m ()) -> m ()
sfoldrRange_ :: Dim o => Tagged n o -> (o -> b -> b) -> b -> b
sfoldlRange_ :: Dim o => Tagged n o -> (b -> o -> b) -> b -> b
swithRangeIndices_ :: Applicative m => Tagged n Int -> (Int -> m ()) -> m ()
sfoldrRangeIndices_ :: Tagged n Int -> (Int -> b -> b) -> b -> b
sfoldlRangeIndices_ :: Tagged n Int -> (b -> Int -> b) -> b -> b
instance Range Zero where
swithRange_ _ _ = pure ()
sfoldrRange_ _ _ z = z
sfoldlRange_ _ _ z = z
swithRangeIndices_ _ _ = pure ()
sfoldrRangeIndices_ _ _ z = z
sfoldlRangeIndices_ _ _ !z = z
instance Range n => Range (Succ n) where
swithRange_ !i f = f (unTagged i) *> swithRange_ (nextTagged i) f
sfoldrRange_ i f z = f (unTagged i) (sfoldrRange_ (nextTagged i) f z)
sfoldlRange_ !i f !z = sfoldlRange_ (nextTagged i) f (f z (unTagged i))
swithRangeIndices_ !i f = f (unTagged i) *> swithRangeIndices_ (nextTaggedI i) f
sfoldrRangeIndices_ i f z =
f (unTagged i) (sfoldrRangeIndices_ (nextTaggedI i) f z)
sfoldlRangeIndices_ !i f !z =
sfoldlRangeIndices_ (nextTaggedI i) f (f z (unTagged i))
nextTagged :: Dim a => Tagged (Succ n) a -> Tagged n a
nextTagged (Tagged a) = Tagged (next a)
nextTaggedI :: Tagged (Succ n) Int -> Tagged n Int
nextTaggedI (Tagged a) = Tagged (a+1)
tagPeano :: Dim n => n -> Tagged (ToPeano (Size n)) n
tagPeano = Tagged
tagPeanoI :: proxy n -> Tagged (ToPeano (Size n)) Int
tagPeanoI _ = Tagged 0
bounds :: (Dim a, Bounded a) => Proxy a -> (a, a)
bounds _ = (zero, maxBound)
cnat :: KnownNat n => proxy (n :: Nat) -> Int
cnat = fromInteger . natVal