Rank.hs

{-# LANGUAGE RankNTypes, ExistentialQuantification, FlexibleInstances, UndecidableInstances, TypeFamilyDependencies, MultiParamTypeClasses, ScopedTypeVariables, FlexibleContexts, FunctionalDependencies, ConstraintKinds, QuantifiedConstraints, TypeOperators #-}
module Futhask.Array.Rank where
import Data.Int
import Foreign.Ptr
import Foreign.Storable
import System.IO.Unsafe

    
class Dim r where
    type Size   r = res | res -> r 
    type Index  r = res | res -> r
    type DimFun r a = res | res -> r
    type List   r a 
    sizeAsArgs      :: DimFun r a -> Size r -> IO a
    sizeFromPointer :: Ptr Int64 -> IO (Size r)
    inBounds        :: Size r -> Index r -> Bool
    linearSize      :: Size r -> Int 
    linearIndex     :: Size r -> Index r -> Int
    sizeString      :: Size r -> String
    indexString     :: Index r -> String
    indices         :: Size r -> [Index r]
    mapToList       :: (Index r -> a) -> Size r -> List r a
ifInBounds s i a = if inBounds s i 
    then a 
    else error ("Index " ++ indexString i ++ " out of bounds of array with size " ++ sizeString s ++ ".")
maybeInBounds s i a = if inBounds s i then Just a else Nothing
data R0
data R1
data R2
data R3
data R4
--data Rx
---- recursive indextype
--instance (Dim R) => Rx R 
--    type Size  (Rx R) = (Int, Size  R)
--    type Index (Rx R) = (Int, Index R)
--    inBounds = 
--    linearSize
--    linearIndex
instance Dim R0 where
    type Size   R0 = ()
    type Index  R0 = ()
    type DimFun R0 a = IO a
    type List R0 a = a
    sizeAsArgs df s = df
    sizeFromPointer _ = return ()
    inBounds _ _ = True
    linearSize _ = 1
    linearIndex _ _ = 0 
    sizeString s = "()"
    indexString i = "()"
    indices s = [()]
    mapToList f s = f () 
instance Dim R1 where
    type Size   R1 = Int
    type Index  R1 = Int
    type DimFun R1 a = Int64 -> IO a
    type List   R1 a = [a]
    sizeAsArgs df s = df (fromIntegral s) 
    sizeFromPointer p = fmap fromIntegral (peek p)
    inBounds s i = 0 <= i && i < s

    linearSize s = s
    linearIndex _ i = i
    sizeString s = show s
    indexString i = show i
    indices s = [0..s-1]
    mapToList f s = map f [0..s-1]
instance Dim R2 where
    type Size   R2 = (Int, Int)
    type Index  R2 = (Int, Int)
    type DimFun R2 a = Int64 -> Int64 -> IO a
    type List   R2 a = [[a]]
    sizeAsArgs df (s0, s1) = df (fromIntegral s0) (fromIntegral s1)
    sizeFromPointer p = do
        s0 <- fmap fromIntegral (peekElemOff p 0)
        s1 <- fmap fromIntegral (peekElemOff p 1)
        return (s0, s1)
    inBounds (s0, s1) (i0, i1) = inBounds s0 i0 && inBounds s1 i1
    linearSize (s0, s1) = linearSize s0 * s1
    linearIndex (s0, s1) (i0, i1) = linearIndex s0 i0 * s1 + i1
    sizeString s = show s
    indexString i = show i
    indices (s0, s1) = (,) <$> [0..s0-1] <*> [0..s1-1]
    mapToList f (s0, s1) = map (\i0 -> map (\i1 -> f (i0, i1))[0..s1-1]) [0..s0-1]
     
instance Dim R3 where
    type Size   R3 = (Int, Int, Int)
    type Index  R3 = (Int, Int, Int)
    type DimFun R3 a = Int64 -> Int64 -> Int64 -> IO a
    type List   R3 a = [[[a]]]
    sizeAsArgs df (s0, s1, s2) = df (fromIntegral s0) (fromIntegral s1) (fromIntegral s2)
    sizeFromPointer p = do
        s0 <- fmap fromIntegral (peekElemOff p 0)
        s1 <- fmap fromIntegral (peekElemOff p 1)
        s2 <- fmap fromIntegral (peekElemOff p 2)
        return (s0, s1, s2)
    inBounds (s0, s1, s2) (i0, i1, i2) = inBounds (s0, s1) (i0, i1) && inBounds s2 i2 
    linearSize (s0, s1, s2) = linearSize (s0, s1) * s2
    linearIndex (s0, s1, s2) (i0, i1, i2) = linearIndex (s0, s1) (i0, i1) * s2 + i2
    sizeString s = show s
    indexString i = show i
    indices (s0, s1, s2) = (,,) <$> [0..s0-1] <*> [0..s1-1] <*> [0..s2-1]
    mapToList f (s0, s1, s2) =
        map (\i0 -> 
            map (\i1 -> 
                map (\i2 -> f (i0, i1, i2)) 
                [0..s2-1]) 
            [0..s1-1]) 
        [0..s0-1]
instance Dim R4 where
    type Size   R4 = (Int, Int, Int, Int)
    type Index  R4 = (Int, Int, Int, Int)
    type DimFun R4 a = Int64 -> Int64 -> Int64 -> Int64 -> IO a
    type List   R4 a = [[[[a]]]]
    sizeAsArgs df (s0, s1, s2, s3) = df (fromIntegral s0) (fromIntegral s1) (fromIntegral s2) (fromIntegral s3)
    sizeFromPointer p = do
        s0 <- fmap fromIntegral (peekElemOff p 0)
        s1 <- fmap fromIntegral (peekElemOff p 1)
        s2 <- fmap fromIntegral (peekElemOff p 2)
        s3 <- fmap fromIntegral (peekElemOff p 3)
        return (s0, s1, s2, s3)
    inBounds (s0, s1, s2, s3) (i0, i1, i2, i3) = inBounds (s0, s1, s2) (i0, i1, i2) && inBounds s3 i3
    linearSize (s0, s1, s2, s3) = s0 * s1 * s2 * s3
    linearIndex (s0, s1, s2, s3) (i0, i1, i2, i3) = linearIndex (s0, s1, s2) (i0, i1, i2) * s3 + i3
    sizeString s = show s
    indexString i = show i
    indices (s0, s1, s2, s3) = (,,,) <$> [0..s0-1] <*> [0..s1-1] <*> [0..s2-1] <*> [0..s3-1]
    mapToList f (s0, s1, s2, s3) =
        map (\i0 -> 

            map (\i1 -> 
                map (\i2 ->
                    map (\i3 -> f (i0, i1, i2, i3)) 
                    [0..s3-1]) 
                [0..s2-1]) 
            [0..s1-1]) 
        [0..s0-1]
--class (Dim (Rank a)) => Indexable a where
--    type Rank a
--    size :: a -> Size (Rank a)
class (Dim (Rank a)) => Array a where
    type Rank a 
    type Elem a
    type Mutable a = res | res -> a
    size :: a -> Size (Rank a)
    mutSize :: Mutable a -> Size (Rank a)
    readElem' :: a -> Index (Rank a) -> Elem a 
    peekElem' :: Mutable a -> Index (Rank a) -> IO (Elem a)
    pokeElem' :: Mutable a -> Index (Rank a) -> Elem a -> IO ()
    copy :: Mutable a -> IO (Mutable a)
    unsafeThaw :: a -> IO (Mutable a)
    unsafeFreeze :: Mutable a -> IO a
    scratch :: Size (Rank a) -> IO (Mutable a)