suffixtree-0.2: Data/SuffixTree.hs
{- Fastest when compiled as follows: ghc -O2 -optc-O3 -funbox-strict-fields -}
-----------------------------------------------------------------------------
-- |
-- Module : Data.SuffixTree
-- Copyright : (c) Bryan O'Sullivan 2007
-- License : BSD-style
-- Maintainer : bos@serpentine.com
-- Stability : experimental
-- Portability : portable
--
-- A lazy, efficient suffix tree implementation.
--
-- Since many function names (but not the type name) clash with
-- "Prelude" names, this module is usually imported @qualified@, e.g.
--
-- > import Data.SuffixTree (STree)
-- > import qualified Data.SuffixTree as T
--
-- The implementation is based on the first of those described in the
-- following paper:
--
-- * Robert Giegerich and Stefan Kurtz, \"/A comparison of
-- imperative and purely functional suffix tree constructions/\",
-- Science of Computer Programming 25(2-3):187-218, 1995,
-- <http://citeseer.ist.psu.edu/giegerich95comparison.html>
--
-- This implementation constructs the suffix tree lazily, so subtrees
-- are not created until they are traversed. Two construction
-- functions are provided, 'constructWith' for sequences composed of
-- small alphabets, and 'construct' for larger alphabets.
--
-- Estimates are given for performance. The value /k/ is a constant;
-- /n/ is the length of a query string; and /t/ is the number of
-- elements (nodes and leaves) in a suffix tree.
module Data.SuffixTree
(
-- * Types
Alphabet
, Edge
, Prefix
, STree(..)
-- * Construction
, constructWith
, construct
-- * Querying
, elem
, findEdge
, findTree
, findPath
, countLeaves
, countRepeats
-- * Traversal
, foldr
, foldl
, fold
-- * Other useful functions
, mkPrefix
, prefix
, suffixes
) where
import Prelude hiding (elem, foldl, foldr)
import qualified Data.Map as M
import Control.Arrow (second)
import qualified Data.ByteString as SB
import qualified Data.ByteString.Lazy as LB
import qualified Data.List as L
import Data.Maybe (listToMaybe, mapMaybe)
-- | The length of a prefix list. This type is formulated to do cheap
-- work eagerly (to avoid constructing a pile of deferred thunks),
-- while deferring potentially expensive work (computing the length of
-- a list).
data Length a = Exactly {-# UNPACK #-} !Int
| Sum {-# UNPACK #-} !Int [a]
deriving (Show)
-- | The list of symbols that 'constructWith' can possibly see in its
-- input.
type Alphabet a = [a]
-- | The prefix string associated with an 'Edge'. Use 'mkPrefix' to
-- create a value of this type, and 'prefix' to deconstruct one.
newtype Prefix a = Prefix ([a], Length a)
instance (Eq a) => Eq (Prefix a) where
a == b = prefix a == prefix b
instance (Ord a) => Ord (Prefix a) where
compare a b = compare (prefix a) (prefix b)
instance (Show a) => Show (Prefix a) where
show a = "mkPrefix " ++ show (prefix a)
type EdgeFunction a = [[a]] -> (Length a, [[a]])
-- | An edge in the suffix tree.
type Edge a = (Prefix a, STree a)
-- | /O(1)/. Construct a 'Prefix' value.
mkPrefix :: [a] -> Prefix a
mkPrefix xs = Prefix (xs, Sum 0 xs)
pmap :: (a -> b) -> Prefix a -> Prefix b
pmap f = mkPrefix . map f . prefix
instance Functor Prefix where
fmap = pmap
-- | The suffix tree type. The implementation is exposed to ease the
-- development of custom traversal functions. Note that @('Prefix' a,
-- 'STree' a)@ pairs are not stored in any order.
data STree a = Node [Edge a]
| Leaf
deriving (Show)
smap :: (a -> b) -> STree a -> STree b
smap _ Leaf = Leaf
smap f (Node es) = Node (map (\(p, t) -> (fmap f p, smap f t)) es)
instance Functor STree where
fmap = smap
-- | /O(n)/. Obtain the list stored in a 'Prefix'.
prefix :: Prefix a -> [a]
prefix (Prefix (ys, Exactly n)) = take n ys
prefix (Prefix (ys, Sum n xs)) = tk n ys
where tk 0 ys = zipWith (const id) xs ys
tk n (y:ys) = y : tk (n-1) ys
-- | /O(t)/. Folds the edges in a tree, using post-order traversal.
-- Suitable for lazy use.
foldr :: (Prefix a -> b -> b) -> b -> STree a -> b
foldr _ z Leaf = z
foldr f z (Node es) = L.foldr (\(p,t) v -> f p (foldr f v t)) z es
-- | /O(t)/. Folds the edges in a tree, using pre-order traversal. The
-- step function is evaluated strictly.
foldl :: (a -> Prefix b -> a) -- ^ step function (evaluated strictly)
-> a -- ^ initial state
-> STree b
-> a
foldl _ z Leaf = z
foldl f z (Node es) = L.foldl' (\v (p,t) -> f (foldl f v t) p) z es
-- | /O(t)/. Generic fold over a tree.
--
-- A few simple examples.
--
-- >countLeaves == fold id id (const const) (1+) 0
--
-- >countEdges = fold id id (\_ a _ -> a+1) id 0
--
-- This more complicated example generates a tree of the same shape,
-- but new type, with annotated leaves.
--
-- @
--data GenTree a b = GenNode [('Prefix' a, GenTree a b)]
-- | GenLeaf b
-- deriving ('Show')
-- @
--
-- @
--gentree :: 'STree' a -> GenTree a Int
--gentree = 'fold' reset id fprefix reset leaf
-- where leaf = GenLeaf 1
-- reset = 'const' leaf
-- fprefix p t (GenLeaf _) = GenNode [(p, t)]
-- fprefix p t (GenNode es) = GenNode ((p, t):es)
-- @
fold :: (a -> a) -- ^ downwards state transformer
-> (a -> a) -- ^ upwards state transformer
-> (Prefix b -> a -> a -> a) -- ^ edge state transformer
-> (a -> a) -- ^ leaf state transformer
-> a -- ^ initial state
-> STree b -- ^ tree
-> a
fold fdown fup fprefix fleaf = go
where go v Leaf = fleaf v
go v (Node es) = fup (L.foldr edge v es)
edge (p, t) v = fprefix p (go (fdown v) t) v
-- | Increments the length of a prefix.
inc :: Length a -> Length a
inc (Exactly n) = Exactly (n+1)
inc (Sum n xs) = Sum (n+1) xs
lazyTreeWith :: (Eq a) => EdgeFunction a -> Alphabet a -> [a] -> STree a
lazyTreeWith edge alphabet = suf . suffixes
where suf [[]] = Leaf
suf ss = Node [(Prefix (a:sa, inc cpl), suf ssr)
| a <- alphabet,
n@(sa:_) <- [ss `clusterBy` a],
(cpl,ssr) <- [edge n]]
clusterBy ss a = [cs | c:cs <- ss, c == a]
-- | /O(n)/. Returns all non-empty suffixes of the argument, longest
-- first. Behaves as follows:
--
-- >suffixes xs == init (tails xs)
suffixes :: [a] -> [[a]]
suffixes xs@(_:xs') = xs : suffixes xs'
suffixes _ = []
lazyTree :: (Ord a) => EdgeFunction a -> [a] -> STree a
lazyTree edge = suf . suffixes
where suf [[]] = Leaf
suf ss = Node [(Prefix (a:sa, inc cpl), suf ssr)
| (a, n@(sa:_)) <- suffixMap ss,
(cpl,ssr) <- [edge n]]
suffixMap :: Ord a => [[a]] -> [(a, [[a]])]
suffixMap = map (second reverse) . M.toList . L.foldl' step M.empty
where step m (x:xs) = M.alter (f xs) x m
step m _ = m
f x Nothing = Just [x]
f x (Just xs) = Just (x:xs)
cst :: Eq a => EdgeFunction a
cst [s] = (Sum 0 s, [[]])
cst awss@((a:w):ss)
| null [c | c:_ <- ss, a /= c] = let cpl' = inc cpl
in seq cpl' (cpl', rss)
| otherwise = (Exactly 0, awss)
where (cpl, rss) = cst (w:[u | _:u <- ss])
pst :: Eq a => EdgeFunction a
pst = g . dropNested
where g [s] = (Sum 0 s, [[]])
g ss = (Exactly 0, ss)
dropNested ss@[_] = ss
dropNested awss@((a:w):ss)
| null [c | c:_ <- ss, a /= c] = [a:s | s <- rss]
| otherwise = awss
where rss = dropNested (w:[u | _:u <- ss])
{-# SPECIALISE constructWith :: [Char] -> [Char] -> STree Char #-}
{-# SPECIALISE constructWith :: [[Char]] -> [[Char]] -> STree [Char] #-}
{-# SPECIALISE constructWith :: [SB.ByteString] -> [SB.ByteString]
-> STree SB.ByteString #-}
{-# SPECIALISE constructWith :: [LB.ByteString] -> [LB.ByteString]
-> STree LB.ByteString #-}
{-# SPECIALISE constructWith :: (Eq a) => [[a]] -> [[a]] -> STree [a] #-}
-- | /O(k n log n)/. Constructs a suffix tree using the given
-- alphabet. The performance of this function is linear in the size
-- /k/ of the alphabet. That makes this function suitable for small
-- alphabets, such as DNA nucleotides. For an alphabet containing
-- more than a few symbols, 'construct' is usually several orders of
-- magnitude faster.
constructWith :: (Eq a) => Alphabet a -> [a] -> STree a
constructWith = lazyTreeWith cst
{-# SPECIALISE construct :: [Char] -> STree Char #-}
{-# SPECIALISE construct :: [[Char]] -> STree [Char] #-}
{-# SPECIALISE construct :: [SB.ByteString] -> STree SB.ByteString #-}
{-# SPECIALISE construct :: [LB.ByteString] -> STree LB.ByteString #-}
{-# SPECIALISE construct :: (Ord a) => [[a]] -> STree [a] #-}
-- | /O(n log n)/. Constructs a suffix tree.
construct :: (Ord a) => [a] -> STree a
construct = lazyTree cst
suffix :: (Eq a) => [a] -> [a] -> Maybe [a]
suffix (p:ps) (x:xs) | p == x = suffix ps xs
| otherwise = Nothing
suffix _ xs = Just xs
{-# SPECIALISE elem :: [Char] -> STree Char -> Bool #-}
{-# SPECIALISE elem :: [[Char]] -> STree [Char] -> Bool #-}
{-# SPECIALISE elem :: [SB.ByteString] -> STree SB.ByteString -> Bool #-}
{-# SPECIALISE elem :: [LB.ByteString] -> STree LB.ByteString -> Bool #-}
{-# SPECIALISE elem :: (Eq a) => [[a]] -> STree [a] -> Bool #-}
-- | /O(n)/. Indicates whether the suffix tree contains the given
-- sublist. Performance is linear in the length /n/ of the
-- sublist.
elem :: (Eq a) => [a] -> STree a -> Bool
elem [] _ = True
elem _ Leaf = False
elem xs (Node es) = any pfx es
where pfx (e, t) = maybe False (`elem` t) (suffix (prefix e) xs)
{-# SPECIALISE findEdge :: [Char] -> STree Char
-> Maybe (Edge Char, Int) #-}
{-# SPECIALISE findEdge :: [String] -> STree String
-> Maybe (Edge String, Int) #-}
{-# SPECIALISE findEdge :: [SB.ByteString] -> STree SB.ByteString
-> Maybe (Edge SB.ByteString, Int) #-}
{-# SPECIALISE findEdge :: [ LB.ByteString] -> STree LB.ByteString
-> Maybe (Edge LB.ByteString, Int) #-}
{-# SPECIALISE findEdge :: (Eq a) => [[a]] -> STree [a]
-> Maybe (Edge [a], Int) #-}
-- | /O(n)/. Finds the given subsequence in the suffix tree. On
-- failure, returns 'Nothing'. On success, returns the 'Edge' in the
-- suffix tree at which the subsequence ends, along with the number of
-- elements to drop from the prefix of the 'Edge' to get the \"real\"
-- remaining prefix.
--
-- Here is an example:
--
-- >> find "ssip" (construct "mississippi")
-- >Just ((mkPrefix "ppi",Leaf),1)
--
-- This indicates that the edge @('mkPrefix' \"ppi\",'Leaf')@ matches,
-- and that we must strip 1 character from the string @\"ppi\"@ to get
-- the remaining prefix string @\"pi\"@.
--
-- Performance is linear in the length /n/ of the query list.
findEdge :: (Eq a) => [a] -> STree a -> Maybe (Edge a, Int)
findEdge _ Leaf = Nothing
findEdge xs (Node es) = listToMaybe (mapMaybe pfx es)
where pfx e@(p, t) = let p' = prefix p
in suffix p' xs >>= \suf ->
case suf of
[] -> return (e, length (zipWith const xs p'))
s -> findEdge s t
-- | /O(n)/. Finds the subtree rooted at the end of the given query
-- sequence. On failure, returns 'Nothing'.
--
-- Performance is linear in the length /n/ of the query list.
findTree :: (Eq a) => [a] -> STree a -> Maybe (STree a)
findTree s t = (snd . fst) `fmap` findEdge s t
-- | /O(n)/. Returns the path from the 'Edge' in the suffix tree at
-- which the given subsequence ends, up to the root of the tree. If
-- the subsequence is not found, returns the empty list.
--
-- Performance is linear in the length of the query list.
findPath :: (Eq a) => [a] -> STree a -> [Edge a]
findPath = go []
where go _ _ Leaf = []
go me xs (Node es) = pfx me es
where pfx _ [] = []
pfx me (e@(p, t):es) =
case suffix (prefix p) xs of
Nothing -> pfx me es
Just [] -> e:me
Just s -> go (e:me) s t
-- | /O(t)/. Count the number of leaves in a tree.
--
-- Performance is linear in the number /t/ of elements in the tree.
countLeaves :: STree a -> Int
countLeaves Leaf = 1
countLeaves (Node es) = L.foldl' (\v (_, t) -> countLeaves t + v) 0 es
-- | /O(n + r)/. Count the number of times a sequence is repeated
-- in the input sequence that was used to construct the suffix tree.
--
-- Performance is linear in the length /n/ of the input sequence, plus
-- the number of times /r/ the sequence is repeated.
countRepeats :: (Eq a) => [a] -> STree a -> Int
countRepeats s t = maybe 0 countLeaves (findTree s t)