{-# OPTIONS_GHC -cpp -fglasgow-exts #-}
{- |
   Module     : Data.Variables
   Copyright  : Copyright (C) 2006 Bulat Ziganshin
   License    : BSD3

   Maintainer : Bulat Ziganshin <bulatz@HotPOP.com>
   Stability  : experimental
   Portability: Hugs/GHC

Fast mutable variables

Based on the: Unboxed mutable Ints (c) The University of Glasgow 2002
-}

module DataVariables where

#include "MachDeps.h"

#ifndef SIZEOF_HSINT
#define SIZEOF_HSINT  INT_SIZE_IN_BYTES
#endif

import Control.Monad.ST         (ST)
import Data.Word
import Data.Bits
import Data.IORef
import Data.STRef
import Foreign.Storable
import Foreign.Ptr

#ifdef __GLASGOW_HASKELL__
import GHC.Exts
import GHC.IOBase
import GHC.ST
#endif


-- -----------------------------------------------------------------------------
-- Syntax sugar for using pointers arithmetic

infixl 6 +:, -:
ptr+:n   = ptr `plusPtr` (fromIntegral n)
ptr-:buf = fromIntegral  (ptr `minusPtr` buf)

{-# INLINE (+:)        #-}
{-# INLINE (-:)        #-}

-- -----------------------------------------------------------------------------
-- Syntax sugar for using mutable variables

infixl 0 =:, +=, -=, .=
ref x   = newMRef x
val ref = readMRef ref
ref=:x  = writeMRef ref x
ref+=x  = modifyMRef  ref (\old -> old+x)
ref-=x  = modifyMRef  ref (\old -> old-x)
ref.=f  = modifyMRef  ref (\old -> f old)
ref.<-f = modifyMRefM ref (\old -> f old)

-- -----------------------------------------------------------------------------
-- The uniform interface for mutable variables in any monad

-- |A mutable variable in some monad 'm'
--   (m->r restriction needed to simplify using of newMRef)
class (Monad m) => MRef m r | r->m, m->r where
    -- |Build a new 'MRef' with given initial value
    newMRef    :: a -> m (r a)
    -- |Read the value of an 'MRef'
    readMRef   :: r a -> m a
    -- |Write a new value into an 'MRef'
    writeMRef  :: r a -> a -> m ()

-- |Modify the contents of an 'MRef' by applying pure function to it
modifyMRef  ref f  =  readMRef ref >>= writeMRef ref . f

-- |Modify the contents of an 'MRef' by applying monadic computation to it
modifyMRefM ref f  =  readMRef ref >>= f >>= writeMRef ref

instance MRef IO IORef where
    newMRef = newIORef
    readMRef = readIORef
    writeMRef = writeIORef

instance MRef (ST s) (STRef s) where
    newMRef = newSTRef
    readMRef = readSTRef
    writeMRef = writeSTRef

-- -----------------------------------------------------------------------------
-- Mutable variables in any monad, with instances specific to value type (`a`)

-- |A mutable in some monad 'm' variable of type 'r' holding values of type 'a'
--    (m a->r restriction is needed to simplify using of newVar)
class (Monad m) => Var m r a | r->a, m a->r where
    -- |Build a new 'Var' with given initial value
    newVar   :: a -> m r
    -- |Read the value of an 'Var'
    readVar  :: r -> m a
    -- |Write a new value into an 'Var'
    writeVar :: r -> a -> m ()

#ifndef __GLASGOW_HASKELL__
instance Var IO (IORef a) a where
    newVar = newIORef
    readVar = readIORef
    writeVar = writeIORef
instance Var (ST s) (STRef s a) a where
    newVar = newSTRef
    readVar = readSTRef
    writeVar = writeSTRef
type IntVar  = IORef Int
type WordVar = IORef Word
type BoolVar = IORef Bool
#else /* __GLASGOW_HASKELL__ */
-- -----------------------------------------------------------------------------
-- Mutable integers

data MutInt s = MutInt (MutableByteArray# s)
type IntVar = MutInt RealWorld

instance Var IO IntVar Int where
    newVar = newInt IO
    readVar = readInt IO
    writeVar = writeInt IO

instance Var (ST s) (MutInt s) Int where
    newVar = newInt ST
    readVar = readInt ST
    writeVar = writeInt ST

newInt c (I# int) = c $ \s ->
  case newByteArray# size s of { (# s, arr #) ->
  case writeIntArray# arr 0# int s of { s ->
  (# s, MutInt arr #) } }
  where I# size = SIZEOF_HSINT

readInt c (MutInt arr) = c $ \s ->
  case readIntArray# arr 0# s of { (# s, int #) ->
  (# s, I# int #) }

writeInt c (MutInt arr) (I# int) = c $ \s ->
  case writeIntArray# arr 0# int s of { s ->
  (# s, () #) }

{-# INLINE readInt     #-}
{-# INLINE writeInt    #-}

-- -----------------------------------------------------------------------------
-- Mutable words (unsigned integers)

data MutWord s = MutWord (MutableByteArray# s)
type WordVar = MutWord RealWorld

instance Var IO WordVar Word where
    newVar = newWord IO
    readVar = readWord IO
    writeVar = writeWord IO

instance Var (ST s) (MutWord s) Word where
    newVar = newWord ST
    readVar = readWord ST
    writeVar = writeWord ST

newWord c (W# word) = c $ \s ->
  case newByteArray# size s of { (# s, arr #) ->
  case writeWordArray# arr 0# word s of { s ->
  (# s, MutWord arr #) } }
  where I# size = SIZEOF_HSWORD

readWord c (MutWord arr) = c $ \s ->
  case readWordArray# arr 0# s of { (# s, word #) ->
  (# s, W# word #) }

writeWord c (MutWord arr) (W# word) = c $ \s ->
  case writeWordArray# arr 0# word s of { s ->
  (# s, () #) }

{-# INLINE readWord    #-}
{-# INLINE writeWord   #-}

-- -----------------------------------------------------------------------------
-- Mutable bools

data MutBool s = MutBool (MutableByteArray# s)
type BoolVar = MutBool RealWorld

instance Var IO BoolVar Bool where
    newVar = newBool IO
    readVar = readBool IO
    writeVar = writeBool IO

instance Var (ST s) (MutBool s) Bool where
    newVar = newBool ST
    readVar = readBool ST
    writeVar = writeBool ST

newBool c b = c $ \s ->
  case newByteArray# size s of { (# s, arr #) ->
  case writeIntArray# arr 0# bool s of { s ->
  (# s, MutBool arr #) } }
  where I# size = SIZEOF_HSINT
        I# bool = fromEnum b

readBool c (MutBool arr) = c $ \s ->
  case readIntArray# arr 0# s of { (# s, int #) ->
  (# s, toEnum (I# int) #) }

writeBool c (MutBool arr) b = c $ \s ->
  case writeIntArray# arr 0# bool s of { s ->
  (# s, () #) }
  where I# bool = fromEnum b

{-# INLINE readBool    #-}
{-# INLINE writeBool   #-}

#endif /* __GLASGOW_HASKELL__ */

-- -----------------------------------------------------------------------------
-- Mutable pointers

newPtr      :: Ptr a -> IO (MutPtr a)
readPtr     :: MutPtr a -> IO (Ptr a)
writePtr    :: MutPtr a -> Ptr a -> IO ()

#ifdef __GLASGOW_HASKELL__

data MutPtr a = MutPtr (MutableByteArray# RealWorld)

newPtr (Ptr addr#) = IO $ \s ->
  case newByteArray# size s of { (# s, arr# #) ->
  case writeAddrArray# arr# 0# addr# s of { s ->
  (# s, MutPtr arr# #) } }
  where I# size = SIZEOF_HSPTR

readPtr (MutPtr arr#) = IO $ \s ->
  case readAddrArray# arr# 0# s of { (# s, addr# #) ->
  (# s, Ptr addr# #) }

writePtr (MutPtr arr#) (Ptr addr#) = IO $ \s ->
  case writeAddrArray# arr# 0# addr# s of { s ->
  (# s, () #) }

#else /* ! __GLASGOW_HASKELL__ */

type MutPtr a = IORef (Ptr a)
newPtr   = newIORef
readPtr  = readIORef
writePtr = writeIORef

#endif /* ! __GLASGOW_HASKELL__ */

{-# INLINE readPtr     #-}
{-# INLINE writePtr    #-}

-- -----------------------------------------------------------------------------
-- Peek/poke bytes in memory

readByteAt  :: Ptr a -> IO Int
writeByteAt :: Ptr a -> Int -> IO ()

#ifdef __GLASGOW_HASKELL__

readByteAt (Ptr addr#) = IO $ \s ->
  case readWord8OffAddr# addr# 0# s of { (# s, byte# #) ->
  (# s, I# (word2Int# byte#) #) }
writeByteAt (Ptr addr#) (I# byte#) = IO $ \s ->
  case writeWord8OffAddr# addr# 0# (int2Word# byte#) s of { s ->
  (# s, () #) }

#else /* ! __GLASGOW_HASKELL__ */

readByteAt  ptr = do byte <- peek (castPtr ptr) :: IO Word8
                     return $! (toEnum (fromIntegral byte))
writeByteAt ptr byte = poke (castPtr ptr) $! (fromIntegral (fromEnum byte)::Word8)

#endif /* ! __GLASGOW_HASKELL__ */

{-# INLINE readByteAt  #-}
{-# INLINE writeByteAt #-}

-- -----------------------------------------------------------------------------
-- Fast Ints & Bools

#ifdef __GLASGOW_HASKELL__

type FastInt = Int#
_ILIT (I# x) = x
iBox x = I# x
iUnbox (I# x) = x
quotFastInt   = quotInt#
negateFastInt = negateInt#
(<<#)  = iShiftL#
(>>#)  = uncheckedIShiftRL#
a &# b = word2Int# (int2Word# a `and#` int2Word# b)
a |# b = word2Int# (int2Word# a `or#`  int2Word# b)

type FastBool = Int#
fastBool True  = 1#
fastBool False = 0#
isFastTrue x = x ==# 1#

fastOr 1# _ = 1#
fastOr 0# x = x

fastAnd 0# x = 0#
fastAnd 1# x = x

#else /* ! __GLASGOW_HASKELL__ */

type FastInt =  Int
_ILIT  x = x :: Int
iBox   x = x :: Int
iUnbox x = x :: Int
a+#b = a+b
a-#b = a-b
a*#b = a*b
a `quotFastInt`   b   = a `quot`   b
a `negateFastInt` b   = a `negate` b
a==#b = a==b
a<#b  = a<b
a<=#b = a<=b
a>=#b = a>=b
a>#b  = a>b
a<<#b = a `shiftL` b
a>>#b = a `shiftR` b
a&#b  = a.&.b
a|#b  = a.|.b

type FastBool = Bool
fastBool x = x
_IS_TRUE_ x = x

#endif /* ! __GLASGOW_HASKELL__ */