diff --git a/Data/AhoCorasick.hs b/Data/AhoCorasick.hs new file mode 100644 --- /dev/null +++ b/Data/AhoCorasick.hs @@ -0,0 +1,187 @@ +{-# LANGUAGE DeriveGeneric #-} +module Data.AhoCorasick + ( ACMachine + , State(..) + , Match(..) + , construct + , constructWithValues + , root + , run + , step + , renderGraph + ) where + +import Control.Applicative +import Control.Monad +import Data.Hashable (Hashable) +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as Map +import Data.List +import Data.Maybe +import Data.Vector (Vector) +import qualified Data.Vector as V +import GHC.Generics (Generic) + + +data ACMachine a v = ACMachine (Goto a) Failure (Output v) + +type Goto a = Vector (HashMap a State) +type Failure = Vector State +type Output v = Vector [(Int, v)] + +type GotoMap a = HashMap State (HashMap a State) +type FailureMap = HashMap State State +type OutputMap v = HashMap State [(Int, v)] + +newtype State = State Int + deriving (Eq, Generic) +instance Hashable State + +data Match v = Match + { matchPos :: Int + , matchValue :: v + } deriving (Show) + +construct :: (Eq a, Hashable a) => [[a]] -> ACMachine a [a] +construct ps = ACMachine (toGotoArray n gotoMap) (toFailureArray n failureMap) (toOutputArray n outputMap) + where + (m, gotoMap) = buildGoto ps + n = m + 1 + failureMap = buildFailure gotoMap + outputMap = buildOutput pvs gotoMap failureMap + pvs = zip ps ps + +constructWithValues :: (Eq a, Hashable a) => [([a], v)] -> ACMachine a v +constructWithValues pvs = ACMachine (toGotoArray n gotoMap) (toFailureArray n failureMap) (toOutputArray n outputMap) + where + (m, gotoMap) = buildGoto ps + n = m + 1 + failureMap = buildFailure gotoMap + outputMap = buildOutput pvs gotoMap failureMap + ps = map fst pvs + +toGotoArray :: Int -> GotoMap a -> Goto a +toGotoArray n m = V.generate n (fromMaybe Map.empty . flip Map.lookup m . State) + +toFailureArray :: Int -> FailureMap -> Failure +toFailureArray n m = V.generate n (fromMaybe (error "failure: ") . flip Map.lookup m . State) + +toOutputArray :: Int -> OutputMap v -> Output v +toOutputArray n m = V.generate n (fromMaybe [] . flip Map.lookup m . State) + +root :: State +root = State 0 + +run :: (Eq a, Hashable a) => ACMachine a v -> [a] -> [Match v] +run acm = go root . zip [1..] + where + go _ [] = [] + go s ((i, x):ixs) = map toMatch vs ++ go s' ixs + where + toMatch (l, v) = Match { matchPos = i - l + 1, matchValue = v } + (s', vs) = step acm x s + +step :: (Eq a, Hashable a) => ACMachine a v -> a -> State -> (State, [(Int, v)]) +step (ACMachine g f o) x s = (s', output s') + where + s' = head $ mapMaybe (flip goto x) $ iterate failure s + goto (State 0) x' = (Map.lookup x' $ g V.! 0) <|> Just root + goto (State i) x' = Map.lookup x' $ g V.! i + failure (State i) = f V.! i + output (State i) = o V.! i + +buildOutput :: (Eq a, Hashable a) => [([a], v)] -> GotoMap a -> FailureMap -> OutputMap v +buildOutput pvs gotoMap failureMap = foldl' build o0 $ tail $ toBFList gotoMap + where + build o s = foldl' (\a (_, s') -> Map.insertWith (flip (++)) s' (lookupDefault [] (failure s') a) a) o ts + where + ts = Map.toList $ lookupDefault Map.empty s gotoMap + failure = fromMaybe (error "failure: ") . flip Map.lookup failureMap + o0 = Map.fromList $ fromJust $ mapM toKV pvs + toKV (p, v) = do + s <- finalState gotoMap root p + return (s, [(length p, v)]) + +finalState :: (Eq a, Hashable a) => GotoMap a -> State -> [a] -> Maybe State +finalState m = foldM (\s x -> Map.lookup s m >>= Map.lookup x) + +buildGoto :: (Eq a, Hashable a) => [[a]] -> (Int, GotoMap a) +buildGoto = foldl' (flip extend) (0, Map.empty) + +extend :: (Eq a, Hashable a) => [a] -> (Int, GotoMap a) -> (Int, GotoMap a) +extend = go root + where + go _ [] nm = nm + go s (x:xs) nm@(n, m) = case Map.lookup x sm of + Nothing -> go s' xs (n', m') + where + s' = State n' -- root is 0 + n' = n + 1 + sm' = Map.insert x s' sm + m' = Map.insert s sm' m + Just s' -> go s' xs nm + where + sm = fromMaybe Map.empty $ Map.lookup s m + +buildFailure :: (Eq a, Hashable a) => GotoMap a -> FailureMap +buildFailure m = foldl' build Map.empty $ toBFList m + where + build f s = foldl' (\a (x, s') -> Map.insert s' (failureState f s x) a) f ts + where + ts = Map.toList $ lookupDefault Map.empty s m + failureState _ (State 0) _ = root + failureState f s x = head $ mapMaybe (flip goto x) $ iterate failure (failure s) + where + failure = fromMaybe (error "failure: ") . flip Map.lookup f + goto (State 0) x = (Map.lookup root m >>= Map.lookup x) <|> Just root + goto s x = Map.lookup s m >>= Map.lookup x + +toBFList :: GotoMap a -> [State] +toBFList m = ss0 + where + ss0 = root : go 1 ss0 + go 0 _ = [] + go n (s:ss) = case Map.lookup s m of + Nothing -> go (n - 1) ss + Just sm -> children ++ go (n - 1 + Map.size sm) ss + where + children = Map.elems sm + go _ _ = error "toBFList: invalid state" + +toBFList' :: Goto a -> [State] +toBFList' v = ss0 + where + ss0 = root : go 1 ss0 + go 0 _ = [] + go n ((State i):ss) = children ++ go (n - 1 + Map.size sm) ss + where + sm = v V.! i + children = Map.elems sm + go _ _ = error "toBFList: invalid state" + +lookupDefault :: (Eq k, Hashable k) => v -> k -> HashMap k v -> v +lookupDefault def k m = fromMaybe def $ Map.lookup k m + +renderGraph :: ACMachine Char [Char] -> String +renderGraph (ACMachine g f o) = + graph "digraph" $ statements [ + attr "graph" [("rankdir", "LR")] + , statements $ map state (toBFList' g) + , statements $ map stateWithOutput $ filter (not . null . snd) $ zip (map State [0..]) (V.toList o) + , statements $ map (\s@(State i) -> statements $ map (uncurry $ transEdge s) $ Map.toList $ g V.! i) (toBFList' g) + , statements $ map (\s@(State i) -> failEdge s $ f V.! i) (tail $ toBFList' g) + ] + where + statements = intercalate " " + graph typ body = typ ++ " { " ++ body ++ " }" + attr typ attrList = typ ++ " " ++ "[" ++ intercalate "," (map kvStr attrList) ++ "];" + node nid attrList = nid ++ " " ++ "[" ++ intercalate "," (map kvStr attrList) ++ "];" + kvStr (k, v) = k ++ "=" ++ v + state s@(State 0) = node (stateID s) [("shape", "doublecircle")] + state s = node (stateID s) [("shape", "circle")] + stateWithOutput (s, xs) = node (stateID s) [("label", "<" ++ tableHTML (stateID s) ("{" ++ intercalate "," (map snd xs) ++ "}") ++ ">"), ("shape", "none")] + tableHTML row1 row2 = "
" ++ row1 ++ "
" ++ row2 ++ "
" + stateID (State 0) = "Root" + stateID (State n) = 'S' : show n + transEdge s x s' = stateID s ++ " -> " ++ stateID s' ++ " [label=\"" ++ [x] ++ "\"];" + failEdge s s' = stateID s ++ " -> " ++ stateID s' ++ " [style=dashed, constraint=false];" diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2013, Yuta Taniguchi + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Yuta Taniguchi nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/ac-machine.cabal b/ac-machine.cabal new file mode 100644 --- /dev/null +++ b/ac-machine.cabal @@ -0,0 +1,22 @@ +name: ac-machine +version: 0.2.0.1 +synopsis: Aho-Corasick string matching algorithm in Haskell +description: An implementation of the Aho-Corasick string matching algorithm written in Haskell. +license: BSD3 +license-file: LICENSE +author: Yuta Taniguchi +maintainer: Yuta Taniguchi +category: Data +build-type: Simple +extra-source-files: examples/Simple.hs +cabal-version: >=1.10 + +library + exposed-modules: Data.AhoCorasick + other-extensions: DeriveGeneric + build-depends: base >=4.6 && <4.7, hashable >=1.2 && <1.3, unordered-containers >=0.2 && <0.3, vector >=0.10 && <0.11 + default-language: Haskell2010 + +source-repository head + type: git + location: git://github.com/yuttie/ACMachine.git diff --git a/examples/Simple.hs b/examples/Simple.hs new file mode 100644 --- /dev/null +++ b/examples/Simple.hs @@ -0,0 +1,11 @@ +module Main where + +import Data.AhoCorasick + + +main :: IO () +main = do + let acm = construct ["he", "she", "his", "hers"] + print $ run acm "ushers" + let acm' = constructWithValues [("he", 1.2), ("she", 3.4), ("his", 5.6), ("hers", 7.8)] + print $ run acm' "ushers"