Memory benchmark for the Text library

doc-store.hs

-- | A simple webserver to store documents. Its primary goal is as a benchmark
-- for the Text library.
{-# LANGUAGE OverloadedStrings #-}

import Data.Char (isPunctuation)
import Data.List (foldl')
import Data.Monoid (mconcat)
import Control.Applicative ((<$>))
import Control.Concurrent.MVar (MVar, modifyMVar_, newMVar, readMVar)
import Control.Monad.Reader (ReaderT, ask, runReaderT)
import Control.Monad.Trans (liftIO)
import Data.Maybe (fromMaybe)

import Data.Map (Map)
import Data.Set (Set)
import Data.Text (Text)
import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as BL
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.Text as T
import qualified Data.Text.Encoding as T

import Text.Blaze (Html, toHtml)
import Text.Blaze.Renderer.Utf8 (renderHtml)
import qualified Text.Blaze.Html5 as H

import Snap.Types ( Snap, getParam, getRequestBody, modifyResponse, route
                  , setContentType, writeLBS
                  )
import Snap.Http.Server (httpServe, defaultConfig)

--------------------------------------------------------------------------------
-- Pure logic                                                                 --
--------------------------------------------------------------------------------

-- | Extract all tokens from a document
tokenize :: Text -> Set Text
tokenize = S.fromList .
    filter (not . T.null) . map stripPunctuation . T.words . T.toLower
  where
    -- | Remove leading and trailing punctuation marks from a token
    stripPunctuation =  T.dropWhileEnd isPunctuation . T.dropWhile isPunctuation

-- | Type used for the document store
type Store = Map Text (Set Int)

-- | Add a document to the store
addDocument :: Int -> Text -> Store -> Store
addDocument id' doc store = foldl' insert store (S.toList $ tokenize doc)
  where
    insert s t = M.insertWith S.union t (S.singleton id') s

--------------------------------------------------------------------------------
-- Web logic                                                                  --
--------------------------------------------------------------------------------

-- | Our application stack
type App = ReaderT (MVar Store) Snap

-- | Utility to respond using a blaze view
blaze :: Html -> App ()
blaze html = do
    modifyResponse $ setContentType "text/html; charset=UTF-8"
    writeLBS $ renderHtml html

-- | Add a document
documentAdd :: App ()
documentAdd = do
    Just id' <- fmap (read . BC.unpack) <$> getParam "id"
    doc <- T.decodeUtf8 . strict <$> getRequestBody
    mvar <- ask
    liftIO $ modifyMVar_ mvar $ return . addDocument id' doc
    blaze $ documentView doc
  where
    strict = mconcat . BL.toChunks

-- | Query for some documents
documentQuery :: App ()
documentQuery = do
    store <- liftIO . readMVar =<< ask
    query <- fmap T.decodeUtf8 <$> getParam "query"
    let results = fromMaybe S.empty $ flip M.lookup store =<< query
    blaze $ resultsView results

--------------------------------------------------------------------------------
-- Web view                                                                   --
--------------------------------------------------------------------------------

documentView :: Text -> Html
documentView = H.p . toHtml

resultsView :: Set Int -> Html
resultsView = H.ul . mconcat . map (H.li . toHtml) . S.toList

--------------------------------------------------------------------------------
-- Glueing together                                                           --
--------------------------------------------------------------------------------

-- | Application routing
app :: App ()
app = route
    [ ("/document/query/:query", documentQuery)
    , ("/document/:id",            documentAdd)
    ]

-- | Main function
main :: IO ()
main = do
    mvar <- newMVar M.empty
    httpServe defaultConfig (runReaderT app mvar)