src/Control/Addressbook/Query.hs

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

src/Control/Addressbook/Streaming.hs

@@ -17,7 +17,7 @@ import System.IO
 
 import qualified Data.Foldable as F
 
-import qualified Data.Map.Strict as Map
+import qualified Data.Trie as Trie
 
 import Data.Maybe
        (fromMaybe)
@@ -36,7 +36,7 @@ run = do
   datDir <- fromMaybe "./" <$> lookupEnv "HOME"
   runResourceT $ do
     x <- runConduit stream
-    runConduit (CL.sourceList (Map.keys x) .| C.map (<> "\n") .| CB.sinkFileCautious (datDir </> ".addressbook.dat"))
+    runConduit (CL.sourceList (Trie.keys x) .| C.map (<> "\n") .| CB.sinkFileCautious (datDir </> ".addressbook.dat"))
   where
     separate = \case
       From x -> [x]
@@ -49,4 +49,4 @@ run = do
         .| combine
         .| C.concatMap separate
         .| CT.encode CT.utf8
-        .| C.foldMap (`Map.singleton` ())
+        .| C.foldMap (`Trie.singleton` ())