module Data.Index
(
Index
, Dim(zero,unit,size,rank,correct,zipMin,toIndex,fromIndex,next'
,prev')
, Rank(..)
, (:.)(..), Z(..)
, Ranged
, Range()
, Peano(..)
, ToPeano
, Size
, sfoldlRange
, sfoldrRange
, swithRange
, foldlRange
, foldrRange
, withRange
, sfoldlRangeIndices
, sfoldrRangeIndices
, swithRangeIndices
, foldlRangeIndices
, foldrRangeIndices
, withRangeIndices
, bounds
, range
, srange
, dimHead, dimTail
, pdimHead, pdimTail
, dim
, 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 (Bounded n, Integral n, Show n, Read n, Ix.Ix n, Dim n) => Index n
instance (Bounded n, Integral n, Show n, Read n, Ix.Ix n, Dim n) => Index n
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
class Ord n => Dim n where
zero :: n
unit :: n
rank :: n -> Int
size :: Proxy n -> Int
reflect' :: Proxy n -> n
reflect :: n
reflect = reflect' Proxy
next :: n -> n
prev :: n -> n
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
overHead :: (Int -> Int) -> n -> n
lastDim :: Proxy n -> n
zipMin :: n -> n -> n
instance Dim Z where
zero = Z
unit = Z
rank _ = 0
size _ = 1
reflect' _ = Z
next _ = error "next: Z"
prev _ = error "prev: Z"
next' _ = Z
prev' _ = Z
toIndex _ = 0
fromIndex' _ i | i > 0 = error "fromIndex: index too large"
| otherwise = Z
correct _ = Z
correctOnce _ = Z
overHead _ _ = Z
lastDim _ = Z
zipMin _ _ = Z
instance (KnownNat x, Dim xs) => Dim (x:.xs) where
zero = 0 :. zero
unit = 1 :. unit
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
| xs' <- next xs, xs' > zero = 0 :. xs'
| otherwise = error "next: index already at maxBound"
prev (x:.xs)
| x > 0 = (x1) :. xs
| xs > zero = 0 :. prev xs
| otherwise = error "prev: index below zero"
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 = lastDim Proxy
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)
| x < dimHead d = x:.correct xs
| otherwise = error "correct: index too large"
correctOnce d@(x:.xs)
| x < dimHead d = x:.xs
| otherwise = error "correctOnce: index too large"
overHead f (x:.xs) = f x :. xs
lastDim d = (pdimHead d 1) :. lastDim (pdimTail d)
zipMin (x:.xs) (y:.ys) = min x y :. zipMin xs ys
range :: Dim n => [n]
range = foldrRange Proxy (:) []
srange :: Ranged n => [n]
srange = sfoldrRange Proxy (:) []
proxyOf :: a -> Proxy a
proxyOf = const Proxy
foldlRange :: Dim a => Proxy a -> (b -> a -> b) -> b -> b
foldlRange _ cons = go zero
where
top = lastDim Proxy
go !i !acc
| i < top = go (next' i) (cons acc i)
| otherwise = cons acc top
foldrRange :: Dim a => Proxy a -> (a -> b -> b) -> b -> b
foldrRange _ cons nil = go zero
where
top = lastDim Proxy
go !i
| i < top = i `cons` go (next' i)
| otherwise = cons top nil
withRange :: (Applicative m, Dim a) => Proxy a -> (a -> m ()) -> m ()
withRange m f = foldrRange m (\d acc -> acc *> f d) (pure ())
foldlRangeIndices :: Dim a => Proxy a -> (b -> Int -> b) -> b -> b
foldlRangeIndices m cons = go 0
where
go !i !acc
| i < size m = go (i+1) (cons acc i)
| otherwise = acc
foldrRangeIndices :: Dim a => Proxy a -> (Int -> b -> b) -> b -> b
foldrRangeIndices m cons nil = go 0
where
go !i
| i < size m = i `cons` go (i+1)
| otherwise = nil
withRangeIndices :: (Applicative m, Dim a) => Proxy a -> (Int -> m ()) -> m ()
withRangeIndices m f = foldrRangeIndices m (\d acc -> acc *> f d) (pure ())
instance Num Z where
_ + _ = Z
_ _ = Z
_ * _ = Z
negate _ = Z
signum _ = Z
abs _ = Z
fromInteger _ = Z
instance (Num xs, Dim (x:.xs)) => Num (x:.xs) where
x:.xs + y:.ys = correctOnce ((x+y) :. (xs+ys))
x:.xs y:.ys = correctOnce ((xy) :. (xsys))
x:.xs * y:.ys = correctOnce ((x*y) :. (xs*ys))
negate (x:.xs) = negate x :. negate xs
signum (x:.xs) = signum x :. signum xs
abs (x:.xs) = abs x :. abs xs
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), Num xs) => Enum (x:.xs) where
toEnum = fromIndex
fromEnum = toIndex
succ = next
pred = prev
enumFrom !l
| l < maxBound = l : enumFrom (next' l)
| otherwise = [l]
enumFromThen !l !j
| l < maxBound = l : enumFromThen (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)) => Integral (x:.xs) where
div (x:.xs) (y:.ys) = div x y :. div xs ys
mod (x:.xs) (y:.ys) = mod x y :. mod xs ys
quot (x:.xs) (y:.ys) = quot x y :. quot xs ys
rem (x:.xs) (y:.ys) = rem x y :. rem xs ys
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 (x:.xs) (y:.ys) = correctOnce ((x+y):.mappend xs ys)
instance KnownNat s => Applicative ((:.) s) where
pure x = 1:.x
d@(ix₁ :. f) <*> ix₂ :. x = ((ix₁*ix₂) `rem` dimHead d) :. f x
(*>) = (>>)
(<*) = flip (>>)
instance KnownNat s => Monad ((:.) s) where
return x = 1:.x
d@(ix₁:.a) >>= f = case f a of
ix₂ :. b -> ((ix₁*ix₂) `mod` dimHead d) :. b
instance Ix.Ix Z where
range _ = [Z]
index _ _ = 0
inRange _ _ = True
rangeSize _ = 0
instance (Dim (x:.xs), Num xs) => Ix.Ix (x:.xs) where
range = uncurry enumFromTo
index _ c = toIndex c
inRange _ c = minBound < c && c < maxBound
rangeSize (l,_) = size (Proxy `proxify` l)
where
proxify :: Proxy a -> a -> Proxy a
proxify p _ = p
dim :: QuasiQuoter
dim = QuasiQuoter
{ quoteExp = \s -> [| Proxy :: $(quoteType dim s) |]
, quoteType = \s ->
let cons a b = [t| $(litT $ numTyLit a) :. $b |]
nil = [t| Z |]
in [t| Proxy $(foldr cons nil . map read . words $ s) |]
, quotePat = \_ -> [p| Proxy |]
, quoteDec = error "dim"
}
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
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
swithRange :: (Applicative m, Ranged o) => Proxy o -> (o -> m ()) -> m ()
swithRange _ (f :: o -> m ()) = swithRange_ (Tagged zero :: Tagged (ToPeano (Size o)) o) f
sfoldrRange :: Ranged o => Proxy o -> (o -> b -> b) -> b -> b
sfoldrRange _ (f :: o -> b -> b) z = sfoldrRange_ (tagPeano zero) f z
sfoldlRange :: Ranged o => Proxy o -> (b -> o -> b) -> b -> b
sfoldlRange _ (f :: b -> o -> b) !z = sfoldlRange_ (tagPeano zero) f z
swithRangeIndices :: (Applicative m, Ranged o) => Proxy o -> (Int -> m ()) -> m ()
swithRangeIndices dim f = swithRangeIndices_ (tagPeanoI dim) f
sfoldrRangeIndices :: Ranged o => Proxy o -> (Int -> b -> b) -> b -> b
sfoldrRangeIndices dim f z = sfoldrRangeIndices_ (tagPeanoI dim) f z
sfoldlRangeIndices :: Ranged o => Proxy o -> (b -> Int -> b) -> b -> b
sfoldlRangeIndices dim f z = sfoldlRangeIndices_ (tagPeanoI dim) f z
bounds :: (Dim a, Bounded a) => Proxy a -> (a, a)
bounds _ = (zero, maxBound)
cnat :: KnownNat n => proxy (n :: Nat) -> Int
cnat = fromInteger . natVal