aboutsummaryrefslogtreecommitdiff
path: root/app/Store.hs
blob: 956359d8ee4f234075b3b88d7682fe237696e76d (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
{-# 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
instance Serialize BSS.ShortByteString where
  put = put . BSS.fromShort
  get = BSS.toShort <$> get

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'