aboutsummaryrefslogtreecommitdiff
path: root/app/Store.hs
blob: 495f1904c169ff3a87681d8387ea519a12f8cfea (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE StrictData #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- |  Read-only on-disk key-value store

module Store where

import Control.DeepSeq (NFData)
import Control.Monad.IO.Class (liftIO)
import qualified Control.Monad.State.Strict as S
import qualified Data.ByteString as BS
import qualified Data.ByteString.Short as BSS
import Data.Either (Either)
import qualified Data.Map.Strict as M
import Data.Serialize (Serialize, decode, encode, get, put)
import GHC.Generics (Generic)
import Prelude hiding (lookup)
import System.Directory (doesFileExist)
import System.FilePath ((</>))
import System.IO (Handle, IOMode(..), hTell, withFile)
import System.IO.MMap (mmapFileByteString)

data Store = Store
  { index :: M.Map BSS.ShortByteString Location
  , values :: BS.ByteString
  }
  
data Location = Location
  { offset :: Int
  , length :: Int
  } deriving (Show, Eq, Ord, Generic, NFData)

instance Serialize Location

#if MIN_VERSION_cereal(0,5,8)
#else
instance Serialize BSS.ShortByteString where
  put = put . BSS.fromShort
  get = BSS.toShort <$> get
#endif  

class StoreItem item where
  toByteString :: item -> BS.ByteString
  fromByteString :: BS.ByteString -> Either String item
  type KeyArgs item = arg | arg -> item
  itemKey :: KeyArgs item -> BSS.ShortByteString

indexFileName :: FilePath
indexFileName = "index"

valuesFileName :: FilePath
valuesFileName = "values"

data ReadMode
  = ReadEntireFile
  | MemoryMapFile
  deriving (Show, Eq)
 
load :: FilePath -> ReadMode -> IO (Either String Store)
load directoryPath readMode = do
  let valuesFilePath = directoryPath </> valuesFileName
      indexFilePath = directoryPath </> indexFileName
  (valuesFileExists, indexFileExists) <-
    (,) <$> doesFileExist indexFilePath <*> doesFileExist valuesFilePath
  case (valuesFileExists, indexFileExists) of
    (True, True) -> do
      indexFile <- BS.readFile indexFilePath
      valuesFile <-
        case readMode of
          ReadEntireFile -> BS.readFile valuesFilePath
          MemoryMapFile -> mmapFileByteString valuesFilePath Nothing
      let eitherIndex = decode @(M.Map BSS.ShortByteString Location) indexFile
      case eitherIndex of
        Right locMap -> return . Right $ Store {index = locMap, values = valuesFile}
        Left err -> return . Left $ "Error while decoding index : " ++ err
    (False, False) ->
      return . Left $ "Cannot find index and values in " ++ directoryPath
    (True, False) -> return . Left $ "Cannot find index in " ++ directoryPath
    (False, True) -> return . Left $ "Cannot find values in " ++ directoryPath

lookup :: (StoreItem item) => KeyArgs item -> Store -> Either String item
lookup keyArgs = lookupByteString (itemKey keyArgs)
      
lookupByteString ::
     (StoreItem item) => BSS.ShortByteString -> Store -> Either String item
lookupByteString key store =
  case M.lookup key (index store) of
    Just (Location off len) ->
      fromByteString . BS.take len . BS.drop off $ values store
    Nothing -> Left $ "Cannot find item with key " ++ show key

data State =
  State (M.Map BSS.ShortByteString Location)
        Handle

add :: (StoreItem item) => item -> KeyArgs item -> S.StateT State IO ()
add item keyArgs = do
  let bs = toByteString item
      len = BS.length bs
  State locMap handle <- S.get
  off <- liftIO $ fromIntegral <$> hTell handle
  _ <- liftIO $ BS.hPut handle bs
  S.put $ State (M.insert (itemKey keyArgs) (Location off len) locMap) handle

createStore ::
     FilePath -> (Handle -> IO (M.Map BSS.ShortByteString Location)) -> IO ()
createStore directoryPath action =
  withFile (directoryPath </> valuesFileName) WriteMode $ \valuesHandle -> do
    locMap <- action valuesHandle
    BS.writeFile (directoryPath </> indexFileName) (encode locMap)
  
writeValues ::
     Handle
  -> M.Map BSS.ShortByteString Location
  -> S.StateT State IO ()
  -> IO (M.Map BSS.ShortByteString Location)
writeValues handle locMap action = do
  State locMap' _ <- S.execStateT action (State locMap handle)
  return locMap'