module Data.Extensible.Struct (
Struct
, set
, get
, new
, newRepeat
, newFor
, newFromHList
, atomicModify
, atomicModify'
, atomicModify_
, atomicModify'_
, (:*)
, unsafeFreeze
, newFrom
, hlookup
, hlength
, hfoldrWithIndex
, thaw
, hfrozen
, hmodify
, toHList) where
import GHC.Prim
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Constraint
import Data.Extensible.Class
import Data.Extensible.Internal
import Control.Comonad
import Data.Profunctor.Rep
import Data.Profunctor.Sieve
import qualified Data.Extensible.HList as L
import GHC.Types
data Struct s (h :: k -> *) (xs :: [k]) = Struct (SmallMutableArray# s Any)
set :: PrimMonad m => Struct (PrimState m) h xs -> Membership xs x -> h x -> m ()
set (Struct m) (getMemberId -> I# i) e = primitive
$ \s -> case unsafeCoerce# writeSmallArray# m i e s of
s' -> (# s', () #)
get :: PrimMonad m => Struct (PrimState m) h xs -> Membership xs x -> m (h x)
get (Struct m) (getMemberId -> I# i) = primitive $ unsafeCoerce# readSmallArray# m i
atomicModify :: PrimMonad m
=> Struct (PrimState m) h xs -> Membership xs x -> (h x -> (h x, a)) -> m a
atomicModify (Struct m) (getMemberId -> I# i) f = primitive
$ \s0 -> case readSmallArray# m i s0 of
(# s, x #) -> retry x s
where
retry x s = let p = unsafeCoerce# f x in
case casSmallArray# m i x (fst p) s of
(# s', b, y #) -> case b of
0# -> (# s', snd p #)
_ -> retry y s'
atomicModify' :: PrimMonad m
=> Struct (PrimState m) h xs -> Membership xs x -> (h x -> (h x, a)) -> m a
atomicModify' s i f = atomicModify s i
(\x -> let (y, a) = f x in (y, y `seq` a))
>>= (return $!)
atomicModify_ :: PrimMonad m
=> Struct (PrimState m) h xs -> Membership xs x -> (h x -> h x) -> m (h x)
atomicModify_ (Struct m) (getMemberId -> I# i) f = primitive
$ \s0 -> case readSmallArray# m i s0 of
(# s, x #) -> retry x s
where
retry x s = case casSmallArray# m i x (unsafeCoerce# f x) s of
(# s', b, y #) -> case b of
0# -> (# s', unsafeCoerce# y #)
_ -> retry y s'
atomicModify'_ :: PrimMonad m
=> Struct (PrimState m) h xs -> Membership xs x -> (h x -> h x) -> m (h x)
atomicModify'_ s i f = atomicModify_ s i f >>= (return $!)
new :: forall h m xs. (PrimMonad m, Generate xs)
=> (forall x. Membership xs x -> h x)
-> m (Struct (PrimState m) h xs)
new = newDict Dict
newDict :: PrimMonad m
=> Dict (Generate xs)
-> (forall x. Membership xs x -> h x)
-> m (Struct (PrimState m) h xs)
newDict Dict k = do
m <- newRepeat undefined
henumerate (\i cont -> set m i (k i) >> cont) $ return m
newRepeat :: forall h m xs. (PrimMonad m, Generate xs)
=> (forall x. h x)
-> m (Struct (PrimState m) h xs)
newRepeat x = do
let !(I# n) = hcount (Proxy :: Proxy xs)
primitive $ \s -> case newSmallArray# n (unsafeCoerce# x) s of
(# s', a #) -> (# s', Struct a #)
newFor :: forall proxy c h m xs. (PrimMonad m, Forall c xs)
=> proxy c
-> (forall x. c x => Membership xs x -> h x)
-> m (Struct (PrimState m) h xs)
newFor = newForDict Dict
newForDict :: forall proxy c h m xs. (PrimMonad m)
=> Dict (Forall c xs)
-> proxy c
-> (forall x. c x => Membership xs x -> h x)
-> m (Struct (PrimState m) h xs)
newForDict Dict p k = do
m <- newRepeat undefined
henumerateFor p (Proxy :: Proxy xs) (\i cont -> set m i (k i) >> cont) $ return m
newFromHList :: forall h m xs. PrimMonad m => L.HList h xs -> m (Struct (PrimState m) h xs)
newFromHList l = do
let !(I# size) = L.hlength l
m <- primitive $ \s -> case newSmallArray# size undefined s of
(# s', a #) -> (# s', Struct a #)
let go :: Int -> L.HList h t -> m ()
go _ L.HNil = return ()
go i (L.HCons x xs) = set m (unsafeMembership i) x >> go (i + 1) xs
go 0 l
return m
data (h :: k -> *) :* (s :: [k]) = HProduct (SmallArray# Any)
unsafeFreeze :: PrimMonad m => Struct (PrimState m) h xs -> m (h :* xs)
unsafeFreeze (Struct m) = primitive $ \s -> case unsafeFreezeSmallArray# m s of
(# s', a #) -> (# s', HProduct a #)
thaw :: PrimMonad m => h :* xs -> m (Struct (PrimState m) h xs)
thaw (HProduct ar) = primitive $ \s -> case thawSmallArray# ar 0# (sizeofSmallArray# ar) s of
(# s', m #) -> (# s', Struct m #)
hlength :: h :* xs -> Int
hlength (HProduct ar) = I# (sizeofSmallArray# ar)
unsafeMembership :: Int -> Membership xs x
unsafeMembership = unsafeCoerce#
hfoldrWithIndex :: (forall x. Membership xs x -> h x -> r -> r) -> r -> h :* xs -> r
hfoldrWithIndex f r p = foldr
(\i -> let m = unsafeMembership i in f m (hlookup m p)) r [0..hlength p 1]
toHList :: forall h xs. h :* xs -> L.HList h xs
toHList p = go 0 where
go :: Int -> L.HList h xs
go i
| i == hlength p = unknownHList L.HNil
| otherwise = unknownHList $ L.HCons (hlookup (unsafeMembership i) p) (go (i + 1))
unknownHList :: L.HList h ys -> L.HList h zs
unknownHList = unsafeCoerce#
newFrom :: forall g h m xs. (PrimMonad m)
=> g :* xs
-> (forall x. Membership xs x -> g x -> h x)
-> m (Struct (PrimState m) h xs)
newFrom hp@(HProduct ar) k = do
let !n = sizeofSmallArray# ar
st <- primitive $ \s -> case newSmallArray# n undefined s of
(# s', a #) -> (# s', Struct a #)
let go i
| i == I# n = return st
| otherwise = do
let !m = unsafeMembership i
set st m $! k m (hlookup m hp)
go (i + 1)
go 0
hlookup :: Membership xs x -> h :* xs -> h x
hlookup (getMemberId -> I# i) (HProduct ar) = case indexSmallArray# ar i of
(# a #) -> unsafeCoerce# a
hfrozen :: (forall s. ST s (Struct s h xs)) -> h :* xs
hfrozen m = runST $ m >>= unsafeFreeze
hmodify :: (forall s. Struct s h xs -> ST s ()) -> h :* xs -> h :* xs
hmodify f m = runST $ do
s <- thaw m
f s
unsafeFreeze s
instance (Corepresentable p, Comonad (Corep p), Functor f) => Extensible f p (:*) where
pieceAt i pafb = cotabulate $ \ws -> sbt (extract ws) <$> cosieve pafb (hlookup i <$> ws) where
sbt xs !x = hmodify (\s -> set s i x) xs