fix instance Control.DeepSeq.NFData BDD for GHC 7.10 with deepseq >= 1.4

-- -*- haskell -*-- -----------------------------------------------
-- |
-- Module      :  Data.Boolean.CUDD
-- Copyright   :  (C) 2002-2005, 2009 University of New South Wales, (C) 2009-2011 Peter Gammie
-- License     :  LGPL (see COPYING.LIB for details)
--
-- A binding for CUDD (Fabio Somenzi, University of Colorado).
--
-- Note this library is not thread-safe.
-------------------------------------------------------------------
module Data.Boolean.CUDD
    (
        BDD
    ,	module Data.Boolean

        -- * CUDD-specific functions
--          ,	gc
    ,	stats
    ,	reorder
    ,	dynamicReordering
    ,	varIndices
    ,	bddSize
    ,	printInfo
    ) where

-------------------------------------------------------------------
-- Dependencies.
-------------------------------------------------------------------

#include "cudd_im.h"

import Control.DeepSeq  ( NFData, rnf )
import Control.Monad	( foldM, liftM, mapAndUnzipM, zipWithM_ )

import Data.IORef	( IORef, newIORef, readIORef, writeIORef )
import Data.List	( genericLength )
import Data.Maybe	( fromJust, isJust )

import Data.Map         ( Map )
import qualified Data.Map as Map

import Foreign		( ForeignPtr, withForeignPtr, newForeignPtr, newForeignPtr_, finalizerFree, touchForeignPtr
                        , Ptr, advancePtr, castPtr, peek, pokeElemOff, ptrToIntPtr
                        , mallocArray )
import Foreign.C

import GHC.ForeignPtr   ( newConcForeignPtr )

import System.IO	( Handle, hIsReadable, hIsWritable )
-- import System.Mem	( performGC )
import System.Posix.IO	( handleToFd )
import System.IO.Unsafe	( unsafePerformIO )

import Data.Boolean

-------------------------------------------------------------------
-- Extra FFI functions.
-------------------------------------------------------------------

