Object.hs

{-# LANGUAGE RankNTypes, ExistentialQuantification, FlexibleInstances, UndecidableInstances, TypeFamilyDependencies, MultiParamTypeClasses, ScopedTypeVariables, FlexibleContexts, FunctionalDependencies, ConstraintKinds, QuantifiedConstraints, TypeOperators #-}
{-|
Defines various classes for futhark objects and some utility functions to manage them.
-}

module Futhask.Object where
import Data.Kind
import Data.Word
import Data.Int
import Data.ByteString as B
import Foreign.C
import Foreign.C.String
import Foreign.Ptr
import Foreign.Storable as S
import Foreign.ForeignPtr
import Foreign.Marshal.Array as A
import Foreign.Marshal.Alloc
import Control.Monad.Base
import Control.Monad.Trans
import Control.Concurrent
import Control.Concurrent.MVar
import System.Mem

import Futhask.Array 
import Futhask.Array.Rank 
import Futhask.Context
import Futhask.Monad

-- * Wrapping Pointers and Managing Memory
-- | The basic functions for wrapping and memory management
class (Context (ObjectContext fo)) => FutharkObject fo where
    type RawObject fo = raw | raw -> fo -- injectivity broke for some reason
    type ObjectContext fo
    wrapRawObject :: ReferenceCounter -> ForeignPtr (RawObject fo) -> fo c
    objectCounter :: fo c -> ReferenceCounter
    objectPointer :: fo c -> ForeignPtr (RawObject fo)
    freeRawObject :: Ptr (RawContext (ObjectContext fo)) -> Ptr (RawObject fo) -> IO Int

-- **Wrapping
-- | Use the raw pointer and return the result    
withRawObject :: (FutharkObject fo) => fo c -> (Ptr (RawObject fo) -> IO a) -> IO a
withRawObject fo = withForeignPtr (objectPointer fo)

-- | Wrap raw pointer with a reference counter and context
wrapFO :: (FutharkObject fo) => ObjectContext fo -> Ptr (RawObject fo) -> IO (fo c) 
wrapFO context ptr = do
    refCounter <- newRefCounter
    pointer <- newForeignPtrWithRef contextRefcounter ptr (withForeignPtr contextPtr (\cptr -> freeRawObject cptr ptr)) 
    return $! wrapRawObject refCounter pointer
    where (contextRefcounter, contextPtr) = unwrapRawContext context


-- **Manual Memory Management

-- | Increments reference counter, useful for concurrent use.
addReferenceFO :: (MonadIO m, FutharkObject fo) => fo c -> FutT (ObjectContext fo) c m ()
addReferenceFO fo = lift . liftIO $ increment (objectCounter fo)

-- | Decrements reference counter, if the counter is zero the object will be finalized, and can no longer be used.
finalizeFO :: (MonadIO m, FutharkObject fo) => fo c -> FutT (ObjectContext fo) c m ()
finalizeFO fo = lift . liftIO $
    let refCounter = objectCounter fo
     in isZero refCounter >>= \b -> if b 
        then finalizeForeignPtr (objectPointer fo) 
        else decrement refCounter




-- *Value composition and decomposition
class (FutharkObject fo) => FutharkComposite fo where
    type Unpacked fo :: Type -> Type
    pack :: (Monad m) => Unpacked fo c -> FutT (ObjectContext fo) c m (fo c)
    unpack :: (Monad m) => fo c -> FutT (ObjectContext fo) c m (Unpacked fo c)

-- *Projection

-- | Takes a projection function from the @Raw@ interface and applies it.
extractWith 
    :: (FutharkObject fo, Storable a) 
    => (Ptr (RawContext (ObjectContext fo)) -> Ptr a -> Ptr (RawObject fo) -> IO Int)
    -> ObjectContext fo -> fo c -> IO a
extractWith f c o = alloca (\ptr -> withRawObject o (\ro -> inContextWithError c (\rc -> f rc ptr ro)) >> S.peek ptr)

-- | Same as `extractWith` but also wraps the extracted object. Only works on futhark objects.
extractFOWith 
    :: (FutharkObject fo1, FutharkObject fo2, ObjectContext fo2 ~ ObjectContext fo1) 
    => (Ptr (RawContext (ObjectContext fo1)) -> Ptr (Ptr (RawObject fo2)) -> Ptr (RawObject fo1) -> IO Int)
    -> ObjectContext fo1 -> fo1 c -> IO (fo2 c)    
extractFOWith f c o = extractWith f c o >>= wrapFO c



-- *Array functions
class (FutharkObject arr) => FutIndexable arr where
    type FutRank arr

-- | Basic array functions for primitive arrays
class (FutharkObject fo, FutIndexable fo, Context ctx, ObjectContext fo ~ ctx, Dim rank, FutRank fo ~ rank, Storable elem, PrimElem fo ~ elem) 
    => FutharkPrimArray' fo ctx rank elem | fo -> ctx rank elem, ctx rank elem -> fo where
    type PrimElem fo
    shapeFPA  :: Ptr (RawContext ctx) -> Ptr (RawObject fo) -> IO (Ptr Int64) 
    valuesFPA :: Ptr (RawContext ctx) -> Ptr (RawObject fo) -> Ptr elem -> IO Int
    newFPA    :: Ptr (RawContext ctx) -> Ptr elem -> DimFun rank (Ptr (RawObject fo))
    indexFPA  :: Ptr (RawContext ctx) -> Ptr elem -> Ptr (RawObject fo) -> DimFun rank Int

class (FutharkObject fo, FutharkPrimArray' fo (ObjectContext fo) (FutRank fo) (PrimElem fo)) => FutharkPrimArray fo
instance (FutharkObject fo, FutharkPrimArray' fo (ObjectContext fo) (FutRank fo) (PrimElem fo)) => FutharkPrimArray fo

-- | Basic array functions for opaque arrays
class (FutharkObject fo, FutIndexable fo, Context ctx, ObjectContext fo ~ ctx, Dim rank, FutRank fo ~ rank, FutharkObject elem, ObjectContext elem ~ ctx, OpaqueElem fo ~ elem) 
    => FutharkOpaqueArray' fo ctx rank elem | fo -> ctx rank elem, ctx rank elem -> fo where
    type OpaqueElem fo :: Type -> Type
    shapeFOA  :: Ptr (RawContext ctx) -> Ptr (RawObject fo) -> IO (Ptr Int64) 
    newFOA    :: Ptr (RawContext ctx) -> Ptr (RawObject elem) -> DimFun rank (Ptr (RawObject fo))
    indexFOA  :: Ptr (RawContext ctx) -> Ptr (Ptr (RawObject elem)) -> Ptr (RawObject fo) -> DimFun rank Int
    setFOA    :: Ptr (RawContext ctx) -> Ptr (RawObject fo) -> ptr (RawObject elem) -> DimFun rank Int

class FutharkOpaqueArray' fo (ObjectContext fo) (FutRank fo) (OpaqueElem fo) => FutharkOpaqueArray fo
instance FutharkOpaqueArray' fo (ObjectContext fo) (FutRank fo) (OpaqueElem fo) => FutharkOpaqueArray fo




-- * Transfer between Contexts

-- | Allows transfer of objects between contexts    
class (FutharkObject fo1, FutharkObject fo2) => Transferable fo1 fo2 where
    transferUp   :: (Monad m) => fo2 c2 -> FutT (ObjectContext fo1) c1 (FutT (ObjectContext fo2) c2 m) (fo1 c1)
    transferDown :: (Monad m) => fo1 c1 -> FutT (ObjectContext fo1) c1 (FutT (ObjectContext fo2) c2 m) (fo2 c2)

instance (FutharkOpaque fo) => Transferable fo fo where
    transferUp fo = unsafeLiftFromIO $ \ctx -> alloca $ \sp -> alloca $ \pp -> withRawObject fo $ \op2 -> do
        S.poke pp nullPtr
        inContextWithError ctx $ \cp -> storeOpaque cp op2 pp sp
        sync ctx
        p <- S.peek pp
        op <- inContext ctx $ \cp -> restoreOpaque cp p
        free p
        if op == nullPtr
            then error "Restoration error in transfer"
            else wrapFO ctx op        
    transferDown fo = unsafeLiftFromIO $ \ctx -> alloca $ \sp -> alloca $ \pp -> withRawObject fo $ \op2 -> do
        S.poke pp nullPtr
        inContextWithError ctx $ \cp -> storeOpaque cp op2 pp sp
        sync ctx
        p <- S.peek pp
        op <- inContext ctx $ \cp -> restoreOpaque cp p
        free p
        if op == nullPtr
            then error "Restoration error in transfer"
            else wrapFO ctx op        

-- *Reading and Writing to Files

-- | Serialization of opaques
class (FutharkObject fo) => FutharkOpaque fo where
    storeOpaque   :: Ptr (RawContext (ObjectContext fo)) -> Ptr (RawObject fo) -> Ptr (Ptr CChar) -> Ptr CSize -> IO Int
    restoreOpaque :: Ptr (RawContext (ObjectContext fo)) -> Ptr CChar -> IO (Ptr (RawObject fo))

-- | encoding and decoding objects
class (Monad m) => Serializable m a where
    encode :: a -> m ByteString
    decode :: ByteString -> m (Maybe a)

-- | write object to file
store fn a = encode a >>= lift . B.writeFile fn

-- | read object from file
load fn = lift (B.readFile fn) >>= decode 

instance (FutharkOpaque fo, FutharkObject fo, ObjectContext fo ~ ctx, Monad m)
    => Serializable (FutT ctx c m) (fo c) where 
    encode fo = unsafeLiftFromIO $ \ctx -> alloca $ \sp -> alloca $ \pp -> withRawObject fo $ \op -> do
        S.poke pp nullPtr
        inContextWithError ctx $ \cp -> storeOpaque cp op pp sp
        sync ctx
        p <- S.peek pp
        s <- S.peek sp
        bs <- packCStringLen (p, fromIntegral s)
        free p
        pure bs
    decode bs = unsafeLiftFromIO $ \ctx -> useAsCString bs $ \cs -> do
        op <- inContext ctx $ \cp -> restoreOpaque cp cs
        if op == nullPtr
            then pure Nothing
            else fmap Just $ wrapFO ctx op

-- * Conversion 

-- | Convert values to and from Futhark
class FutharkObject fo => Convertible fo where
    type Native fo
    toFuthark :: Monad m => Native fo -> FutT (ObjectContext fo) c m (fo c)
    fromFuthark :: Monad m => fo c -> FutT (ObjectContext fo) c m (Native fo)