{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE UnboxedTuples         #-}

module Control.Lens.Mutable.Internal where

-- external
import           Control.Concurrent.STM.TMVar (TMVar)
import           Control.Lens                 (Lens')
import           Data.Primitive.MutVar        (MutVar (..))
import           GHC.Conc                     (TVar (..))
import           GHC.Exts                     (MVar#, RealWorld, State#,
                                               newMVar#, newMutVar#, newTVar#,
                                               putMVar#, readMutVar#, readTVar#,
                                               retry#, takeMVar#, writeMutVar#,
                                               writeTVar#)
import           GHC.IORef                    (IORef (..))
import           GHC.MVar                     (MVar (..))
import           GHC.STRef                    (STRef (..))
import           GHC.Stack                    (HasCallStack)
import           Unsafe.Coerce                (unsafeCoerce)

-- internal
import           Control.Lens.Mutable.Types


-- | Convert a reference type to a 'Lens''.
class AsLens s a ref where
  asLens :: ref a -> Lens' s a

instance AsLens (S 'OpST s) a (MutVar s) where
  asLens :: MutVar s a -> Lens' (S 'OpST s) a
asLens (MutVar MutVar# s a
var#) a -> f a
f (S State# s
s1#) =
    let !(# State# s
s2#, a
valr #) = MutVar# s a -> State# s -> (# State# s, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
readMutVar# MutVar# s a
var# State# s
s1#
    in  (a -> S 'OpST s) -> f a -> f (S 'OpST s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
valw -> State# s -> S 'OpST s
forall (p :: PrimOpGroup) s. State# s -> S p s
S (MutVar# s a -> a -> State# s -> State# s
forall d a. MutVar# d a -> a -> State# d -> State# d
writeMutVar# MutVar# s a
var# a
valw State# s
s2#)) (a -> f a
f a
valr)
  {-# INLINE asLens #-}

instance AsLens (S 'OpST s) a (STRef s) where
  asLens :: STRef s a -> Lens' (S 'OpST s) a
asLens (STRef MutVar# s a
var#) a -> f a
f (S State# s
s1#) =
    let !(# State# s
s2#, a
valr #) = MutVar# s a -> State# s -> (# State# s, a #)
forall d a. MutVar# d a -> State# d -> (# State# d, a #)
readMutVar# MutVar# s a
var# State# s
s1#
    in  (a -> S 'OpST s) -> f a -> f (S 'OpST s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
valw -> State# s -> S 'OpST s
forall (p :: PrimOpGroup) s. State# s -> S p s
S (MutVar# s a -> a -> State# s -> State# s
forall d a. MutVar# d a -> a -> State# d -> State# d
writeMutVar# MutVar# s a
var# a
valw State# s
s2#)) (a -> f a
f a
valr)
  {-# INLINE asLens #-}

instance AsLens (S 'OpST RealWorld) a IORef where
  asLens :: IORef a -> Lens' (S 'OpST RealWorld) a
asLens (IORef STRef RealWorld a
stref) = STRef RealWorld a -> Lens' (S 'OpST RealWorld) a
forall s a (ref :: * -> *). AsLens s a ref => ref a -> Lens' s a
asLens STRef RealWorld a
stref
  {-# INLINE asLens #-}

-- | View a @'MVar' a@ as a @'SLens' \''OpST' 'RealWorld' a@.
--
-- Note: when this is eventually run in 'IO', the action will block the thread
-- until there is a value present, as per the semantics of 'takeMVar#'. It will
-- then put a value into the 'MVar', which will block the thread if the value
-- is absent. GHC doesn't give atomicity guarantees for 'MVar' so it's possible
-- this does happen, e.g. if another producer managed to "get in there" during
-- the intervening period between the two operations. Unfortunately GHC does
-- not provide an atomic @modifyMVar@ function or primop.
--
-- If you don't want to deal with this, don't use an 'MVar', use a 'TMVar'.
instance AsLens (S 'OpMVar RealWorld) a MVar where
  asLens :: MVar a -> Lens' (S 'OpMVar RealWorld) a
asLens (MVar MVar# RealWorld a
var#) a -> f a
f (S State# RealWorld
s1#) =
    let !(# State# RealWorld
s2#, a
valr #) = MVar# RealWorld a -> State# RealWorld -> (# State# RealWorld, a #)
forall d a. MVar# d a -> State# d -> (# State# d, a #)
takeMVar# MVar# RealWorld a
var# State# RealWorld
s1#
    in  (a -> S 'OpMVar RealWorld) -> f a -> f (S 'OpMVar RealWorld)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
valw -> State# RealWorld -> S 'OpMVar RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S (MVar# RealWorld a -> a -> State# RealWorld -> State# RealWorld
forall d a. MVar# d a -> a -> State# d -> State# d
putMVar# MVar# RealWorld a
var# a
valw State# RealWorld
s2#)) (a -> f a
f a
valr)
  {-# INLINE asLens #-}

instance AsLens (S 'OpSTM RealWorld) a TVar where
  asLens :: TVar a -> Lens' (S 'OpSTM RealWorld) a
asLens (TVar TVar# RealWorld a
var#) a -> f a
f (S State# RealWorld
s1#) =
    let !(# State# RealWorld
s2#, a
valr #) = TVar# RealWorld a -> State# RealWorld -> (# State# RealWorld, a #)
forall d a. TVar# d a -> State# d -> (# State# d, a #)
readTVar# TVar# RealWorld a
var# State# RealWorld
s1#
    in  (a -> S 'OpSTM RealWorld) -> f a -> f (S 'OpSTM RealWorld)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
valw -> State# RealWorld -> S 'OpSTM RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S (TVar# RealWorld a -> a -> State# RealWorld -> State# RealWorld
forall d a. TVar# d a -> a -> State# d -> State# d
writeTVar# TVar# RealWorld a
var# a
valw State# RealWorld
s2#)) (a -> f a
f a
valr)
  {-# INLINE asLens #-}

instance AsLens (S 'OpSTM RealWorld) a TMVar where
  asLens :: TMVar a -> Lens' (S 'OpSTM RealWorld) a
asLens (TMVar a
tmvar :: TMVar a) a -> f a
f (S State# RealWorld
s1#) =
    -- they hide the constructor (no exports) but it's just a newtype
    let !(TVar TVar# RealWorld (Maybe a)
var#)      = (TMVar a -> TVar (Maybe a)
forall a b. a -> b
unsafeCoerce TMVar a
tmvar :: TVar (Maybe a))
        !(# State# RealWorld
s2#, Maybe a
valr' #) = TVar# RealWorld (Maybe a)
-> State# RealWorld -> (# State# RealWorld, Maybe a #)
forall d a. TVar# d a -> State# d -> (# State# d, a #)
readTVar# TVar# RealWorld (Maybe a)
var# State# RealWorld
s1#
        valr :: a
valr              = case Maybe a
valr' of
          Just a
v  -> a
v
          Maybe a
Nothing -> let (# State# RealWorld
_, b
a #) = State# RealWorld -> (# State# RealWorld, b #)
forall a. State# RealWorld -> (# State# RealWorld, a #)
retry# State# RealWorld
s1# in a
forall b. b
a
    in  (a -> S 'OpSTM RealWorld) -> f a -> f (S 'OpSTM RealWorld)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
valw -> State# RealWorld -> S 'OpSTM RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S (TVar# RealWorld (Maybe a)
-> Maybe a -> State# RealWorld -> State# RealWorld
forall d a. TVar# d a -> a -> State# d -> State# d
writeTVar# TVar# RealWorld (Maybe a)
var# (a -> Maybe a
forall a. a -> Maybe a
Just a
valw) State# RealWorld
s2#)) (a -> f a
f a
valr)
  {-# INLINE asLens #-}


-- | A state in which you can allocate new references.
--
-- This can be defined on either pure or impure references. For pure references
-- one could e.g. define an instance of this on @Map k v@ with @Const k@ as the
-- reference type - see unit tests for an example.
class AsLens s a ref => Allocable s a ref where

  -- | Allocate a new reference with the given value.
  alloc :: a -> s -> (ref a, s)

  {- | Deallocate an existing reference, and return its value.

  The default implementation simply writes 'error' into the reference and
  returns the old value. The caller is responsible for actually throwing away
  the reference and never using it again, as per Haskell's GC semantics.
  -}
  free :: HasCallStack => ref a -> s -> (a, s)
  free ref a
ref = ref a -> (a -> (a, a)) -> s -> (a, s)
forall s a (ref :: * -> *). AsLens s a ref => ref a -> Lens' s a
asLens ref a
ref (, [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"use-after-free")
  {-# INLINE free #-}

  {- | Check if a reference is valid.

  The default implementation simply forces the reference and returns 'True'. If
  the reference has already been freed (via 'free') then an error will be
  raised, which you can catch in the 'IO' monad as per usual. In other words,
  the default implementation will never return 'False'.
  -}
  isValid :: ref a -> s -> (Bool, s)
  isValid ref a
ref = ref a -> Lens' s a
forall s a (ref :: * -> *). AsLens s a ref => ref a -> Lens' s a
asLens ref a
ref ((a -> (Bool, a)) -> s -> (Bool, s))
-> (a -> (Bool, a)) -> s -> (Bool, s)
forall a b. (a -> b) -> a -> b
$ \a
r -> (a
r a -> Bool -> Bool
`seq` Bool
True, a
r)
  {-# INLINE isValid #-}

instance Allocable (S 'OpST s) a (MutVar s) where
  alloc :: a -> S 'OpST s -> (MutVar s a, S 'OpST s)
alloc a
val (S State# s
s1#) =
    let !(# State# s
s2#, MutVar# s a
var# #) = a -> State# s -> (# State# s, MutVar# s a #)
forall a d. a -> State# d -> (# State# d, MutVar# d a #)
newMutVar# a
val State# s
s1# in (MutVar# s a -> MutVar s a
forall s a. MutVar# s a -> MutVar s a
MutVar MutVar# s a
var#, State# s -> S 'OpST s
forall (p :: PrimOpGroup) s. State# s -> S p s
S State# s
s2#)
  {-# INLINE alloc #-}

instance Allocable (S 'OpST s) a (STRef s) where
  alloc :: a -> S 'OpST s -> (STRef s a, S 'OpST s)
alloc a
val (S State# s
s1#) =
    let !(# State# s
s2#, MutVar# s a
var# #) = a -> State# s -> (# State# s, MutVar# s a #)
forall a d. a -> State# d -> (# State# d, MutVar# d a #)
newMutVar# a
val State# s
s1# in (MutVar# s a -> STRef s a
forall s a. MutVar# s a -> STRef s a
STRef MutVar# s a
var#, State# s -> S 'OpST s
forall (p :: PrimOpGroup) s. State# s -> S p s
S State# s
s2#)
  {-# INLINE alloc #-}

instance Allocable (S 'OpST RealWorld) a IORef where
  alloc :: a -> S 'OpST RealWorld -> (IORef a, S 'OpST RealWorld)
alloc a
val S 'OpST RealWorld
s = let (STRef RealWorld a
r, S 'OpST RealWorld
s') = a -> S 'OpST RealWorld -> (STRef RealWorld a, S 'OpST RealWorld)
forall s a (ref :: * -> *).
Allocable s a ref =>
a -> s -> (ref a, s)
alloc a
val S 'OpST RealWorld
s in (STRef RealWorld a -> IORef a
forall a. STRef RealWorld a -> IORef a
IORef STRef RealWorld a
r, S 'OpST RealWorld
s')
  {-# INLINE alloc #-}

instance Allocable (S 'OpMVar RealWorld) a MVar where
  alloc :: a -> S 'OpMVar RealWorld -> (MVar a, S 'OpMVar RealWorld)
alloc (a
val :: a) (S State# RealWorld
s1#) =
    let !(# State# RealWorld
s2#, MVar# RealWorld a
var# #) =
          State# RealWorld -> (# State# RealWorld, MVar# RealWorld a #)
forall d a. State# d -> (# State# d, MVar# d a #)
newMVar# State# RealWorld
s1# :: (# State# RealWorld, MVar# RealWorld a #)
    in  (MVar# RealWorld a -> MVar a
forall a. MVar# RealWorld a -> MVar a
MVar MVar# RealWorld a
var#, State# RealWorld -> S 'OpMVar RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S (MVar# RealWorld a -> a -> State# RealWorld -> State# RealWorld
forall d a. MVar# d a -> a -> State# d -> State# d
putMVar# MVar# RealWorld a
var# a
val State# RealWorld
s2#))
  {-# INLINE alloc #-}

instance Allocable (S 'OpSTM RealWorld) a TVar where
  alloc :: a -> S 'OpSTM RealWorld -> (TVar a, S 'OpSTM RealWorld)
alloc a
val (S State# RealWorld
s1#) =
    let !(# State# RealWorld
s2#, TVar# RealWorld a
var# #) = a -> State# RealWorld -> (# State# RealWorld, TVar# RealWorld a #)
forall a d. a -> State# d -> (# State# d, TVar# d a #)
newTVar# a
val State# RealWorld
s1# in (TVar# RealWorld a -> TVar a
forall a. TVar# RealWorld a -> TVar a
TVar TVar# RealWorld a
var#, State# RealWorld -> S 'OpSTM RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S State# RealWorld
s2#)
  {-# INLINE alloc #-}

instance Allocable (S 'OpSTM RealWorld) a TMVar where
  alloc :: a -> S 'OpSTM RealWorld -> (TMVar a, S 'OpSTM RealWorld)
alloc a
val (S State# RealWorld
s1#) =
    let !(# State# RealWorld
s2#, TVar# RealWorld (Maybe a)
var# #) = Maybe a
-> State# RealWorld
-> (# State# RealWorld, TVar# RealWorld (Maybe a) #)
forall a d. a -> State# d -> (# State# d, TVar# d a #)
newTVar# (a -> Maybe a
forall a. a -> Maybe a
Just a
val) State# RealWorld
s1#
    in  (TVar (Maybe a) -> TMVar a
forall a b. a -> b
unsafeCoerce (TVar# RealWorld (Maybe a) -> TVar (Maybe a)
forall a. TVar# RealWorld a -> TVar a
TVar TVar# RealWorld (Maybe a)
var#), State# RealWorld -> S 'OpSTM RealWorld
forall (p :: PrimOpGroup) s. State# s -> S p s
S State# RealWorld
s2#)
  {-# INLINE alloc #-}