module Numeric.Dimensions.Idx
(
Idx (..)
, appendIdx, splitIdx
) where
import Control.Arrow (first)
import GHC.Exts (IsList (..))
import Unsafe.Coerce (unsafeCoerce)
import Numeric.Dimensions.Dim
import Numeric.Dimensions.List
data Idx (ds :: [Nat]) where
Z :: Idx '[]
(:!) :: !Int -> !(Idx ds) -> Idx (d ': ds)
infixr 5 :!
idxToList :: Idx ds -> [Int]
idxToList Z = []
idxToList (x :! xs) = x : idxToList xs
idxFromList :: [Int] -> Idx ds
idxFromList [] = unsafeCoerce Z
idxFromList (x:xs) = unsafeCoerce $ x :! unsafeCoerce (idxFromList xs)
succIdx :: Dim xs -> Idx xs -> Idx xs
succIdx _ Z = Z
succIdx ((Dn :: Dim d) :* ds) (i :! is) | i >= dimVal' @d = 1 :! succIdx ds is
| otherwise = succ i :! is
predIdx :: Dim xs -> Idx xs -> Idx xs
predIdx _ Z = Z
predIdx ((Dn :: Dim d) :* ds) (i :! is) | i <= 1 = dimVal' @d :! predIdx ds is
| otherwise = pred i :! is
toIdx :: Dim xs -> Int -> Idx xs
toIdx D _ = Z
toIdx ((Dn :: Dim d) :* ds) off = case divMod off (dimVal' @d) of
(off', i) -> i+1 :! toIdx ds off'
fromIdx :: Dim xs -> Idx xs -> Int
fromIdx _ Z = 0
fromIdx ((Dn :: Dim d) :* ds) (i :! is) = i 1 + dimVal' @d * fromIdx ds is
diffIdx :: Dim xs -> Idx xs -> Idx xs -> Int
diffIdx _ Z _ = 0
diffIdx ((Dn :: Dim d) :* ds) (i1:!is1) (i2:!is2) = i1 i2
+ dimVal' @d * diffIdx ds is1 is2
stepIdx :: Dim ds -> Int -> Idx ds -> Idx ds
stepIdx _ _ Z = Z
stepIdx ((Dn :: Dim d) :* ds) di (i:!is)
= case divMod (di + i 1) (dimVal' @d) of
(0 , i') -> i'+1 :! is
(di', i') -> i'+1 :! stepIdx ds di' is
appendIdx :: forall (as :: [Nat]) (b :: Nat)
. Idx as -> Int -> Idx (as +: b)
appendIdx Z i = i :! Z
appendIdx (j :! js) i = unsafeCoerce $ j :! (unsafeCoerce (appendIdx js i) :: Idx (Tail (as +: b)))
splitIdx :: forall (as :: [Nat]) (bs :: [Nat])
. FiniteList as => Idx (as ++ bs) -> (Idx as, Idx bs)
splitIdx = splitN (order @_ @as)
where
splitN :: Int -> Idx (as ++ bs) -> (Idx as, Idx bs)
splitN 0 js = unsafeCoerce (Z, js)
splitN n (j :! js) = first (unsafeCoerce . (j :!))
$ splitN (n1) (unsafeCoerce js)
splitN _ Z = unsafeCoerce (Z, Z)
instance Show (Idx ds) where
show Z = "Idx Ø"
show xs = "Idx" ++ foldr (\i s -> " " ++ show i ++ s) "" (idxToList xs)
instance Eq (Idx ds) where
Z == Z = True
(a:!as) == (b:!bs) = a == b && as == bs
Z /= Z = False
(a:!as) /= (b:!bs) = a /= b || as /= bs
instance Num (Idx '[n]) where
(a:!Z) + (b:!Z) = (a+b) :! Z
(a:!Z) (b:!Z) = (ab) :! Z
(a:!Z) * (b:!Z) = (a*b) :! Z
signum (a:!Z) = signum a :! Z
abs (a:!Z) = abs a :! Z
fromInteger i = fromInteger i :! Z
instance Ord (Idx ds) where
compare Z Z = EQ
compare (a:!as) (b:!bs) = compare as bs `mappend` compare a b
instance Dimensions ds => Bounded (Idx ds) where
maxBound = f (dim @ds)
where
f :: forall ns . Dim ns -> Idx ns
f D = Z
f ((Dn :: Dim n) :* ds) = dimVal' @n :! f ds
minBound = f (dim @ds)
where
f :: forall (ns :: [Nat]) . Dim ns -> Idx ns
f D = Z
f (Dn :* ds) = 1 :! f ds
instance IsList (Idx ds) where
type Item (Idx ds) = Int
fromList = idxFromList
toList = idxToList
instance Dimensions ds => Enum (Idx ds) where
succ = succIdx (dim @ds)
pred = predIdx (dim @ds)
toEnum = toIdx (dim @ds)
fromEnum = fromIdx (dim @ds)
enumFrom x = take (diffIdx ds maxBound x + 1) $ iterate (succIdx ds) x
where
ds = dim @ds
enumFromTo x y | x >= y = take (diffIdx ds x y + 1) $ iterate (predIdx ds) x
| otherwise = take (diffIdx ds y x + 1) $ iterate (succIdx ds) x
where
ds = dim @ds
enumFromThen x x' = take n $ iterate (stepIdx ds dn) x
where
ds = dim @ds
dn = diffIdx ds x' x
n = 1 + if dn == 0 then 0
else if dn > 0 then diffIdx ds maxBound x `div` dn
else diffIdx ds x minBound `div` negate dn
enumFromThenTo x x' y = take n $ iterate (stepIdx ds dn) x
where
ds = dim @ds
dn = diffIdx ds x' x
n = 1 + if dn == 0 then 0
else diffIdx ds y x `div` dn