markov.hs

#!/usr/bin/env stack runhaskell --package array --package bytestring --package containers --package time
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PartialTypeSignatures #-}

import           Control.Monad

import           Data.Array.IArray
import           Data.Array (Array)
import           Data.Array.Unboxed (UArray)

import           Data.ByteString (ByteString)
import           Data.ByteString.Internal (c2w)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL

import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import           Data.List.Extra
import           Data.Maybe
import           Data.Word
import           Data.Time.Clock.System


data Trie c a = Trie
    { trieValue :: Maybe a
    , trieTrans :: (Map c (Trie c a))
    } deriving Show

emptyTrie :: Trie c a
emptyTrie = Trie Nothing M.empty

trieAppend :: Ord c => [c] -> a -> Trie c a -> Trie c a
trieAppend [] a (Trie _ m) = Trie (Just a) m
trieAppend (c:cs) a (Trie b m) = Trie b (M.alter append c m) where
    append t = Just $ trieAppend cs a (maybe emptyTrie id t)

trieNodeCount :: Trie i a -> Int
trieNodeCount (Trie f m) = 1 + M.foldr ((+) . trieNodeCount) 0 m


type AhoStateIx = Int

initStateIx :: AhoStateIx
initStateIx = 0

invalidStateIx :: AhoStateIx
invalidStateIx = -1

data Aho c a = Aho
    { ahoValue :: Array AhoStateIx (Maybe a)
    , ahoTrans :: UArray (AhoStateIx,c) AhoStateIx
    , ahoSuffix :: UArray AhoStateIx AhoStateIx
    , ahoFSuffix :: UArray AhoStateIx AhoStateIx
    } deriving (Show, Functor)

ahoFromTrie :: (Ix c, Enum c, Bounded c) => Trie c a -> Aho c a
ahoFromTrie trie = Aho value (toUArray trans) (toUArray suffix) (toUArray fSuffix) where
    nc = trieNodeCount trie

    bfs :: AhoStateIx -> [(AhoStateIx, c, Trie c a)] -> [(AhoStateIx, AhoStateIx, c, Maybe a)]
    bfs _ [] = []
    bfs curIx ((prevIx, c, (Trie a m)):ts) = (curIx, prevIx, c, a) : bfs (curIx+1) (ts ++ M.foldMapWithKey (\c t -> [(curIx, c, t)]) m)

    trieTrans = tail $ bfs initStateIx [(invalidStateIx, minBound, trie)]
    (rootTrans, otherTrans) = span (\(_, prevIx, _, _) -> prevIx == initStateIx) trieTrans

    assocMap :: (IArray a e, Ix i) => (i -> e -> e) -> a i e -> a i e
    assocMap f a = array (bounds a) $ map (\(i, e) -> (i, f i e)) $ assocs a

    toUArray :: (IArray a1 e, IArray a2 e, Ix i) => a1 i e -> a2 i e
    toUArray a = array (bounds a) (assocs a)

    value :: Array _ _
    value = array (0, nc-1) $ [(initStateIx, trieValue trie)] ++ map (\(curIx, _, _, a) -> (curIx, a)) trieTrans

    ahoTrieTrans :: Array _ _
    ahoTrieTrans = array ((0,minBound), (nc-1,maxBound)) [((ix,c), invalidStateIx) | ix <- [0 .. nc-1], c <- [minBound .. maxBound]]
        // map (\(curIx, prevIx, c, _) -> ((prevIx,c),curIx)) trieTrans

    trans = assocMap transF ahoTrieTrans where
        transF (prevIx,c) curIx | curIx /= invalidStateIx = curIx
        transF (prevIx,c) _     | prevIx == initStateIx && ahoTrieTrans ! (initStateIx,c) == invalidStateIx = initStateIx
        transF (prevIx,c) _     | prevIx == initStateIx = ahoTrieTrans ! (initStateIx,c)
        transF (prevIx,c) _     = trans ! (suffix ! prevIx, c)

    suffix :: Array _ _
    suffix = array (0, nc-1) $ [(initStateIx,invalidStateIx)] ++
        map (\(curIx, _, _, _) -> (curIx, initStateIx)) rootTrans ++
        map (\(curIx, prevIx, c, _) -> (curIx, trans ! (suffix ! prevIx, c))) otherTrans

    fSuffix :: Array _ _
    fSuffix = array (0, nc-1) $ [(initStateIx, invalidStateIx)] ++
        map (\(curIx, _, _, _) -> (curIx, if isJust (value ! (suffix ! curIx)) then suffix ! curIx else fSuffix ! (suffix ! curIx))) trieTrans

ahoNextState :: Ix c => Aho c a -> AhoStateIx -> c -> AhoStateIx
ahoNextState aho ix c = ahoTrans aho ! (ix,c)

ahoStateValues :: Aho c a -> AhoStateIx -> [a]
ahoStateValues aho ix = h ++ t (ahoFSuffix aho ! ix) where
    h = case ahoValue aho ! ix of
        Nothing -> []
        Just a -> [a]
    t ix | ix == invalidStateIx = []
    t ix = fromJust (ahoValue aho ! ix) : t (ahoFSuffix aho ! ix)


type Str = [Word8]

str :: Str -> ByteString
str = BS.pack

type Subst = (Int, Str)

getSubsts :: Trie Word8 Subst -> IO (Trie Word8 Subst)
getSubsts trie = do
    line <- BS.getLine
    case BS.null line of
        True -> return trie
        False ->
            let (search, repl) = BS.break (== c2w '/') line
                trie' = trieAppend (BS.unpack search) (BS.length search, BS.unpack $ BS.tail repl) trie
            in getSubsts trie'

runSubst :: Aho Word8 Subst -> [AhoStateIx] -> (Str, Str) -> [ByteString]
runSubst aho ixs@(ix:_) (before, after) | notNull (ahoStateValues aho ix) =
    let ((len,sub):_) = ahoStateValues aho ix
        before1 = drop len before
    in str (reverse before1) <> str sub <> str after : runSubst aho (drop len ixs) (before1, sub ++ after)
runSubst aho ixs@(ix:_) (before, after) | null after = []
runSubst aho ixs@(ix:_) (before, after@(c : after1)) = runSubst aho (ahoNextState aho ix c : ixs) (c : before, after1)

seconds :: SystemTime -> Double
seconds (MkSystemTime s ns) = (fromInteger $ toInteger s) + (fromInteger $ toInteger ns) * 1e-9

main = do
    trie <- getSubsts emptyTrie
    init <- BS.unpack <$> BS.getLine
    let aho = ahoFromTrie trie
    let subs@(fst:_) = runSubst aho [initStateIx] ([], init)
    start <- seconds <$> (fst `seq` getSystemTime)
    BS.putStrLn $ str init
    mapM_ BS.putStrLn $ filter ((== "PG") . BS.take 2) subs
    end <- seconds <$> getSystemTime
    print $ end - start