module Hoogle.Store.Type(
Once, once, fromOnce, putOnce, getOnce,
SPut, runSPut, putByteString, putStorable, putDefer,
SGet, runSGet, getByteString, getStorable, getDefer, getLazyList
) where
import General.Base
import General.System
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader
import qualified Data.IntMap as IntMap
import Data.IORef
import Data.Typeable
import Foreign
import System.IO.Unsafe
import qualified Hoogle.Store.ReadBuffer as R
import qualified Hoogle.Store.WriteBuffer as W
stats = False
data Once a = Once {_onceKey :: Int, valueOnce :: a}
deriving Typeable
fromOnce :: Once a -> a
fromOnce = valueOnce
onceKeys :: Int -> IO Int
onceKeys = System.IO.Unsafe.unsafePerformIO $ do
ref <- newIORef 0
return $ \n -> atomicModifyIORef ref $ \x -> (x+n, x)
data SPutS = SPutS
{putBuffer :: W.Buffer
,putOnces :: IORef (IntMap.IntMap PutOnce)
,putPending :: IORef [SPut ()]
}
type SPut a = ReaderT SPutS IO a
modifyRef f x = liftIO . (`modifyIORef` x) =<< asks f
readPos = liftIO . W.getPos =<< asks putBuffer
runSPut :: FilePath -> SPut () -> IO ()
runSPut file act = withBinaryFile file WriteMode $ \h -> do
pending <- newIORef [act]
once <- newIORef IntMap.empty
W.withBuffer h $ \buffer -> do
let flush = do
xs <- liftIO $ readIORef pending
liftIO $ writeIORef pending []
forM_ xs $ \x -> do
x
flush
runReaderT flush $ SPutS buffer once pending
putByteString :: BString -> SPut ()
putByteString x = do
buf <- asks putBuffer
liftIO $ W.putByteString buf x
putStorable :: Storable a => a -> SPut ()
putStorable x = do
buf <- asks putBuffer
liftIO $ W.putStorable buf x
putDefer :: SPut () -> SPut ()
putDefer act = do
pos <- readPos
putStorable (0 :: Word32)
modifyRef putPending $ (:) $ do
val <- readPos
buf <- asks putBuffer
liftIO $ W.patch buf pos val
act
once :: a -> Once a
once x = System.IO.Unsafe.unsafePerformIO $ do
key <- onceKeys 1
return $ Once key x
type PutOnce = Either [Word32] Word32
putOnce :: (a -> SPut ()) -> Once a -> SPut ()
putOnce act (Once key x) = do
ref <- asks putOnces
mp <- liftIO $ readIORef ref
case fromMaybe (Left []) $ IntMap.lookup key mp of
Right val -> putStorable val
Left poss -> do
pos <- readPos
liftIO $ writeIORef ref $ IntMap.insert key (Left $ pos:poss) mp
putStorable (0 :: Word32)
when (null poss) $ modifyRef putPending $ (:) $ do
val <- readPos
mp <- liftIO $ readIORef ref
let Left poss = mp IntMap.! key
buf <- asks putBuffer
liftIO $ forM_ poss $ \pos -> W.patch buf pos val
liftIO $ writeIORef ref $ IntMap.insert key (Right val) mp
act x
data SGetS = SGetS {getBuffer :: R.Buffer, onceBase :: Int}
type SGet a = ReaderT SGetS IO a
runSGet :: Typeable a => FilePath -> SGet a -> IO a
runSGet file m = do
h <- openBinaryFile file ReadMode
sz <- hFileSize h
buf <- R.newBuffer h
one <- onceKeys $ fromIntegral sz
runReaderT (getDeferFrom 0 m) $ SGetS buf one
getStorable :: Typeable a => Storable a => SGet a
getStorable = do
buf <- asks getBuffer
res <- liftIO $ R.getStorable buf
when stats $ liftIO $ putStrLn $ "Reading storable " ++ show (sizeOf res)
return res
getByteString :: Word32 -> SGet BString
getByteString len = do
buf <- asks getBuffer
when stats $ liftIO $ putStrLn $ "Reading bytestring " ++ show len
liftIO $ R.getByteString buf $ fromIntegral len
getDefer :: Typeable a => SGet a -> SGet a
getDefer get = do
pos :: Word32 <- getStorable
getDeferFrom pos get
getDeferFrom :: forall a . Typeable a => Word32 -> SGet a -> SGet a
getDeferFrom pos get = do
s <- ask
liftIO $ unsafeInterleaveIO $ do
when stats $ putStrLn $ "Read at " ++ show (typeOf (undefined :: a))
R.setPos (getBuffer s) pos
runReaderT get s
getOnce :: Typeable a => SGet a -> SGet (Once a)
getOnce get = do
pos :: Word32 <- getStorable
x <- getDeferFrom pos get
one <- asks onceBase
return $ Once (fromIntegral pos + one) x
getLazyList :: SGet a -> Int -> Int -> SGet [a]
getLazyList get size n = do
s <- ask
pos <- liftIO $ R.getPos $ getBuffer s
liftIO $ forM [0..n1] $ \i -> unsafeInterleaveIO $ do
R.setPos (getBuffer s) (pos + fromIntegral (i * size))
runReaderT get s