1 {-# LANGUAGE CPP, ScopedTypeVariables, Rank2Types #-}
    2 
    3 -- | Mostly internal utilities for use by the rest of the library. Subject to
    4 -- removal without further notice.
    5 module Storage.Hashed.Utils where
    6 
    7 import Prelude hiding ( lookup, catch )
    8 import System.Mem( performGC )
    9 import System.IO.MMap( mmapFileByteString )
   10 import Bundled.Posix( getFileStatus, fileSize )
   11 import System.Directory( getCurrentDirectory, setCurrentDirectory )
   12 import Data.Int( Int64 )
   13 import Data.Maybe( catMaybes )
   14 import Control.Exception.Extensible( catch, bracket, SomeException(..) )
   15 import Control.Monad( when )
   16 import Control.Monad.Identity( runIdentity )
   17 import Control.Applicative( (<$>) )
   18 
   19 import Foreign.ForeignPtr( withForeignPtr )
   20 import Foreign.Ptr( plusPtr )
   21 import Data.ByteString.Internal( toForeignPtr, memcpy )
   22 import System.IO (withFile, IOMode(ReadMode), hSeek, SeekMode(AbsoluteSeek))
   23 import Data.Bits( Bits )
   24 #ifdef BIGENDIAN
   25 import Data.Bits( (.&.), (.|.), shift, shiftL, shiftR )
   26 #endif
   27 import Data.Path
   28 
   29 import qualified Data.ByteString.Lazy.Char8 as BL
   30 import qualified Data.ByteString.Char8 as BS8
   31 import qualified Data.ByteString as BS
   32 
   33 import qualified Data.Set as S
   34 import qualified Data.Map as M
   35 
   36 -- | Pointer to a filesystem, possibly with start/end offsets. Supposed to be
   37 -- fed to (uncurry mmapFileByteString) or similar.
   38 type FileSegment = (Absolute, Maybe (Int64, Int))
   39 
   40 -- | Read in a FileSegment into a Lazy ByteString. Implemented using mmap.
   41 readSegment :: FileSegment -> IO BL.ByteString
   42 readSegment (f,range) = do
   43     bs <- tryToRead
   44        `catch` (\(_::SomeException) -> do
   45                      size <- fileSize `fmap` getFileStatus (pathToString f)
   46                      if size == 0
   47                         then return BS8.empty
   48                         else performGC >> tryToRead)
   49     return $ BL.fromChunks [bs]
   50   where
   51     tryToRead = do 
   52         case range of
   53             Nothing -> BS.readFile (pathToString f)
   54             Just (off, size) -> withFile (pathToString f) ReadMode $ \h -> do
   55                 hSeek h AbsoluteSeek $ fromIntegral off
   56                 BS.hGet h size 
   57 {-# INLINE readSegment #-}
   58 
   59 -- | Run an IO action with @path@ as a working directory. Does neccessary
   60 -- bracketing.
   61 withCurrentDirectory :: (AbsRel x) => Path x -> IO a -> IO a
   62 withCurrentDirectory name =
   63     bracket
   64         (do cwd <- getCurrentDirectory
   65             (setCurrentDirectory $ pathToString name)
   66             return cwd)
   67         (\oldwd -> setCurrentDirectory oldwd
   68                      `catch` \(_::SomeException) -> return ())
   69         . const
   70 
   71 -- Wow, unsafe.
   72 unsafePokeBS :: BS8.ByteString -> BS8.ByteString -> IO ()
   73 unsafePokeBS to from =
   74     do let (fp_to, off_to, len_to) = toForeignPtr to
   75            (fp_from, off_from, len_from) = toForeignPtr from
   76        when (len_to /= len_from) $ fail $ "Length mismatch in unsafePokeBS: from = "
   77             ++ show len_from ++ " /= to = " ++ show len_to
   78        withForeignPtr fp_from $ \p_from ->
   79          withForeignPtr fp_to $ \p_to ->
   80            memcpy (plusPtr p_to off_to)
   81                   (plusPtr p_from off_from)
   82                   (fromIntegral len_to)
   83 
   84 align :: Integral a => a -> a -> a
   85 align boundary i = case i `rem` boundary of
   86                      0 -> i
   87                      x -> i + boundary - x
   88 {-# INLINE align #-}
   89 
   90 xlate32 :: (Bits a) => a -> a
   91 xlate64 :: (Bits a) => a -> a
   92 
   93 #ifdef LITTLEENDIAN
   94 xlate32 = id
   95 xlate64 = id
   96 #endif
   97 
   98 #ifdef BIGENDIAN
   99 bytemask :: (Bits a) => a
  100 bytemask = 255
  101 
  102 xlate32 a = ((a .&. (bytemask `shift`  0) `shiftL` 24)) .|.
  103             ((a .&. (bytemask `shift`  8) `shiftL`  8)) .|.
  104             ((a .&. (bytemask `shift` 16) `shiftR`  8)) .|.
  105             ((a .&. (bytemask `shift` 24) `shiftR` 24))
  106 xlate64 a = ((a .&. (bytemask `shift`  0) `shiftL` 56)) .|.
  107             ((a .&. (bytemask `shift`  8) `shiftL` 40)) .|.
  108             ((a .&. (bytemask `shift` 16) `shiftL` 24)) .|.
  109             ((a .&. (bytemask `shift` 24) `shiftL`  8)) .|.
  110             ((a .&. (bytemask `shift` 32) `shiftR`  8)) .|.
  111             ((a .&. (bytemask `shift` 40) `shiftR` 24)) .|.
  112             ((a .&. (bytemask `shift` 48) `shiftR` 40)) .|.
  113             ((a .&. (bytemask `shift` 56) `shiftR` 56))
  114 #endif
  115 
  116 -- | Find a monadic fixed point of @f@ that is the least above @i@. (Will
  117 -- happily diverge if there is none.)
  118 mfixFrom :: (Eq a, Functor m, Monad m) => (a -> m a) -> a -> m a
  119 mfixFrom f i = do x <- f i
  120                   if x == i then return i
  121                             else mfixFrom f x
  122 
  123 -- | Find a fixed point of @f@ that is the least above @i@. (Will happily
  124 -- diverge if there is none.)
  125 fixFrom :: (Eq a) => (a -> a) -> a -> a
  126 fixFrom f i = runIdentity $ mfixFrom (return . f) i
  127 
  128 -- | For a @refs@ function, a @map@ (@key@ -> @value@) and a @rootSet@, find a
  129 -- submap of @map@ such that all items in @map@ are reachable, through @refs@
  130 -- from @rootSet@.
  131 reachable :: forall monad key value. (Functor monad, Monad monad, Ord key, Eq value) =>
  132               (value -> monad [key])
  133            -> (key -> monad (Maybe (key, value)))
  134            -> S.Set key -> monad (M.Map key value)
  135 reachable refs lookup rootSet =
  136     do lookupSet rootSet >>= mfixFrom expand
  137     where lookupSet :: S.Set key -> monad (M.Map key value)
  138           expand :: M.Map key value -> monad (M.Map key value)
  139 
  140           lookupSet s = do list <- mapM lookup (S.toAscList s)
  141                            return $ M.fromAscList (catMaybes list)
  142           expand from = do refd <- concat <$> mapM refs (M.elems from)
  143                            M.union from <$> lookupSet (S.fromList refd)