Main.hs

module Main(main) where

import Util

import System.Environment (getArgs)
import Control.Monad (forever)
import Data.List (intercalate, scanl1)

main :: IO ()
main = do
  args <- getArgs
  if null args
    then usage
    else concatFiles args >>= mainProc

usage :: IO ()
usage = putStrLn "usage: $ ./Main <input files>"

concatFiles :: [FilePath] -> IO (String, Lookup)
concatFiles fps = do
  list <- mapM readFile fps
  let filePos       = zip fps $ scanl1 (+) $ map ((+separatorLen).length) list
      lookupFilePos = conv filePos
  return $ (intercalate separator list, lookupFilePos)
  where
    separator    = "\0"
    separatorLen = length separator

mainProc :: (String, Lookup) -> IO ()
mainProc (text,idx2fpos) = forever loop
  where
    loop = getLine >>= print . idx2fpos . searchText text

------------------------------------------------------------
type Lookup = [Index] -> [(FilePath, Index)]

conv :: [(FilePath,Index)] -> Lookup
conv lst results = conv' (lst ++ [("nothing",maxBound)]) 0 results

conv' :: [(FilePath,Index)] -> Index -> Lookup
conv' [] _  _ = error "text overflow"
conv' _  _ [] = []
conv' poslist prev (idx:rest)
  | idx < len  = (fname, idx - prev) : conv' poslist prev rest
  | idx >= len = (fst $ head $ tail poslist, idx - len) : conv' (tail poslist) len rest
  where
    (fname,len) = head poslist

Util.hs

module Util
    (makeSuffixArray
    ,searchText
    ,Text,Index) where

import Data.List (sort, sortBy, tails, isPrefixOf)
import Data.Function (on)

type Text = String
type Index = Int

makeSuffixArray :: Text -> [Index]
makeSuffixArray = map fst . makeSuffixArray'

makeSuffixArray' :: Text -> [(Index,Text)]
makeSuffixArray' = sortBy (compare `on` snd) . zip [1..] . init . tails

searchText :: Text -> Text -> [Index]
searchText t subt = sort $ searchText' (makeSuffixArray' t) subt

searchText' :: [(Index,Text)] -> Text -> [Index]
searchText' []  _ = []
searchText' ary subt
  | suffix < subt          = right
  | suffix == subt         = index : right
  | isPrefixOf subt suffix = left ++ [index] ++ right
  | otherwise              = left
  where
    center = length ary `div` 2
    left   = searchText' left' subt
    right  = searchText' (tail right') subt
    (left',right') = splitAt center ary
    (index,suffix) = ary !! center