-- | A C file handle.
{#pointer *FILE -> CFile#}

-- | Convert a Haskell Handle into a C FILE *.
-- - FIXME: throw exception on error.
-- - suggested by Simon Marlow.
handleToCFile :: Handle -> IO FILE
handleToCFile h =
    do r <- hIsReadable h
       w <- hIsWritable h
       modestr <- newCString $ mode r w
       fd <- handleToFd h
       {#call unsafe fdopen#} (cToNum fd) modestr
    where mode :: Bool -> Bool -> String
          mode False False = error "Handle not readable or writable!"
          mode False True  = "w"
          mode True  False = "r"
          mode True  True  = "r+" -- FIXME

-- | The Haskell Handle is unusable after calling 'handleToCFile',
-- so provide a neutered "fprintf"-style thing here.
printCFile :: FILE -> String -> IO ()
printCFile file str =
    withCString str ({#call unsafe fprintf_neutered#} file)

--  Close a C FILE *.
-- - FIXME: throw exception on error.
-- closeCFile :: FILE -> IO ()
-- closeCFile cfile = do {#call unsafe fclose#} cfile
--                    return ()

cToNum :: (Num i, Integral e) => e -> i
cToNum  = fromIntegral . toInteger
{-# INLINE cToNum #-}

cFromEnum :: (Integral i, Enum e) => e -> i
cFromEnum  = fromIntegral . fromEnum
{-# INLINE cFromEnum #-}

-------------------------------------------------------------------
-- Types.
-------------------------------------------------------------------

-- | The BDD manager, which we hide from clients.
{#pointer *DdManager as DDManager newtype#}

-- | The arguments to 'Cudd_Init'.
{#enum CuddSubtableSize {} #}
{#enum CuddCacheSize {} #}

-- | Variable reordering tree options.
{#enum CuddMTRParams {} #}

-- | The abstract type of BDDs.
{#pointer *DdNode as BDD foreign newtype#}

{-# INLINE withBDD #-}

-- BDDs are just pointers, so there's no work to do when we @deepseq@ them.
instance Control.DeepSeq.NFData BDD where
  rnf x = seq x ()

-- Belt-and-suspenders equality checking and ordering.
instance Eq BDD where
    bdd0 == bdd1 = unsafePerformIO $
      withBDD bdd0 $ \bdd0p -> withBDD bdd1 $ \bdd1p ->
        return $ bdd0p == bdd1p

instance Ord BDD where
    bdd0 <= bdd1 = unsafePerformIO $
      withBDD bdd0 $ \bdd0p -> withBDD bdd1 $ \bdd1p ->
        return $ bdd0p == bdd1p || bdd0p < bdd1p

instance Show BDD where
-- Print out the BDD as a sum-of-products.
    showsPrec _ = sop

-------------------------------------------------------------------
-- Administrative functions.
-------------------------------------------------------------------

-- | The BDD manager. This needs to be invoked before any of the
-- following functions... which we arrange for automagically.
ddmanager :: DDManager
ddmanager = unsafePerformIO $
  {#call unsafe Cudd_Init as _cudd_Init #}
    0 0 (cFromEnum UNIQUE_SLOTS) (cFromEnum CACHE_SLOTS) 0
{-# NOINLINE ddmanager #-}

-- | A map to and from BDD variable numbers.
-- FIXME: the second one would be more efficiently Data.IntMap, but we don't use it often.
type VarMap = (Map String BDD, Map CUInt String)

-- | Tracks existing variables.
bdd_vars :: IORef VarMap
bdd_vars = unsafePerformIO $ newIORef (Map.empty, Map.empty)
{-# NOINLINE bdd_vars #-}

-- | Attaches a finalizer to a BDD object.
-- The call to "Cudd_Ref" happens in C to ensure atomicity.
-- Returning objects with initial refcount 0 is a bad design decision.
addBDDfinalizer :: Ptr BDD -> IO BDD
addBDDfinalizer bddp = liftM BDD $ newConcForeignPtr bddp bddf
    where bddf = {#call unsafe Cudd_RecursiveDeref as _cudd_RecursiveDeref#} ddmanager bddp
                   >> return ()
{-# INLINE addBDDfinalizer #-}

-- | Attaches a null finalizer to a 'BDD' object.
-- Used for variables and constants.
addBDDnullFinalizer :: Ptr BDD -> IO BDD
addBDDnullFinalizer bddp = fmap BDD $ newForeignPtr_ bddp
{-# INLINE addBDDnullFinalizer #-}

-- Simulate an "apply" function to simplify atomic refcount incrementing.
{#enum CuddBinOp {} #}

-------------------------------------------------------------------
-- Logical operations.
-------------------------------------------------------------------

-- | Allocate a variable.
newVar :: IO (Ptr BDD) -> String -> IO (Maybe CUInt, BDD)
newVar allocVar label =
      do (toBDD, fromBDD) <- readIORef bdd_vars
         case label `Map.lookup` toBDD of
           Just bdd -> return (Nothing, bdd)
           Nothing ->
             do bddp <- allocVar
                vid <- fmap cToNum $
                         {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#} bddp
                --putStrLn $ label ++ " -> " ++ (show (vid, bdd))
                bdd <- addBDDnullFinalizer bddp
                writeIORef bdd_vars ( Map.insert label bdd   toBDD
                                    , Map.insert vid   label fromBDD )
                return (Just vid, bdd)

instance BooleanVariable BDD where
    bvar label = unsafePerformIO $
      fmap snd $ newVar ({#call unsafe Cudd_bddNewVar as _cudd_bddNewVar#} ddmanager)  label

    bvars labels =
      case labels of
        [] -> []
        ls -> unsafePerformIO $
          do (vids, vars) <- mapAndUnzipM (newVar ({#call unsafe Cudd_bddNewVar as _cudd_bddNewVar#} ddmanager)) ls
             if all isJust vids
               then groupVars (fromJust (head vids)) (genericLength vids)
               else putStrLn $ "hBDD-CUDD warning: not grouping variables " ++ show ls
             return vars
      where
        -- FIXME why one or the other?
        groupVars vid len = makeTreeNode ddmanager vid len (cFromEnum CUDD_MTR_DEFAULT) >> return ()
        -- groupVars vid len = makeTreeNode ddmanager vid len (cFromEnum CUDD_MTR_FIXED) >> return ()
        makeTreeNode = {#call unsafe Cudd_MakeTreeNode as _cudd_MakeTreeNode#}

    unbvar bdd = unsafePerformIO $ bdd `seq`
                   withBDD bdd $ \bddp ->
                     do vid <- {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#} bddp
                        (_, fromBDD) <- readIORef bdd_vars
                        return $ case cToNum vid `Map.lookup` fromBDD of
                                   Nothing -> "(VID: " ++ show vid ++ ")"
                                   Just v  -> v

instance Boolean BDD where
    true  = cudd_constant {#call unsafe Cudd_ReadOne as _cudd_ReadOne#}
    false = cudd_constant {#call unsafe Cudd_ReadLogicZero as _cudd_ReadLogicZero#}

    (/\)  = bddBinOp AND
    neg x = unsafePerformIO $ withBDD x $ \xp ->
      {#call unsafe cudd_bddNot#} xp >>= addBDDfinalizer

    nand  = bddBinOp NAND
    (\/)  = bddBinOp OR
    nor   = bddBinOp NOR
    xor   = bddBinOp XOR
    -- (-->) = FIXME ???
    (<->) = bddBinOp XNOR

    ifthenelse i t e = unsafePerformIO $
      withBDD i $ \ip -> withBDD t $ \tp -> withBDD e $ \ep ->
        {#call unsafe cudd_bddIte#} ddmanager ip tp ep
          >>= addBDDfinalizer

cudd_constant :: (DDManager -> IO (Ptr BDD)) -> BDD
cudd_constant f = unsafePerformIO $ f ddmanager >>= addBDDnullFinalizer
{-# INLINE cudd_constant #-}

bddBinOp :: CuddBinOp -> BDD -> BDD -> BDD
bddBinOp op x y = unsafePerformIO $
  withBDD x $ \xp -> withBDD y $ \yp ->
    {#call unsafe cudd_BinOp#} ddmanager (cFromEnum op) xp yp >>= addBDDfinalizer
{-# INLINE bddBinOp #-}

instance QBF BDD where
    data Group BDD = MkGroup !BDD

    mkGroup = MkGroup . conjoin -- FIXME maybe use Cudd_bddComputeCube
    exists = cudd_exists
    forall = cudd_forall
    rel_product = cudd_rel_product

instance Substitution BDD where
  -- Cache substitution arrays.
  data Subst BDD = MkSubst {-# UNPACK #-} !Int
                           {-# UNPACK #-} !(ForeignPtr (Ptr BDD))
                           {-# UNPACK #-} !(ForeignPtr (Ptr BDD))

  mkSubst = cudd_mkSubst
  rename = cudd_rename
  substitute = error "CUDD substitute" -- cudd_substitute

cudd_mkSubst :: [(BDD, BDD)] -> Subst BDD
cudd_mkSubst subst = unsafePerformIO $
    do arrayx  <- mallocArray len
       arrayx' <- mallocArray len
       zipWithM_ (pokeSubst arrayx arrayx') subst [0..]
       fpx  <- newForeignPtr finalizerFree arrayx
       fpx' <- newForeignPtr finalizerFree arrayx'
       return (MkSubst len fpx fpx')
  where
    pokeSubst :: Ptr (Ptr BDD) -> Ptr (Ptr BDD) -> (BDD, BDD) -> Int -> IO ()
    pokeSubst arrayx arrayx' (v, v') i =
      do withBDD v  $ pokeElemOff arrayx  i
         withBDD v' $ pokeElemOff arrayx' i

    len = length subst
{-# INLINE cudd_mkSubst #-}

instance BDDOps BDD where
    get_bdd_ptr bdd = unsafePerformIO $ withBDD bdd (return . ptrToIntPtr)
    bif bdd = unsafePerformIO $
      do vid <- withBDD bdd {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#}
         {#call unsafe Cudd_bddIthVar as _cudd_bddIthVar#} ddmanager (cToNum vid)
           >>= addBDDnullFinalizer
    bthen bdd = unsafePerformIO $
      withBDD bdd $ \bddp ->
        {#call unsafe cudd_bddT#} bddp >>= addBDDfinalizer
    belse bdd = unsafePerformIO $
      withBDD bdd $ \bddp ->
        {#call unsafe cudd_bddE#} bddp >>= addBDDfinalizer

    reduce = cudd_reduce
    satisfy = cudd_satisfy
    support = cudd_support

-------------------------------------------------------------------
-- Implementations.
-------------------------------------------------------------------

withGroup :: Group BDD -> (Ptr BDD -> IO a) -> IO a
withGroup (MkGroup g) = withBDD g

cudd_exists :: Group BDD -> BDD -> BDD
cudd_exists group bdd = unsafePerformIO $ bdd `seq`
  withGroup group $ \groupp -> withBDD bdd $ \bddp ->
    {#call unsafe cudd_bddExistAbstract#} ddmanager bddp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_exists #-}

cudd_forall :: Group BDD -> BDD -> BDD
cudd_forall group bdd = unsafePerformIO $ bdd `seq`
  withGroup group $ \groupp -> withBDD bdd $ \bddp ->
    {#call unsafe cudd_bddUnivAbstract#} ddmanager bddp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_forall #-}

cudd_rel_product :: Group BDD -> BDD -> BDD -> BDD
cudd_rel_product group f g = unsafePerformIO $
  withGroup group $ \groupp -> withBDD f $ \fp -> withBDD g $ \gp ->
    {#call unsafe cudd_bddAndAbstract#} ddmanager fp gp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_rel_product #-}

-- | This function swaps variables, as it uses
-- @cudd_bddSwapVariables@. This is not exactly a "rename" behaviour.
cudd_rename :: Subst BDD -> BDD -> BDD
cudd_rename s@(MkSubst len fpx fpx') f = unsafePerformIO $ s `seq` f `seq`
  withBDD f $ \fp ->
    withForeignPtr fpx $ \arrayx -> withForeignPtr fpx' $ \arrayx' ->
      {#call unsafe cudd_bddSwapVariables#} ddmanager fp arrayx arrayx' (fromIntegral len)
        >>= addBDDfinalizer
{-# INLINE cudd_rename #-}

-- FIXME verify implementation.
cudd_reduce :: BDD -> BDD -> BDD
cudd_reduce f g = unsafePerformIO $
  withBDD f $ \fp -> withBDD g $ \gp ->
    {#call unsafe cudd_bddLICompaction#} ddmanager fp gp
      >>= addBDDfinalizer

-- FIXME verify implementation.
cudd_satisfy :: BDD -> BDD
cudd_satisfy bdd = unsafePerformIO $
  withBDD bdd ({#call unsafe cudd_satone#} ddmanager) >>= addBDDfinalizer

cudd_support :: BDD -> [BDD]
cudd_support bdd = unsafePerformIO $
  do varBitArray <- withBDD bdd ({#call unsafe Cudd_SupportIndex as _cudd_SupportIndex#} ddmanager)
     touchForeignPtr (case bdd of BDD fp -> fp)
     -- The array is as long as the number of variables allocated.
     (_toBDD, fromBDD) <- readIORef bdd_vars
     bdds <- foldM (toBDDs varBitArray) [] (Map.toList fromBDD)
     {#call unsafe cFree#} (castPtr varBitArray)
     return bdds
  where
    toBDDs varBitArray bdds (vid, var) =
      do varInSupport <- peek (advancePtr varBitArray (cToNum vid))
         return $ if varInSupport /= 0
                    then bvar var : bdds
                    else bdds

-------------------------------------------------------------------
-- Operations specific to this BDD binding.
-------------------------------------------------------------------

-- | Dump usage statistics to the given 'Handle'.
stats :: Handle -> IO ()
stats handle =
    do cfile <- handleToCFile handle
       printCFile cfile "CUDD stats\n"
       _ <- {#call unsafe Cudd_PrintInfo as _cudd_PrintInfo#} ddmanager cfile
       printCFile cfile "\nVariable groups\n"
       {#call unsafe cudd_printVarGroups#} ddmanager

----------------------------------------
-- Variable Reordering.
----------------------------------------

{#enum Cudd_ReorderingType as CUDDReorderingMethod {} deriving (Eq, Ord, Show)#}

decodeROM :: ReorderingMethod -> CInt
decodeROM rom = cFromEnum $ case rom of
  ReorderNone -> CUDD_REORDER_NONE
  ReorderSift -> CUDD_REORDER_SIFT
  ReorderSiftSym -> CUDD_REORDER_SYMM_SIFT
  ReorderStableWindow3 -> CUDD_REORDER_WINDOW3

-- | Reorder the variables now.
reorder :: ReorderingMethod -> IO ()
reorder rom = {#call unsafe Cudd_ReduceHeap as _cudd_ReduceHeap#} ddmanager (decodeROM rom) 1 >> return ()

-- | Set the dynamic variable ordering heuristic.
dynamicReordering :: ReorderingMethod -> IO ()
dynamicReordering rom = {#call unsafe Cudd_AutodynEnable as _cudd_AutodynEnable#} ddmanager (decodeROM rom) >> return ()

----------------------------------------
-- | Returns the relationship between BDD indices, BDD ids and
-- variable names. This may be useful for discovering what the
-- dynamic reordering is doing.
--
-- The intention of this function is to return
--   @[(position in variable order, immutable variable id, label)]@
-- so that the variable order can be saved.
----------------------------------------

varIndices :: IO [(CInt, CInt, String)]
varIndices = do (toBDD, _fromBDD) <- readIORef bdd_vars
                mapM procVar $ Map.toList toBDD
    where -- procVar :: (String, BDD) -> IO (CUInt, CUInt, String)
          procVar (var, bdd) =
              do -- FIXME CUDD inconsistently uses signed and unsigned ints.
                 vid <- fmap cToNum $
                          withBDD bdd {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#}
                 level <- {#call unsafe Cudd_ReadPerm as _cudd_ReadPerm#} ddmanager vid
                 return (level, vid, var)
{-# NOINLINE varIndices #-}

----------------------------------------
-- BDD Statistics.
----------------------------------------

-- | Determine the size of a BDD.
bddSize :: BDD -> Int
bddSize bdd = unsafePerformIO $
  do size <- withBDD bdd ({#call unsafe Cudd_DagSize as _cudd_DagSize#})
     return $ cToNum size

printInfo :: IO ()
printInfo = {#call unsafe print_stats#} ddmanager

-- -*- haskell -*-- -----------------------------------------------
-- |
-- Module      :  Data.Boolean.CUDD
-- Copyright   :  (C) 2002-2005, 2009 University of New South Wales, (C) 2009-2011 Peter Gammie
-- License     :  LGPL (see COPYING.LIB for details)
--
-- A binding for CUDD (Fabio Somenzi, University of Colorado).
--
-- Note this library is not thread-safe.
-------------------------------------------------------------------
module Data.Boolean.CUDD
    (
        BDD
    ,	module Data.Boolean

        -- * CUDD-specific functions
--          ,	gc
    ,	stats
    ,	reorder
    ,	dynamicReordering
    ,	varIndices
    ,	bddSize
    ,	printInfo
    ) where

-------------------------------------------------------------------
-- Dependencies.
-------------------------------------------------------------------

#include "cudd_im.h"

import Control.DeepSeq  ( NFData )
import Control.Monad	( foldM, liftM, mapAndUnzipM, zipWithM_ )

import Data.IORef	( IORef, newIORef, readIORef, writeIORef )
import Data.List	( genericLength )
import Data.Maybe	( fromJust, isJust )

import Data.Map         ( Map )
import qualified Data.Map as Map

import Foreign		( ForeignPtr, withForeignPtr, newForeignPtr, newForeignPtr_, finalizerFree, touchForeignPtr
                        , Ptr, advancePtr, castPtr, peek, pokeElemOff, ptrToIntPtr
                        , mallocArray )
import Foreign.C

import GHC.ForeignPtr   ( newConcForeignPtr )

import System.IO	( Handle, hIsReadable, hIsWritable )
-- import System.Mem	( performGC )
import System.Posix.IO	( handleToFd )
import System.IO.Unsafe	( unsafePerformIO )

import Data.Boolean

-------------------------------------------------------------------
-- Extra FFI functions.
-------------------------------------------------------------------

-- | A C file handle.
{#pointer *FILE -> CFile#}

-- | Convert a Haskell Handle into a C FILE *.
-- - FIXME: throw exception on error.
-- - suggested by Simon Marlow.
handleToCFile :: Handle -> IO FILE
handleToCFile h =
    do r <- hIsReadable h
       w <- hIsWritable h
       modestr <- newCString $ mode r w
       fd <- handleToFd h
       {#call unsafe fdopen#} (cToNum fd) modestr
    where mode :: Bool -> Bool -> String
          mode False False = error "Handle not readable or writable!"
          mode False True  = "w"
          mode True  False = "r"
          mode True  True  = "r+" -- FIXME

-- | The Haskell Handle is unusable after calling 'handleToCFile',
-- so provide a neutered "fprintf"-style thing here.
printCFile :: FILE -> String -> IO ()
printCFile file str =
    withCString str ({#call unsafe fprintf_neutered#} file)

--  Close a C FILE *.
-- - FIXME: throw exception on error.
-- closeCFile :: FILE -> IO ()
-- closeCFile cfile = do {#call unsafe fclose#} cfile
--                    return ()

cToNum :: (Num i, Integral e) => e -> i
cToNum  = fromIntegral . toInteger
{-# INLINE cToNum #-}

cFromEnum :: (Integral i, Enum e) => e -> i
cFromEnum  = fromIntegral . fromEnum
{-# INLINE cFromEnum #-}

-------------------------------------------------------------------
-- Types.
-------------------------------------------------------------------

-- | The BDD manager, which we hide from clients.
{#pointer *DdManager as DDManager newtype#}

-- | The arguments to 'Cudd_Init'.
{#enum CuddSubtableSize {} #}
{#enum CuddCacheSize {} #}

-- | Variable reordering tree options.
{#enum CuddMTRParams {} #}

-- | The abstract type of BDDs.
{#pointer *DdNode as BDD foreign newtype#}

{-# INLINE withBDD #-}

-- BDDs are just pointers, so there's no work to do when we @deepseq@ them.
instance Control.DeepSeq.NFData BDD

-- Belt-and-suspenders equality checking and ordering.
instance Eq BDD where
    bdd0 == bdd1 = unsafePerformIO $
      withBDD bdd0 $ \bdd0p -> withBDD bdd1 $ \bdd1p ->
        return $ bdd0p == bdd1p

instance Ord BDD where
    bdd0 <= bdd1 = unsafePerformIO $
      withBDD bdd0 $ \bdd0p -> withBDD bdd1 $ \bdd1p ->
        return $ bdd0p == bdd1p || bdd0p < bdd1p

instance Show BDD where
-- Print out the BDD as a sum-of-products.
    showsPrec _ = sop

-------------------------------------------------------------------
-- Administrative functions.
-------------------------------------------------------------------

-- | The BDD manager. This needs to be invoked before any of the
-- following functions... which we arrange for automagically.
ddmanager :: DDManager
ddmanager = unsafePerformIO $
  {#call unsafe Cudd_Init as _cudd_Init #}
    0 0 (cFromEnum UNIQUE_SLOTS) (cFromEnum CACHE_SLOTS) 0
{-# NOINLINE ddmanager #-}

-- | A map to and from BDD variable numbers.
-- FIXME: the second one would be more efficiently Data.IntMap, but we don't use it often.
type VarMap = (Map String BDD, Map CUInt String)

-- | Tracks existing variables.
bdd_vars :: IORef VarMap
bdd_vars = unsafePerformIO $ newIORef (Map.empty, Map.empty)
{-# NOINLINE bdd_vars #-}

-- | Attaches a finalizer to a BDD object.
-- The call to "Cudd_Ref" happens in C to ensure atomicity.
-- Returning objects with initial refcount 0 is a bad design decision.
addBDDfinalizer :: Ptr BDD -> IO BDD
addBDDfinalizer bddp = liftM BDD $ newConcForeignPtr bddp bddf
    where bddf = {#call unsafe Cudd_RecursiveDeref as _cudd_RecursiveDeref#} ddmanager bddp
                   >> return ()
{-# INLINE addBDDfinalizer #-}

-- | Attaches a null finalizer to a 'BDD' object.
-- Used for variables and constants.
addBDDnullFinalizer :: Ptr BDD -> IO BDD
addBDDnullFinalizer bddp = fmap BDD $ newForeignPtr_ bddp
{-# INLINE addBDDnullFinalizer #-}

-- Simulate an "apply" function to simplify atomic refcount incrementing.
{#enum CuddBinOp {} #}

-------------------------------------------------------------------
-- Logical operations.
-------------------------------------------------------------------

-- | Allocate a variable.
newVar :: IO (Ptr BDD) -> String -> IO (Maybe CUInt, BDD)
newVar allocVar label =
      do (toBDD, fromBDD) <- readIORef bdd_vars
         case label `Map.lookup` toBDD of
           Just bdd -> return (Nothing, bdd)
           Nothing ->
             do bddp <- allocVar
                vid <- fmap cToNum $
                         {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#} bddp
                --putStrLn $ label ++ " -> " ++ (show (vid, bdd))
                bdd <- addBDDnullFinalizer bddp
                writeIORef bdd_vars ( Map.insert label bdd   toBDD
                                    , Map.insert vid   label fromBDD )
                return (Just vid, bdd)

instance BooleanVariable BDD where
    bvar label = unsafePerformIO $
      fmap snd $ newVar ({#call unsafe Cudd_bddNewVar as _cudd_bddNewVar#} ddmanager)  label

    bvars labels =
      case labels of
        [] -> []
        ls -> unsafePerformIO $
          do (vids, vars) <- mapAndUnzipM (newVar ({#call unsafe Cudd_bddNewVar as _cudd_bddNewVar#} ddmanager)) ls
             if all isJust vids
               then groupVars (fromJust (head vids)) (genericLength vids)
               else putStrLn $ "hBDD-CUDD warning: not grouping variables " ++ show ls
             return vars
      where
        -- FIXME why one or the other?
        groupVars vid len = makeTreeNode ddmanager vid len (cFromEnum CUDD_MTR_DEFAULT) >> return ()
        -- groupVars vid len = makeTreeNode ddmanager vid len (cFromEnum CUDD_MTR_FIXED) >> return ()
        makeTreeNode = {#call unsafe Cudd_MakeTreeNode as _cudd_MakeTreeNode#}

    unbvar bdd = unsafePerformIO $ bdd `seq`
                   withBDD bdd $ \bddp ->
                     do vid <- {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#} bddp
                        (_, fromBDD) <- readIORef bdd_vars
                        return $ case cToNum vid `Map.lookup` fromBDD of
                                   Nothing -> "(VID: " ++ show vid ++ ")"
                                   Just v  -> v

instance Boolean BDD where
    true  = cudd_constant {#call unsafe Cudd_ReadOne as _cudd_ReadOne#}
    false = cudd_constant {#call unsafe Cudd_ReadLogicZero as _cudd_ReadLogicZero#}

    (/\)  = bddBinOp AND
    neg x = unsafePerformIO $ withBDD x $ \xp ->
      {#call unsafe cudd_bddNot#} xp >>= addBDDfinalizer

    nand  = bddBinOp NAND
    (\/)  = bddBinOp OR
    nor   = bddBinOp NOR
    xor   = bddBinOp XOR
    -- (-->) = FIXME ???
    (<->) = bddBinOp XNOR

    ifthenelse i t e = unsafePerformIO $
      withBDD i $ \ip -> withBDD t $ \tp -> withBDD e $ \ep ->
        {#call unsafe cudd_bddIte#} ddmanager ip tp ep
          >>= addBDDfinalizer

cudd_constant :: (DDManager -> IO (Ptr BDD)) -> BDD
cudd_constant f = unsafePerformIO $ f ddmanager >>= addBDDnullFinalizer
{-# INLINE cudd_constant #-}

bddBinOp :: CuddBinOp -> BDD -> BDD -> BDD
bddBinOp op x y = unsafePerformIO $
  withBDD x $ \xp -> withBDD y $ \yp ->
    {#call unsafe cudd_BinOp#} ddmanager (cFromEnum op) xp yp >>= addBDDfinalizer
{-# INLINE bddBinOp #-}

instance QBF BDD where
    data Group BDD = MkGroup !BDD

    mkGroup = MkGroup . conjoin -- FIXME maybe use Cudd_bddComputeCube
    exists = cudd_exists
    forall = cudd_forall
    rel_product = cudd_rel_product

instance Substitution BDD where
  -- Cache substitution arrays.
  data Subst BDD = MkSubst {-# UNPACK #-} !Int
                           {-# UNPACK #-} !(ForeignPtr (Ptr BDD))
                           {-# UNPACK #-} !(ForeignPtr (Ptr BDD))

  mkSubst = cudd_mkSubst
  rename = cudd_rename
  substitute = error "CUDD substitute" -- cudd_substitute

cudd_mkSubst :: [(BDD, BDD)] -> Subst BDD
cudd_mkSubst subst = unsafePerformIO $
    do arrayx  <- mallocArray len
       arrayx' <- mallocArray len
       zipWithM_ (pokeSubst arrayx arrayx') subst [0..]
       fpx  <- newForeignPtr finalizerFree arrayx
       fpx' <- newForeignPtr finalizerFree arrayx'
       return (MkSubst len fpx fpx')
  where
    pokeSubst :: Ptr (Ptr BDD) -> Ptr (Ptr BDD) -> (BDD, BDD) -> Int -> IO ()
    pokeSubst arrayx arrayx' (v, v') i =
      do withBDD v  $ pokeElemOff arrayx  i
         withBDD v' $ pokeElemOff arrayx' i

    len = length subst
{-# INLINE cudd_mkSubst #-}

instance BDDOps BDD where
    get_bdd_ptr bdd = unsafePerformIO $ withBDD bdd (return . ptrToIntPtr)
    bif bdd = unsafePerformIO $
      do vid <- withBDD bdd {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#}
         {#call unsafe Cudd_bddIthVar as _cudd_bddIthVar#} ddmanager (cToNum vid)
           >>= addBDDnullFinalizer
    bthen bdd = unsafePerformIO $
      withBDD bdd $ \bddp ->
        {#call unsafe cudd_bddT#} bddp >>= addBDDfinalizer
    belse bdd = unsafePerformIO $
      withBDD bdd $ \bddp ->
        {#call unsafe cudd_bddE#} bddp >>= addBDDfinalizer

    reduce = cudd_reduce
    satisfy = cudd_satisfy
    support = cudd_support

-------------------------------------------------------------------
-- Implementations.
-------------------------------------------------------------------

withGroup :: Group BDD -> (Ptr BDD -> IO a) -> IO a
withGroup (MkGroup g) = withBDD g

cudd_exists :: Group BDD -> BDD -> BDD
cudd_exists group bdd = unsafePerformIO $ bdd `seq`
  withGroup group $ \groupp -> withBDD bdd $ \bddp ->
    {#call unsafe cudd_bddExistAbstract#} ddmanager bddp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_exists #-}

cudd_forall :: Group BDD -> BDD -> BDD
cudd_forall group bdd = unsafePerformIO $ bdd `seq`
  withGroup group $ \groupp -> withBDD bdd $ \bddp ->
    {#call unsafe cudd_bddUnivAbstract#} ddmanager bddp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_forall #-}

cudd_rel_product :: Group BDD -> BDD -> BDD -> BDD
cudd_rel_product group f g = unsafePerformIO $
  withGroup group $ \groupp -> withBDD f $ \fp -> withBDD g $ \gp ->
    {#call unsafe cudd_bddAndAbstract#} ddmanager fp gp groupp
      >>= addBDDfinalizer
{-# INLINE cudd_rel_product #-}

-- | This function swaps variables, as it uses
-- @cudd_bddSwapVariables@. This is not exactly a "rename" behaviour.
cudd_rename :: Subst BDD -> BDD -> BDD
cudd_rename s@(MkSubst len fpx fpx') f = unsafePerformIO $ s `seq` f `seq`
  withBDD f $ \fp ->
    withForeignPtr fpx $ \arrayx -> withForeignPtr fpx' $ \arrayx' ->
      {#call unsafe cudd_bddSwapVariables#} ddmanager fp arrayx arrayx' (fromIntegral len)
        >>= addBDDfinalizer
{-# INLINE cudd_rename #-}

-- FIXME verify implementation.
cudd_reduce :: BDD -> BDD -> BDD
cudd_reduce f g = unsafePerformIO $
  withBDD f $ \fp -> withBDD g $ \gp ->
    {#call unsafe cudd_bddLICompaction#} ddmanager fp gp
      >>= addBDDfinalizer

-- FIXME verify implementation.
cudd_satisfy :: BDD -> BDD
cudd_satisfy bdd = unsafePerformIO $
  withBDD bdd ({#call unsafe cudd_satone#} ddmanager) >>= addBDDfinalizer

cudd_support :: BDD -> [BDD]
cudd_support bdd = unsafePerformIO $
  do varBitArray <- withBDD bdd ({#call unsafe Cudd_SupportIndex as _cudd_SupportIndex#} ddmanager)
     touchForeignPtr (case bdd of BDD fp -> fp)
     -- The array is as long as the number of variables allocated.
     (_toBDD, fromBDD) <- readIORef bdd_vars
     bdds <- foldM (toBDDs varBitArray) [] (Map.toList fromBDD)
     {#call unsafe cFree#} (castPtr varBitArray)
     return bdds
  where
    toBDDs varBitArray bdds (vid, var) =
      do varInSupport <- peek (advancePtr varBitArray (cToNum vid))
         return $ if varInSupport /= 0
                    then bvar var : bdds
                    else bdds

-------------------------------------------------------------------
-- Operations specific to this BDD binding.
-------------------------------------------------------------------

-- | Dump usage statistics to the given 'Handle'.
stats :: Handle -> IO ()
stats handle =
    do cfile <- handleToCFile handle
       printCFile cfile "CUDD stats\n"
       _ <- {#call unsafe Cudd_PrintInfo as _cudd_PrintInfo#} ddmanager cfile
       printCFile cfile "\nVariable groups\n"
       {#call unsafe cudd_printVarGroups#} ddmanager

----------------------------------------
-- Variable Reordering.
----------------------------------------

{#enum Cudd_ReorderingType as CUDDReorderingMethod {} deriving (Eq, Ord, Show)#}

decodeROM :: ReorderingMethod -> CInt
decodeROM rom = cFromEnum $ case rom of
  ReorderNone -> CUDD_REORDER_NONE
  ReorderSift -> CUDD_REORDER_SIFT
  ReorderSiftSym -> CUDD_REORDER_SYMM_SIFT
  ReorderStableWindow3 -> CUDD_REORDER_WINDOW3

-- | Reorder the variables now.
reorder :: ReorderingMethod -> IO ()
reorder rom = {#call unsafe Cudd_ReduceHeap as _cudd_ReduceHeap#} ddmanager (decodeROM rom) 1 >> return ()

-- | Set the dynamic variable ordering heuristic.
dynamicReordering :: ReorderingMethod -> IO ()
dynamicReordering rom = {#call unsafe Cudd_AutodynEnable as _cudd_AutodynEnable#} ddmanager (decodeROM rom) >> return ()

----------------------------------------
-- | Returns the relationship between BDD indices, BDD ids and
-- variable names. This may be useful for discovering what the
-- dynamic reordering is doing.
--
-- The intention of this function is to return
--   @[(position in variable order, immutable variable id, label)]@
-- so that the variable order can be saved.
----------------------------------------

varIndices :: IO [(CInt, CInt, String)]
varIndices = do (toBDD, _fromBDD) <- readIORef bdd_vars
                mapM procVar $ Map.toList toBDD
    where -- procVar :: (String, BDD) -> IO (CUInt, CUInt, String)
          procVar (var, bdd) =
              do -- FIXME CUDD inconsistently uses signed and unsigned ints.
                 vid <- fmap cToNum $
                          withBDD bdd {#call unsafe Cudd_NodeReadIndex as _cudd_NodeReadIndex#}
                 level <- {#call unsafe Cudd_ReadPerm as _cudd_ReadPerm#} ddmanager vid
                 return (level, vid, var)
{-# NOINLINE varIndices #-}

----------------------------------------
-- BDD Statistics.
----------------------------------------

-- | Determine the size of a BDD.
bddSize :: BDD -> Int
bddSize bdd = unsafePerformIO $
  do size <- withBDD bdd ({#call unsafe Cudd_DagSize as _cudd_DagSize#})
     return $ cToNum size

printInfo :: IO ()
printInfo = {#call unsafe print_stats#} ddmanager


commsec-test: fixup interopTest to build against cipher-aes-0.2.3

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE  EmptyDataDecls #-}

{- interopTest.hs
 - This interoperability test is not ment to replace KATs and other
 - validation techniques that should be applied prior to fielding of the
 - associated AES-GCM code and 'commsec' encapsulation.
 -
 - However, this test routine does serve to demonstrate interoperability
 - with a completely separate implementation of AES that, in at least one
 - version, was shown to be correct against a large set of KATs.
 -
 - Function:                    Test:                   Notes:
 - securePkg_init               correct keyExpand       implicit
 - securePkg_zero               nada
 - securePkg_enc_in_place       correctness             interop
 - securePkg_dec                correctness             interop
 - securePkg_enc                correctness             interop
 -}
module Main where

import Control.Monad
import Data.Word
import Foreign.Ptr
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Storable
import Foreign.Marshal (allocaBytes)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Marshal.Alloc (mallocBytes)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Monadic
import Test.QuickCheck

import qualified Data.ByteString as B
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Unsafe as BC

import Crypto.Cipher.AES                -- Vincent's GCM routine
import Crypto.Cipher.Types
import Control.Concurrent.MVar
import Data.Serialize

import Debug.Trace

#include "aeslib/commsec.h"

----------------------------------------------------------------------
-- Foreign import the securePkg_* calls which should be used
-- by HACMS/SMACCMS

-- If this were production code you would want to track the SecureContext
-- being zeroed and throw a run-time error when using a zeroed Ctx.
newtype SecureContext = SCtx (ForeignPtr Ctx)

data Ctx

newtype BaseId = BaseId Word32
                 deriving (Eq, Ord, Show)

instance Storable BaseId where
    sizeOf ~(BaseId x) = sizeOf x
    peek = fmap BaseId . peek . castPtr
    poke ptr (BaseId x) = poke (castPtr ptr) x
    alignment ~(BaseId x) = alignment x

instance Arbitrary BaseId where
    arbitrary = (BaseId . (`rem` 16) . abs) `fmap` arbitrary

instance Arbitrary ByteString where
    arbitrary = do
        len0 <- arbitrary
        let len = abs len0 `rem` 128
        B.pack `fmap` replicateM len arbitrary

instance Serialize BaseId where
    get = BaseId `fmap` getWord32be
    put (BaseId x) = putWord32be x

foreign import ccall "securePkg_size_of_message"
    c_securePkg_size_of_message :: CInt -> CInt

secPkg_size_of_message :: Int -> Int
secPkg_size_of_message =
  fromIntegral . c_securePkg_size_of_message . fromIntegral

foreign import ccall "securePkg_size_of_package"
    c_securePkg_size_of_package :: CInt -> CInt
secPkg_size_of_package :: Int -> Int
secPkg_size_of_package =
  fromIntegral . c_securePkg_size_of_package . fromIntegral

foreign import ccall "securePkg_init"
    c_securePkg_init :: Ptr Ctx -> BaseId ->
                        Word32 -> Ptr Word8 -> -- Decrypt Key
                        Word32 -> Ptr Word8 -> -- Encrypt Key
                        IO Word32

secPkgInit :: BaseId -> Word32 -> ByteString -> Word32 -> ByteString
           -> IO SecureContext
secPkgInit bid dsalt dkey esalt ekey = do
    -- FIXME sizeof is giving bad values!
    fptr <- mallocForeignPtrBytes ({#sizeof commsec_ctx#} + 10000)
    b <- withForeignPtr fptr $ \ctx ->
          BC.unsafeUseAsCString dkey $ \ptrD ->
           BC.unsafeUseAsCString ekey $ \ptrE ->
            c_securePkg_init ctx bid dsalt (castPtr ptrD) esalt (castPtr ptrE)
    if b == 0 then return (SCtx fptr) else error "secPkgInit failed"

foreign import ccall "securePkg_zero"
    c_securePkg_zero :: Ptr Ctx -> IO ()

secPkgZero :: SecureContext -> IO ()
secPkgZero (SCtx ptr) = withForeignPtr ptr c_securePkg_zero

foreign import ccall "securePkg_enc_in_place"
    c_securePkg_enc_in_place :: Ptr Ctx ->   -- Cipher ctx
                                Ptr Word8 -> -- Buffer
                                Word32 ->    -- Msg start index
                                Word32 ->    -- Msg Length
                                IO Word32

-- Encrypting data in place requires room for the header (8 bytes)
-- and room for the trailing tag (8 bytes)
secPkgEncInPlace :: SecureContext -> ByteString -> IO (Maybe ByteString)
secPkgEncInPlace (SCtx c) pt = do
    let pkgSz = secPkg_size_of_package . B.length $ pt
    pkg <- BI.mallocByteString pkgSz
    b <- withForeignPtr pkg $ \ptrCT ->
          BC.unsafeUseAsCStringLen pt $ \(ptrPT,lenPT) -> do
            copyBytes (castPtr ptrCT `plusPtr` 8) ptrPT lenPT
            withForeignPtr c $ \ctxPtr ->
             c_securePkg_enc_in_place
               ctxPtr
               (castPtr ptrCT)
               8
               (fromIntegral lenPT)
    let b = 0
    return (if b == 0 then Just (BI.fromForeignPtr pkg 0 pkgSz)
                      else Nothing)

foreign import ccall "securePkg_enc"
    c_securePkg_enc :: Ptr Ctx ->   -- Cipher ctx
                       Ptr Word8 -> -- PT Buffer (in) CT Buffer (out)
                       Word32    -> -- PT Len
                       Ptr Word8 -> -- Header Buffer (8 bytes, min)
                       Ptr Word8 -> -- Tag buffer (8 bytes, min)
                       IO Word32

-- Encrypting data in place requires room for the header (8 bytes)
-- and room for the trailing tag (8 bytes)
secPkgEnc :: SecureContext -> ByteString
          -> IO (Maybe (ByteString,ByteString,ByteString))
secPkgEnc (SCtx c) pt = do
    let tagLen = 8 -- C and IVORY could should use the C #define
        hdrLen = 8
        lenCT  = B.length pt
    ct  <- BI.mallocByteString lenCT
    hdr <- BI.mallocByteString hdrLen
    tag <- BI.mallocByteString tagLen
    b <- withForeignPtr ct $ \ptrCT ->
          BC.unsafeUseAsCStringLen pt $ \(ptrPT,lenPT) ->
           withForeignPtr tag $ \ptrTag ->
            withForeignPtr hdr $ \ptrHdr ->
             withForeignPtr c $ \ctxPtr -> do
               copyBytes (castPtr ptrCT) (castPtr ptrPT) lenPT
               c_securePkg_enc ctxPtr ptrCT (fromIntegral lenPT) ptrHdr ptrTag
    let tagBS = BI.fromForeignPtr tag 0 tagLen
        hdrBS = BI.fromForeignPtr hdr 0 hdrLen
        ctBS  = BI.fromForeignPtr ct 0 lenCT
    return (if b == 0 then Just (hdrBS,ctBS,tagBS) else Nothing)

foreign import ccall "securePkg_dec"
    c_securePkg_dec :: Ptr Ctx   ->
                       Ptr Word8 ->
                       Word32    ->
                       IO Word32

secPkgDec :: SecureContext -> ByteString -> IO (Maybe ByteString)
secPkgDec (SCtx c) ct =
  BC.unsafeUseAsCStringLen ct $ \(ptrPKG,lenPKG) -> do
    let lenW32 = fromIntegral lenPKG
        lenPT  = secPkg_size_of_message lenPKG
    fpPT <- BI.mallocByteString lenPT
    b <- allocaBytes lenPKG $ \tmpPtr -> do
         copyBytes tmpPtr ptrPKG lenPKG
         withForeignPtr c $ \ctxPtr -> do
         res <- c_securePkg_dec ctxPtr (castPtr tmpPtr) lenW32
         withForeignPtr fpPT $ \ptrPT -> do
           copyBytes ptrPT (tmpPtr `plusPtr` 8) lenPT
         return res
    return (if b == 0 then Just (BI.fromForeignPtr fpPT 0 lenPT)
                      else Nothing)

----------------------------------------------------------------------
-- Haskell re-implementation of the SecurePkg format.  This code uses
-- a completely separate C AES routine.  This is NOT a reference
-- in that it is easy to read, it is simply a re-implementation which
-- serves as a light sanity check on the functionallity of the prior
-- implementation.

-- Key, salt, station ID, counter
type OutContext = (AES,Word32,BaseId,Word32)
type InContext = (AES,Word32,[(BaseId,Word32)])
data SecureContext_HS = SC { inbound    :: InContext
                           , outbound   :: OutContext }
type ReferenceContext = MVar SecureContext_HS

refPkg :: OutContext -> ByteString
       -> (OutContext, Maybe (ByteString,ByteString,ByteString))
refPkg ctx@(key,salt,bid,ctr) pt
    | ctr == maxBound = (ctx, Nothing)
    | otherwise       =
  let iv  = runPut (putWord32be salt >> put bid >> putWord32be ctr)
      new = (key,salt,bid,ctr+1)
      aad = B.empty
      (ct, AuthTag tag) = encryptGCM key iv aad pt
      tagLen   = 8
      header   = runPut (put bid >> putWord32be ctr)
  in (new, Just (header, ct, B.take tagLen tag))

refDec :: InContext -> ByteString -> (InContext, Maybe ByteString)
refDec old@(key,salt,bidList) pkg =
    let aad = B.empty
        -- Get the iv out of the send package.
        Right (bid,newCtr,ct,tag) =
           runGet (do bid'    <- get
                      newCtr' <- getWord32be
                      pt'     <- getByteString . (subtract 8) =<< remaining
                      tag'    <- getByteString =<< remaining
                      return (bid',newCtr',pt',tag')
                  ) pkg
        -- Serialize the salt I have with the sent baseID and updated counter.
        iv  = runPut (putWord32be salt >> put bid >> putWord32be newCtr)
        -- Replace the BaseID/Counter pair in the list with the updated pair.
        newBids = (bid,newCtr) : filter ((/= bid) . fst) bidList
        new     = (key, salt, newBids)
        -- Decrypt the message.  I get the sent message and auth. tag, if all
        -- goes well.
        (pt, AuthTag decTag) = decryptGCM key iv aad ct
    in case lookup bid bidList of
        Nothing -> (old, Nothing)
        Just cnt
          -- Counter is too high or old.
          | cnt == maxBound || cnt >= newCtr    -> (old,Nothing)
          -- Return my updated inContext with the decrypted message.
          | B.take (B.length tag) decTag == tag -> (new,Just pt)
          | otherwise                           -> (old,Nothing)

refInitOutContext :: ByteString -> Word32 -> BaseId -> OutContext
refInitOutContext key salt bid = (initAES key, salt, bid, 1)

refInitInContext :: ByteString -> Word32 -> InContext
refInitInContext key salt
  = (initAES key, salt, zip (map BaseId [0..]) (replicate 16 0))

secPkgInit_HS :: BaseId -> Word32 -> ByteString -> Word32 -> ByteString
              -> IO ReferenceContext
secPkgInit_HS b dsalt dkey esalt ekey
  = newMVar (SC (refInitInContext dkey dsalt)
                (refInitOutContext ekey esalt b))

secPkgEncInPlace_HS :: ReferenceContext -> ByteString -> IO (Maybe ByteString)
secPkgEncInPlace_HS rc pt = do
    r <- secPkgEnc_HS rc pt
    case r of
        Just (a,b,c) -> return (Just (B.concat [a,b,c]))
        Nothing -> return Nothing

----------------------------------------------------------------------
-- Lifting the reference implementation to match the API
secPkgEnc_HS :: ReferenceContext -> ByteString
             -> IO (Maybe (ByteString,ByteString,ByteString))
secPkgEnc_HS rc pt =
  modifyMVar rc $ \(SC _in outbound) ->
    let (newOut,res) = refPkg outbound pt in do
    return (SC _in newOut,res)

secPkgDec_HS :: ReferenceContext -> ByteString -> IO (Maybe ByteString)
secPkgDec_HS rc pkg =
  modifyMVar rc $ \(SC inbound _out) ->
    let (newIn, res) = refDec inbound pkg in return (SC newIn _out, res)

-----------------------------------------------------------------------
-- Quickcheck tests showing compatability of the two implementations.

-- FIXME more tests, use monadic guard to not test trivial pt==null case.

prop_encEq :: ByteString -> ByteString -> ByteString -> Word32 -> Word32
           -> Word32 -> Property
prop_encEq dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit    (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    rPKG <- run $ secPkgEnc_HS rctx pt
    sPKG <- run $ secPkgEnc sctx pt
    assert $ sPKG == rPKG

prop_decCompat :: ByteString -> ByteString -> ByteString -> Word32
               -> Word32 -> Word32 -> Property
prop_decCompat dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    Just rPKG <- run $ secPkgEncInPlace_HS rctx pt
    Just sPKG <- run $ secPkgEncInPlace sctx pt
    dec1 <- run $ secPkgDec_HS rctx sPKG
    dec2 <- run $ secPkgDec sctx rPKG
    assert $ sPKG == rPKG && dec1 == dec2

prop_decFail :: ByteString -> ByteString -> ByteString -> Word32 -> Word32
             -> Word32 -> Property
prop_decFail dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    Just sPKG0 <- run $ secPkgEncInPlace sctx pt
    i <- pick $ fmap ((`rem` B.length sPKG0) . abs) arbitrary
    let sPKG = B.pack . incIth i . B.unpack $ sPKG0
    dec1 <- run $ secPkgDec_HS rctx sPKG
    dec2 <- run $ secPkgDec sctx sPKG
    assert $ dec1 == dec2 && dec1 == Nothing
 where
 incIth i ls = take i ls ++ [(ls !! i) + 1] ++ drop (i+1) ls

data Test = forall a. Testable a => T a String

tests = [T prop_encEq "prop_encEq"
        ,T prop_decFail "prop_decFail"
        ,T prop_decCompat "prop_decCompat"
        ]

runTest :: Test -> IO ()
runTest (T a s) = putStrLn ("Testing: " ++ s) >> quickCheck a

runTests :: [Test] -> IO ()
runTests = mapM_ runTest

main :: IO ()
main = runTests tests
{-
main = do
    putStrLn "Notice the commsec_ctx structure is ~9k, depending on configuration.  However, c2hs believes it is smaller (see the below number) so we manually ad 10K to make this test work.   This is a definite FIXME."
    print {#sizeof commsec_ctx#}
    let key = B.replicate 16 0
    ctx <- secPkgInit (BaseId 0) 1 key 1 key
    ctx2 <- secPkgInit_HS (BaseId 0) 1 key 1 key
    Just pkg1 <- secPkgEncInPlace_HS ctx2 (B.pack [1..100])
    Just pkg2 <- secPkgEncInPlace ctx (B.pack [1..100])
    Just pkg3 <- secPkgEncInPlace ctx (B.pack [1..100])
    print pkg1
    print pkg2
    print pkg3
    secPkgDec_HS ctx2 pkg1 >>= print
    secPkgDec ctx pkg2 >>= print
    secPkgDec ctx pkg3 >>= print
    return ()
-}

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE  EmptyDataDecls #-}

{- interopTest.hs
 - This interoperability test is not ment to replace KATs and other
 - validation techniques that should be applied prior to fielding of the
 - associated AES-GCM code and 'commsec' encapsulation.
 -
 - However, this test routine does serve to demonstrate interoperability
 - with a completely separate implementation of AES that, in at least one
 - version, was shown to be correct against a large set of KATs.
 -
 - Function:                    Test:                   Notes:
 - securePkg_init               correct keyExpand       implicit
 - securePkg_zero               nada
 - securePkg_enc_in_place       correctness             interop
 - securePkg_dec                correctness             interop
 - securePkg_enc                correctness             interop
 -}
module Main where

import Control.Monad
import Data.Word
import Foreign.Ptr
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Storable
import Foreign.Marshal (allocaBytes)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Marshal.Alloc (mallocBytes)
import Test.QuickCheck.Arbitrary
import Test.QuickCheck.Monadic
import Test.QuickCheck

import qualified Data.ByteString as B
import Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Unsafe as BC

import Crypto.Cipher.AES                -- Vincent's GCM routine
import Control.Concurrent.MVar
import Data.Serialize

import Debug.Trace

#include "aeslib/commsec.h"

----------------------------------------------------------------------
-- Foreign import the securePkg_* calls which should be used
-- by HACMS/SMACCMS

-- If this were production code you would want to track the SecureContext
-- being zeroed and throw a run-time error when using a zeroed Ctx.
newtype SecureContext = SCtx (ForeignPtr Ctx)

data Ctx

newtype BaseId = BaseId Word32
                 deriving (Eq, Ord, Show)

instance Storable BaseId where
    sizeOf ~(BaseId x) = sizeOf x
    peek = fmap BaseId . peek . castPtr
    poke ptr (BaseId x) = poke (castPtr ptr) x
    alignment ~(BaseId x) = alignment x

instance Arbitrary BaseId where
    arbitrary = (BaseId . (`rem` 16) . abs) `fmap` arbitrary

instance Arbitrary ByteString where
    arbitrary = do
        len0 <- arbitrary
        let len = abs len0 `rem` 128
        B.pack `fmap` replicateM len arbitrary

instance Serialize BaseId where
    get = BaseId `fmap` getWord32be
    put (BaseId x) = putWord32be x

foreign import ccall "securePkg_size_of_message"
    c_securePkg_size_of_message :: CInt -> CInt

secPkg_size_of_message :: Int -> Int
secPkg_size_of_message =
  fromIntegral . c_securePkg_size_of_message . fromIntegral

foreign import ccall "securePkg_size_of_package"
    c_securePkg_size_of_package :: CInt -> CInt
secPkg_size_of_package :: Int -> Int
secPkg_size_of_package =
  fromIntegral . c_securePkg_size_of_package . fromIntegral

foreign import ccall "securePkg_init"
    c_securePkg_init :: Ptr Ctx -> BaseId ->
                        Word32 -> Ptr Word8 -> -- Decrypt Key
                        Word32 -> Ptr Word8 -> -- Encrypt Key
                        IO Word32

secPkgInit :: BaseId -> Word32 -> ByteString -> Word32 -> ByteString
           -> IO SecureContext
secPkgInit bid dsalt dkey esalt ekey = do
    -- FIXME sizeof is giving bad values!
    fptr <- mallocForeignPtrBytes ({#sizeof commsec_ctx#} + 10000)
    b <- withForeignPtr fptr $ \ctx ->
          BC.unsafeUseAsCString dkey $ \ptrD ->
           BC.unsafeUseAsCString ekey $ \ptrE ->
            c_securePkg_init ctx bid dsalt (castPtr ptrD) esalt (castPtr ptrE)
    if b == 0 then return (SCtx fptr) else error "secPkgInit failed"

foreign import ccall "securePkg_zero"
    c_securePkg_zero :: Ptr Ctx -> IO ()

secPkgZero :: SecureContext -> IO ()
secPkgZero (SCtx ptr) = withForeignPtr ptr c_securePkg_zero

foreign import ccall "securePkg_enc_in_place"
    c_securePkg_enc_in_place :: Ptr Ctx ->   -- Cipher ctx
                                Ptr Word8 -> -- Buffer
                                Word32 ->    -- Msg start index
                                Word32 ->    -- Msg Length
                                IO Word32

-- Encrypting data in place requires room for the header (8 bytes)
-- and room for the trailing tag (8 bytes)
secPkgEncInPlace :: SecureContext -> ByteString -> IO (Maybe ByteString)
secPkgEncInPlace (SCtx c) pt = do
    let pkgSz = secPkg_size_of_package . B.length $ pt
    pkg <- BI.mallocByteString pkgSz
    b <- withForeignPtr pkg $ \ptrCT ->
          BC.unsafeUseAsCStringLen pt $ \(ptrPT,lenPT) -> do
            copyBytes (castPtr ptrCT `plusPtr` 8) ptrPT lenPT
            withForeignPtr c $ \ctxPtr ->
             c_securePkg_enc_in_place
               ctxPtr
               (castPtr ptrCT)
               8
               (fromIntegral lenPT)
    let b = 0
    return (if b == 0 then Just (BI.fromForeignPtr pkg 0 pkgSz)
                      else Nothing)

foreign import ccall "securePkg_enc"
    c_securePkg_enc :: Ptr Ctx ->   -- Cipher ctx
                       Ptr Word8 -> -- PT Buffer (in) CT Buffer (out)
                       Word32    -> -- PT Len
                       Ptr Word8 -> -- Header Buffer (8 bytes, min)
                       Ptr Word8 -> -- Tag buffer (8 bytes, min)
                       IO Word32

-- Encrypting data in place requires room for the header (8 bytes)
-- and room for the trailing tag (8 bytes)
secPkgEnc :: SecureContext -> ByteString
          -> IO (Maybe (ByteString,ByteString,ByteString))
secPkgEnc (SCtx c) pt = do
    let tagLen = 8 -- C and IVORY could should use the C #define
        hdrLen = 8
        lenCT  = B.length pt
    ct  <- BI.mallocByteString lenCT
    hdr <- BI.mallocByteString hdrLen
    tag <- BI.mallocByteString tagLen
    b <- withForeignPtr ct $ \ptrCT ->
          BC.unsafeUseAsCStringLen pt $ \(ptrPT,lenPT) ->
           withForeignPtr tag $ \ptrTag ->
            withForeignPtr hdr $ \ptrHdr ->
             withForeignPtr c $ \ctxPtr -> do
               copyBytes (castPtr ptrCT) (castPtr ptrPT) lenPT
               c_securePkg_enc ctxPtr ptrCT (fromIntegral lenPT) ptrHdr ptrTag
    let tagBS = BI.fromForeignPtr tag 0 tagLen
        hdrBS = BI.fromForeignPtr hdr 0 hdrLen
        ctBS  = BI.fromForeignPtr ct 0 lenCT
    return (if b == 0 then Just (hdrBS,ctBS,tagBS) else Nothing)

foreign import ccall "securePkg_dec"
    c_securePkg_dec :: Ptr Ctx   ->
                       Ptr Word8 ->
                       Word32    ->
                       IO Word32

secPkgDec :: SecureContext -> ByteString -> IO (Maybe ByteString)
secPkgDec (SCtx c) ct =
  BC.unsafeUseAsCStringLen ct $ \(ptrPKG,lenPKG) -> do
    let lenW32 = fromIntegral lenPKG
        lenPT  = secPkg_size_of_message lenPKG
    fpPT <- BI.mallocByteString lenPT
    b <- allocaBytes lenPKG $ \tmpPtr -> do
         copyBytes tmpPtr ptrPKG lenPKG
         withForeignPtr c $ \ctxPtr -> do
         res <- c_securePkg_dec ctxPtr (castPtr tmpPtr) lenW32
         withForeignPtr fpPT $ \ptrPT -> do
           copyBytes ptrPT (tmpPtr `plusPtr` 8) lenPT
         return res
    return (if b == 0 then Just (BI.fromForeignPtr fpPT 0 lenPT)
                      else Nothing)

----------------------------------------------------------------------
-- Haskell re-implementation of the SecurePkg format.  This code uses
-- a completely separate C AES routine.  This is NOT a reference
-- in that it is easy to read, it is simply a re-implementation which
-- serves as a light sanity check on the functionallity of the prior
-- implementation.

-- Key, salt, station ID, counter
type OutContext = (Key,Word32,BaseId,Word32)
type InContext = (Key,Word32,[(BaseId,Word32)])
data SecureContext_HS = SC { inbound    :: InContext
                           , outbound   :: OutContext }
type ReferenceContext = MVar SecureContext_HS

refPkg :: OutContext -> ByteString
       -> (OutContext, Maybe (ByteString,ByteString,ByteString))
refPkg ctx@(key,salt,bid,ctr) pt
    | ctr == maxBound = (ctx, Nothing)
    | otherwise       =
  let iv  = runPut (putWord32be salt >> put bid >> putWord32be ctr)
      new = (key,salt,bid,ctr+1)
      aad = B.empty
      (ct,tag) = encryptGCM key (IV iv) aad pt
      tagLen   = 8
      header   = runPut (put bid >> putWord32be ctr)
  in (new, Just (header, ct, B.take tagLen tag))

refDec :: InContext -> ByteString -> (InContext, Maybe ByteString)
refDec old@(key,salt,bidList) pkg =
    let iv  = runPut (putWord32be salt >> put bid >> putWord32be newCtr)
        aad = B.empty
        Right (bid,newCtr,ct,tag) =
           runGet (do bid    <- get
                      newCtr <- getWord32be
                      pt <- getByteString . (subtract 8) =<< remaining
                      tag <- getByteString =<< remaining
                      return (bid,newCtr,pt,tag)
                      ) pkg
        newBids = (bid,newCtr) : filter ((/= bid) . fst) bidList
        new     = (key,salt,newBids)
        (pt,decTag) = decryptGCM key (IV iv) aad ct
    in case lookup bid bidList of
        Nothing -> (old, Nothing)
        Just cnt | cnt == maxBound || cnt >= newCtr    -> (old,Nothing)
                 | B.take (B.length tag) decTag == tag -> (new,Just pt)
                 | otherwise                           -> (old,Nothing)

refInitOutContext :: ByteString -> Word32 -> BaseId -> OutContext
refInitOutContext key salt bid = (initKey key, salt, bid, 1)

refInitInContext :: ByteString -> Word32 -> InContext
refInitInContext key salt
  = (initKey key, salt, zip (map BaseId [0..]) (replicate 16 0))

secPkgInit_HS :: BaseId -> Word32 -> ByteString -> Word32 -> ByteString
              -> IO ReferenceContext
secPkgInit_HS b dsalt dkey esalt ekey
  = newMVar (SC (refInitInContext dkey dsalt)
                (refInitOutContext ekey esalt b))

secPkgEncInPlace_HS :: ReferenceContext -> ByteString -> IO (Maybe ByteString)
secPkgEncInPlace_HS rc pt = do
    r <- secPkgEnc_HS rc pt
    case r of
        Just (a,b,c) -> return (Just (B.concat [a,b,c]))
        Nothing -> return Nothing

----------------------------------------------------------------------
-- Lifting the reference implementation to match the API
secPkgEnc_HS :: ReferenceContext -> ByteString
             -> IO (Maybe (ByteString,ByteString,ByteString))
secPkgEnc_HS rc pt =
  modifyMVar rc $ \(SC _in outbound) ->
    let (newOut,res) = refPkg outbound pt in do
    return (SC _in newOut,res)

secPkgDec_HS :: ReferenceContext -> ByteString -> IO (Maybe ByteString)
secPkgDec_HS rc pkg =
  modifyMVar rc $ \(SC inbound _out) ->
    let (newIn, res) = refDec inbound pkg in return (SC newIn _out, res)

-----------------------------------------------------------------------
-- Quickcheck tests showing compatability of the two implementations.

-- FIXME more tests, use monadic guard to not test trivial pt==null case.

prop_encEq :: ByteString -> ByteString -> ByteString -> Word32 -> Word32
           -> Word32 -> Property
prop_encEq dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit    (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    rPKG <- run $ secPkgEnc_HS rctx pt
    sPKG <- run $ secPkgEnc sctx pt
    assert $ sPKG == rPKG

prop_decCompat :: ByteString -> ByteString -> ByteString -> Word32
               -> Word32 -> Word32 -> Property
prop_decCompat dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    Just rPKG <- run $ secPkgEncInPlace_HS rctx pt
    Just sPKG <- run $ secPkgEncInPlace sctx pt
    dec1 <- run $ secPkgDec_HS rctx sPKG
    dec2 <- run $ secPkgDec sctx rPKG
    assert $ sPKG == rPKG && dec1 == dec2

prop_decFail :: ByteString -> ByteString -> ByteString -> Word32 -> Word32
             -> Word32 -> Property
prop_decFail dKeySmall eKeySmall ptSmall dsalt esalt ctr
  = monadicIO $ do
    let dkey = B.take 16 (B.append dKeySmall (B.replicate 16 0))
    let ekey = B.take 16 (B.append eKeySmall (B.replicate 16 0))
    let pt   = B.take (max (B.length ptSmall) 1)
                      (B.append ptSmall (B.replicate 1 1))
    sctx <- run $ secPkgInit (BaseId 0) dsalt dkey esalt ekey
    rctx <- run $ secPkgInit_HS (BaseId 0) dsalt dkey esalt ekey
    Just sPKG0 <- run $ secPkgEncInPlace sctx pt
    i <- pick $ fmap ((`rem` B.length sPKG0) . abs) arbitrary
    let sPKG = B.pack . incIth i . B.unpack $ sPKG0
    dec1 <- run $ secPkgDec_HS rctx sPKG
    dec2 <- run $ secPkgDec sctx sPKG
    assert $ dec1 == dec2 && dec1 == Nothing
 where
 incIth i ls = take i ls ++ [(ls !! i) + 1] ++ drop (i+1) ls

data Test = forall a. Testable a => T a String

tests = [T prop_encEq "prop_encEq"
        ,T prop_decFail "prop_decFail"
        ,T prop_decCompat "prop_decCompat"
        ]

runTest :: Test -> IO ()
runTest (T a s) = putStrLn ("Testing: " ++ s) >> quickCheck a

runTests :: [Test] -> IO ()
runTests = mapM_ runTest

main :: IO ()
main = runTests tests
{-
main = do
    putStrLn "Notice the commsec_ctx structure is ~9k, depending on configuration.  However, c2hs believes it is smaller (see the below number) so we manually ad 10K to make this test work.   This is a definite FIXME."
    print {#sizeof commsec_ctx#}
    let key = B.replicate 16 0
    ctx <- secPkgInit (BaseId 0) 1 key 1 key
    ctx2 <- secPkgInit_HS (BaseId 0) 1 key 1 key
    Just pkg1 <- secPkgEncInPlace_HS ctx2 (B.pack [1..100])
    Just pkg2 <- secPkgEncInPlace ctx (B.pack [1..100])
    Just pkg3 <- secPkgEncInPlace ctx (B.pack [1..100])
    print pkg1
    print pkg2
    print pkg3
    secPkgDec_HS ctx2 pkg1 >>= print
    secPkgDec ctx pkg2 >>= print
    secPkgDec ctx pkg3 >>= print
    return ()
-}


storable instance for DevicePtr

storable instance for DevicePtr

Ignore-this: b5f768604b0ed1581e6f7f0408d1f761
- to be used when setting function parameters

darcs-hash:20091210065234-9241b-d5cbb224cd1b05dae8ef6c1d2c5b5183c5021427.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory management for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Marshal
  (
    DevicePtr, HostPtr, AllocFlag(..), withHostPtr,
    malloc, free, mallocHost, freeHost,
    peekArray, peekArrayAsync, pokeArray, pokeArrayAsync, memset,
    getDevicePtr
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Internal.C2HS
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Driver.Stream               (Stream, withStream)

-- System
import Foreign                           hiding (peekArray, pokeArray, malloc, free)
import Foreign.C
import Control.Monad                            (liftM)
import Unsafe.Coerce

#c
typedef enum CUmemhostalloc_option_enum {
    CU_MEMHOSTALLOC_OPTION_PORTABLE       = CU_MEMHOSTALLOC_PORTABLE,
    CU_MEMHOSTALLOC_OPTION_DEVICE_MAPPED  = CU_MEMHOSTALLOC_DEVICEMAP,
    CU_MEMHOSTALLOC_OPTION_WRITE_COMBINED = CU_MEMHOSTALLOC_WRITECOMBINED
} CUmemhostalloc_option;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to memory allocated on the device
--
newtype DevicePtr a = DevicePtr { useDevicePtr :: {# type CUdeviceptr #}}
  deriving (Show)

instance Storable (DevicePtr a) where
  sizeOf _      = sizeOf    (undefined :: {# type CUdeviceptr #})
  alignment _   = alignment (undefined :: {# type CUdeviceptr #})
  peek p        = DevicePtr `fmap` peek (castPtr p)
  poke p v      = poke (castPtr p) (useDevicePtr v)

-- |
-- A reference to memory on the host that is page-locked and directly-accessible
-- from the device. Since the memory can be accessed directly, it can be read or
-- written at a much higher bandwidth than pageable memory from the traditional
-- malloc.
--
-- The driver automatically accelerates calls to functions such as `memcpy'
-- which reference page-locked memory.
--
newtype HostPtr a = HostPtr (ForeignPtr a)

withHostPtr :: HostPtr a -> (Ptr a -> IO b) -> IO b
withHostPtr (HostPtr hptr) = withForeignPtr hptr

newHostPtr :: IO (HostPtr a)
newHostPtr = HostPtr `fmap` mallocForeignPtrBytes (sizeOf (undefined :: Ptr ()))

-- |
-- Options for host allocation
--
{# enum CUmemhostalloc_option as AllocFlag
    { underscoreToCase }
    with prefix="CU_MEMHOSTALLOC_OPTION" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Host Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the host which is page-locked and
-- directly accessible from the device. Note that since the amount of pageable
-- memory is thusly reduced, overall system performance may suffer. This is best
-- used sparingly to allocate staging areas for data exchange.
--
mallocHost :: [AllocFlag] -> Int -> IO (Either String (HostPtr a))
mallocHost flags bytes =
  newHostPtr >>= \hp -> withHostPtr hp $ \p ->
  cuMemHostAlloc p bytes flags >>= \rv ->
  return $ case nothingIfOk rv of
    Nothing  -> Right hp
    Just err -> Left err

{# fun unsafe cuMemHostAlloc
  { with'* `Ptr a'
  ,        `Int'
  , combineBitMasks `[AllocFlag]' } -> `Status' cToEnum #}
  where with' = with . castPtr

-- |
-- Free a section of page-locked host memory
--
freeHost :: HostPtr a -> IO (Maybe String)
freeHost hp = withHostPtr hp $ \p -> (nothingIfOk `fmap` cuMemFreeHost p)

{# fun unsafe cuMemFreeHost
  { castPtr `Ptr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Device Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the device, and return a reference to
-- it. The memory is suitably aligned for any type, and is not cleared.
--
malloc :: Int -> IO (Either String (DevicePtr a))
malloc bytes = resultIfOk `fmap` cuMemAlloc bytes

{# fun unsafe cuMemAlloc
  { alloca-  `DevicePtr a' dptr*
  ,          `Int'               } -> `Status' cToEnum #}
  where dptr = liftM DevicePtr . peekIntConv

-- |
-- Release a section of device memory
--
free :: DevicePtr a -> IO (Maybe String)
free dp = nothingIfOk `fmap` cuMemFree dp

{# fun unsafe cuMemFree
  { useDevicePtr `DevicePtr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Transfer
--------------------------------------------------------------------------------

-- |
-- Copy a number of elements from the device to host memory. This is a
-- synchronous operation
--
peekArray :: Storable a => Int -> DevicePtr a -> Ptr a -> IO (Maybe String)
peekArray n dptr hptr = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPeek x _ = nothingIfOk `fmap` cuMemcpyDtoH hptr dptr (n * sizeOf x)

{# fun unsafe cuMemcpyDtoH
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory from the device asynchronously, possibly associated with a
-- particular stream. The destination host memory must be page-locked (allocated
-- by mallocHost)
--
peekArrayAsync :: Storable a => Int -> DevicePtr a -> Ptr a -> Maybe Stream -> IO (Maybe String)
peekArrayAsync n dptr hptr mst = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPeek x _ =
      nothingIfOk `fmap` case mst of
        Nothing -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) nullPtr
        Just st -> (withStream st $ \s -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) s)

{# fun unsafe cuMemcpyDtoHAsync
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'
  , castPtr      `Ptr Stream'  } -> `Status' cToEnum #}

-- |
-- Copy a number of elements onto the device. This is a synchronous operation
--
pokeArray :: Storable a => Int -> Ptr a -> DevicePtr a -> IO (Maybe String)
pokeArray n hptr dptr = doPoke undefined dptr
  where
    doPoke :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPoke x _ = nothingIfOk `fmap` cuMemcpyHtoD dptr hptr (n * sizeOf x)

{# fun unsafe cuMemcpyHtoD
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory onto the device asynchronously, possibly associated with a
-- particular stream. The source host memory must be page-locked (allocated by
-- mallocHost)
--
pokeArrayAsync :: Storable a => Int -> Ptr a -> DevicePtr a -> Maybe Stream -> IO (Maybe String)
pokeArrayAsync n hptr dptr mst = dopoke undefined dptr
  where
    dopoke :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    dopoke x _ =
      nothingIfOk `fmap` case mst of
        Nothing -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) nullPtr
        Just st -> (withStream st $ \s -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) s)

{# fun unsafe cuMemcpyHtoDAsync
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'
  , castPtr      `Ptr Stream'  } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Utility
--------------------------------------------------------------------------------

-- |
-- Set a number of data elements to the specified value, which may be either 8-,
-- 16-, or 32-bits wide.
--
memset :: Storable a => DevicePtr a -> Int -> a -> IO (Maybe String)
memset dptr n val = case sizeOf val of
    1 -> nothingIfOk `fmap` cuMemsetD8  dptr val n
    2 -> nothingIfOk `fmap` cuMemsetD16 dptr val n
    4 -> nothingIfOk `fmap` cuMemsetD32 dptr val n
    _ -> return $ Just "can only memset 8-, 16-, and 32-bit values"

--
-- We use unsafe coerce below to reinterpret the bits of the value to memset as,
-- into the integer type required by the setting functions.
--
{# fun unsafe cuMemsetD8
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD16
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD32
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Return the device pointer associated with a mapped, pinned host buffer, which
-- was allocated with the DEVICE_MAPPED option.
--
-- Currently, no options are supported and this must be empty.
--
getDevicePtr :: [AllocFlag] -> HostPtr a -> IO (Either String (DevicePtr a))
getDevicePtr flags hptr = withHostPtr hptr $ \hp -> (resultIfOk `fmap` cuMemHostGetDevicePointer hp flags)

{# fun unsafe cuMemHostGetDevicePointer
  { alloca-         `DevicePtr a' dptr*
  , castPtr         `Ptr a'
  , combineBitMasks `[AllocFlag]'       } -> `Status' cToEnum #}
  where dptr = liftM DevicePtr . peekIntConv

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory management for low-level driver interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Marshal
  (
    DevicePtr, HostPtr, AllocFlag(..), withHostPtr,
    malloc, free, mallocHost, freeHost,
    peekArray, peekArrayAsync, pokeArray, pokeArrayAsync, memset,
    getDevicePtr
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Internal.C2HS
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Driver.Stream               (Stream, withStream)

-- System
import Foreign                           hiding (peekArray, pokeArray, malloc, free)
import Foreign.C
import Control.Monad                            (liftM)
import Unsafe.Coerce

#c
typedef enum CUmemhostalloc_option_enum {
    CU_MEMHOSTALLOC_OPTION_PORTABLE       = CU_MEMHOSTALLOC_PORTABLE,
    CU_MEMHOSTALLOC_OPTION_DEVICE_MAPPED  = CU_MEMHOSTALLOC_DEVICEMAP,
    CU_MEMHOSTALLOC_OPTION_WRITE_COMBINED = CU_MEMHOSTALLOC_WRITECOMBINED
} CUmemhostalloc_option;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to memory allocated on the device
--
newtype DevicePtr a = DevicePtr { useDevicePtr :: {# type CUdeviceptr #}}
  deriving (Show)

-- |
-- A reference to memory on the host that is page-locked and directly-accessible
-- from the device. Since the memory can be accessed directly, it can be read or
-- written at a much higher bandwidth than pageable memory from the traditional
-- malloc.
--
-- The driver automatically accelerates calls to functions such as `memcpy'
-- which reference page-locked memory.
--
newtype HostPtr a = HostPtr (ForeignPtr a)

withHostPtr :: HostPtr a -> (Ptr a -> IO b) -> IO b
withHostPtr (HostPtr hptr) = withForeignPtr hptr

newHostPtr :: IO (HostPtr a)
newHostPtr = HostPtr `fmap` mallocForeignPtrBytes (sizeOf (undefined :: Ptr ()))

-- |
-- Options for host allocation
--
{# enum CUmemhostalloc_option as AllocFlag
    { underscoreToCase }
    with prefix="CU_MEMHOSTALLOC_OPTION" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Host Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the host which is page-locked and
-- directly accessible from the device. Note that since the amount of pageable
-- memory is thusly reduced, overall system performance may suffer. This is best
-- used sparingly to allocate staging areas for data exchange.
--
mallocHost :: [AllocFlag] -> Int -> IO (Either String (HostPtr a))
mallocHost flags bytes =
  newHostPtr >>= \hp -> withHostPtr hp $ \p ->
  cuMemHostAlloc p bytes flags >>= \rv ->
  return $ case nothingIfOk rv of
    Nothing  -> Right hp
    Just err -> Left err

{# fun unsafe cuMemHostAlloc
  { with'* `Ptr a'
  ,        `Int'
  , combineBitMasks `[AllocFlag]' } -> `Status' cToEnum #}
  where with' = with . castPtr

-- |
-- Free a section of page-locked host memory
--
freeHost :: HostPtr a -> IO (Maybe String)
freeHost hp = withHostPtr hp $ \p -> (nothingIfOk `fmap` cuMemFreeHost p)

{# fun unsafe cuMemFreeHost
  { castPtr `Ptr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Device Allocation
--------------------------------------------------------------------------------

-- |
-- Allocate a section of linear memory on the device, and return a reference to
-- it. The memory is suitably aligned for any type, and is not cleared.
--
malloc :: Int -> IO (Either String (DevicePtr a))
malloc bytes = resultIfOk `fmap` cuMemAlloc bytes

{# fun unsafe cuMemAlloc
  { alloca-  `DevicePtr a' dptr*
  ,          `Int'               } -> `Status' cToEnum #}
  where dptr = liftM DevicePtr . peekIntConv

-- |
-- Release a section of device memory
--
free :: DevicePtr a -> IO (Maybe String)
free dp = nothingIfOk `fmap` cuMemFree dp

{# fun unsafe cuMemFree
  { useDevicePtr `DevicePtr a' } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Transfer
--------------------------------------------------------------------------------

-- |
-- Copy a number of elements from the device to host memory. This is a
-- synchronous operation
--
peekArray :: Storable a => Int -> DevicePtr a -> Ptr a -> IO (Maybe String)
peekArray n dptr hptr = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPeek x _ = nothingIfOk `fmap` cuMemcpyDtoH hptr dptr (n * sizeOf x)

{# fun unsafe cuMemcpyDtoH
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory from the device asynchronously, possibly associated with a
-- particular stream. The destination host memory must be page-locked (allocated
-- by mallocHost)
--
peekArrayAsync :: Storable a => Int -> DevicePtr a -> Ptr a -> Maybe Stream -> IO (Maybe String)
peekArrayAsync n dptr hptr mst = doPeek undefined dptr
  where
    doPeek :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPeek x _ =
      nothingIfOk `fmap` case mst of
        Nothing -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) nullPtr
        Just st -> (withStream st $ \s -> cuMemcpyDtoHAsync hptr dptr (n * sizeOf x) s)

{# fun unsafe cuMemcpyDtoHAsync
  { castPtr      `Ptr a'
  , useDevicePtr `DevicePtr a'
  ,              `Int'
  , castPtr      `Ptr Stream'  } -> `Status' cToEnum #}

-- |
-- Copy a number of elements onto the device. This is a synchronous operation
--
pokeArray :: Storable a => Int -> Ptr a -> DevicePtr a -> IO (Maybe String)
pokeArray n hptr dptr = doPoke undefined dptr
  where
    doPoke :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    doPoke x _ = nothingIfOk `fmap` cuMemcpyHtoD dptr hptr (n * sizeOf x)

{# fun unsafe cuMemcpyHtoD
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Copy memory onto the device asynchronously, possibly associated with a
-- particular stream. The source host memory must be page-locked (allocated by
-- mallocHost)
--
pokeArrayAsync :: Storable a => Int -> Ptr a -> DevicePtr a -> Maybe Stream -> IO (Maybe String)
pokeArrayAsync n hptr dptr mst = dopoke undefined dptr
  where
    dopoke :: Storable a' => a' -> DevicePtr a' -> IO (Maybe String)
    dopoke x _ =
      nothingIfOk `fmap` case mst of
        Nothing -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) nullPtr
        Just st -> (withStream st $ \s -> cuMemcpyHtoDAsync dptr hptr (n * sizeOf x) s)

{# fun unsafe cuMemcpyHtoDAsync
  { useDevicePtr `DevicePtr a'
  , castPtr      `Ptr a'
  ,              `Int'
  , castPtr      `Ptr Stream'  } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Utility
--------------------------------------------------------------------------------

-- |
-- Set a number of data elements to the specified value, which may be either 8-,
-- 16-, or 32-bits wide.
--
memset :: Storable a => DevicePtr a -> Int -> a -> IO (Maybe String)
memset dptr n val = case sizeOf val of
    1 -> nothingIfOk `fmap` cuMemsetD8  dptr val n
    2 -> nothingIfOk `fmap` cuMemsetD16 dptr val n
    4 -> nothingIfOk `fmap` cuMemsetD32 dptr val n
    _ -> return $ Just "can only memset 8-, 16-, and 32-bit values"

--
-- We use unsafe coerce below to reinterpret the bits of the value to memset as,
-- into the integer type required by the setting functions.
--
{# fun unsafe cuMemsetD8
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD16
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

{# fun unsafe cuMemsetD32
  { useDevicePtr `DevicePtr a'
  , unsafeCoerce `a'
  ,              `Int'         } -> `Status' cToEnum #}

-- |
-- Return the device pointer associated with a mapped, pinned host buffer, which
-- was allocated with the DEVICE_MAPPED option.
--
-- Currently, no options are supported and this must be empty.
--
getDevicePtr :: [AllocFlag] -> HostPtr a -> IO (Either String (DevicePtr a))
getDevicePtr flags hptr = withHostPtr hptr $ \hp -> (resultIfOk `fmap` cuMemHostGetDevicePointer hp flags)

{# fun unsafe cuMemHostGetDevicePointer
  { alloca-         `DevicePtr a' dptr*
  , castPtr         `Ptr a'
  , combineBitMasks `[AllocFlag]'       } -> `Status' cToEnum #}
  where dptr = liftM DevicePtr . peekIntConv