Simplify querying

This commit is contained in:
Mats Rauhala 2021-10-29 17:34:41 +03:00
parent e189f87dd7
commit 0c75d66122

View File

@ -4,15 +4,12 @@ module Control.Addressbook.Query where
import Data.Text import Data.Text
(Text) (Text)
import qualified Data.Text as T import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import Conduit import Conduit
import qualified Data.Conduit.Binary as CB import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.Combinators as C import qualified Data.Conduit.Combinators as C
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.Text as CT import qualified Data.Conduit.Text as CT
import qualified Data.Map.Strict as Map
import System.IO import System.IO
(stdout) (stdout)
@ -26,27 +23,20 @@ import System.FilePath
import Control.Exception import Control.Exception
(catch) (catch)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
-- XXX: The current Data.Trie implementation is much slower than Data.Set -- XXX: The current Data.Trie implementation is much slower than Data.Set
query :: Text -> IO () query :: Text -> IO ()
query prefix = do query prefix = do
datDir <- fromMaybe "./" <$> lookupEnv "HOME" datDir <- fromMaybe "./" <$> lookupEnv "HOME"
state <- catch @IOError (runResourceT $ runConduit $ readState datDir) (\_ -> pure Map.empty) catch @IOError (runResourceT $ runConduit $ filterAddresses datDir) $ \e -> print e
runConduit $ outputResults state
where where
readState :: FilePath -> ConduitM () Void (ResourceT IO) (Map.Map ByteString [Text]) prefixLC = T.toLower prefix
readState dir = filterAddresses :: FilePath -> ConduitT () Void (ResourceT IO) ()
filterAddresses dir =
CB.sourceFile (dir </> ".addressbook.dat") CB.sourceFile (dir </> ".addressbook.dat")
.| CT.decode CT.utf8 .| CT.decode CT.utf8
.| CT.lines .| CT.lines
.| C.foldMap (\s -> Map.singleton (TE.encodeUtf8 $ T.toLower s) [s]) .| C.filter (\v -> prefixLC `T.isPrefixOf` (T.toLower v))
prefixB = TE.encodeUtf8 prefix
outputResults :: Map.Map ByteString [Text] -> ConduitM () Void IO ()
outputResults state =
CL.sourceList (Map.elems $ Map.filterWithKey (\k _ -> prefixB `B.isPrefixOf` k) state)
.| C.concat
.| C.map (<> "\n")
.| CT.encode CT.utf8 .| CT.encode CT.utf8
.| C.map (<> "\n")
.| CB.sinkHandle stdout .| CB.sinkHandle stdout