{-# LANGUAGE ScopedTypeVariables #-}
{-|
== Suffix tree

A moderately simple and flexible suffix tree.
A suffix tree of a string is a compressed trie of all the suffixes of the string, useful for fast
matching on substrings of the string or to calculate certain properties of the string.
Other suffix structures such as suffix arrays and suffix automata may serve as alternates to suffix
trees.

The implementation here constructs a SuffixTree a from a given string index function, with string
elements of type Int. Ukkonen's algorithm is used for construction. This is not lazy, the entire
tree is constructed right away.
Only the implicit suffix tree is constructed, which is a suffix tree where suffixes that are
prefixes of other suffixes do not end in leaves. To make them end in leaves, set the last element
of the string to a unique value.
A SuffixTree a also stores accumulated values of type a at every node. These values are calculated
using user supplied functions, which can be chosen based on what the tree will be used for.

Sources:

* Esko Ukkonen, "On–line construction of suffix trees", 1995
  https://www.cs.helsinki.fi/u/ukkonen/SuffixT1withFigs.pdf
* Dan Gusfield, "Algorithms on Strings, Trees, and Sequences", 1997
  https://doi.org/10.1017/CBO9780511574931

String indexing is assumed to take O(1). Let k be the alphabet size. Let the complexity of IntMap
operations be f(n), where n is the size of the map. f(n) is O(min(n, word size)), see IntMap
documentation for details.

-}

{-
Implementation notes:
* Ukkonen's algorithm is used to construct the suffix tree. The implementation here is similar to
  the one in his paper. The state maintained is (u, pos). For index i the substring [pos..i] must be
  inserted at node u, then at the suffix link of u, and so on.
* buildSufT builds a nicer-to-use tree on top of arrays calculated by Ukkonen's algorithm.
* Suffix links, which can be useful in some situations, are not retained for simplicity.
* There is no setup for fast LCA queries, which are also useful in some situations.
* matchSufT can be easily modified to return the value for the longest prefix of the pattern
  matched, if it does not match completely.
-}

module SuffixTree
    ( SufTNode(..)
    , SufTEdge(..)
    , Chr
    , buildSufT
    , matchSufT
    , buildMatchSufT
    , drawSufT
    ) where

import Control.DeepSeq
import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed
import Data.Array.Unsafe
import Data.List
import qualified Data.IntMap.Strict as IM

import TreeDraw ( draw )

type Chr = Int

data SufTNode a = SufTNode !a !(IM.IntMap (SufTEdge a))
data SufTEdge a = SufTEdge !Int !Int !(SufTNode a)

-- | Builds a suffix tree. n is the length of the string. at is a 0-based indexing function into the
-- string. fromLeaf constructs a value from a leaf index. updEdge constructs a value from an existing
-- value and a length of an edge leading to it. merge combines two values.
-- Examples:
--   To query the number of times a pattern occurs in the string, use const 1, const, and (+).
--   To calculate the number of distinct substrings, use const 0, (+), and (+).
-- O(n * f(k)), plus the above functions are each called O(n) times.
buildSufT :: (Int -> a) -> (a -> Int -> a) -> (a -> a -> a)
          -> Int -> (Int -> Chr)
          -> SufTNode a
buildSufT :: (Int -> a)
-> (a -> Int -> a)
-> (a -> a -> a)
-> Int
-> (Int -> Int)
-> SufTNode a
buildSufT Int -> a
fromLeaf a -> Int -> a
updEdge a -> a -> a
merge Int
n Int -> Int
at = Int -> Int -> SufTNode a
mkNode Int
n Int
0 where
    (Array Int (IntMap Int)
nxta, UArray Int Int
lefta, UArray Int Int
lena) = Int
-> (Int -> Int)
-> (Array Int (IntMap Int), UArray Int Int, UArray Int Int)
ukkonen Int
n Int -> Int
at
    mkNode :: Int -> Int -> SufTNode a
mkNode Int
dep Int
i = a -> IntMap (SufTEdge a) -> SufTNode a
forall a. a -> IntMap (SufTEdge a) -> SufTNode a
SufTNode a
a IntMap (SufTEdge a)
nxt where
        nxt :: IntMap (SufTEdge a)
nxt = (Int -> SufTEdge a) -> IntMap Int -> IntMap (SufTEdge a)
forall a b. (a -> b) -> IntMap a -> IntMap b
IM.map (\Int
j -> Int -> Int -> SufTNode a -> SufTEdge a
forall a. Int -> Int -> SufTNode a -> SufTEdge a
SufTEdge (UArray Int Int
leftaUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
j) (UArray Int Int
lenaUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
j) (Int -> Int -> SufTNode a
mkNode (Int
dep Int -> Int -> Int
forall a. Num a => a -> a -> a
- UArray Int Int
lenaUArray Int Int -> Int -> Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
j) Int
j)) (Array Int (IntMap Int)
nxtaArray Int (IntMap Int) -> Int -> IntMap Int
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
!Int
i)
        a :: a
a | IntMap (SufTEdge a) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (SufTEdge a)
nxt = Int -> a
fromLeaf Int
dep
          | Bool
otherwise   = (a -> a -> a) -> [a] -> a
forall a. (a -> a -> a) -> [a] -> a
foldl1' a -> a -> a
merge [a -> Int -> a
updEdge a
a' Int
len | SufTEdge Int
_ Int
len (SufTNode a
a' IntMap (SufTEdge a)
_) <- IntMap (SufTEdge a) -> [SufTEdge a]
forall a. IntMap a -> [a]
IM.elems IntMap (SufTEdge a)
nxt]

ukkonen :: Int -> (Int -> Chr) -> (Array Int (IM.IntMap Int), UArray Int Int, UArray Int Int)
ukkonen :: Int
-> (Int -> Int)
-> (Array Int (IntMap Int), UArray Int Int, UArray Int Int)
ukkonen Int
n Int -> Int
at = (forall s.
 ST s (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> (Array Int (IntMap Int), UArray Int Int, UArray Int Int)
forall a. (forall s. ST s a) -> a
runST ((forall s.
  ST s (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
 -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> (forall s.
    ST s (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> (Array Int (IntMap Int), UArray Int Int, UArray Int Int)
forall a b. (a -> b) -> a -> b
$ do
    let sz :: Int
sz = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    STArray s Int (IntMap Int)
nxt :: STArray s Int (IM.IntMap Int) <- (Int, Int) -> IntMap Int -> ST s (STArray s Int (IntMap Int))
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) IntMap Int
forall a. IntMap a
IM.empty
    [STUArray s Int Int
suf, STUArray s Int Int
left, STUArray s Int Int
len] :: [STUArray s Int Int] <- Int -> ST s (STUArray s Int Int) -> ST s [STUArray s Int Int]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
3 ((Int, Int) -> Int -> ST s (STUArray s Int Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray (Int
0, Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
0)
    STUArray s () Int
cur :: STUArray s () Int <- ((), ()) -> Int -> ST s (STUArray s () Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray ((), ()) Int
1
    let root :: Int
root = Int
0
        nxtId :: ST s Int
nxtId = STUArray s () Int -> () -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s () Int
cur () ST s Int -> (Int -> ST s Int) -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
i -> Int
i Int -> ST s () -> ST s Int
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ STUArray s () Int -> () -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s () Int
cur () (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        step :: Int -> Int -> Int -> ST s (Int, Int)
step Int
i = Int -> Int -> Int -> ST s (Int, Int)
go Int
root where
            go :: Int -> Int -> Int -> ST s (Int, Int)
go Int
prv Int
pos Int
u = do
                IntMap Int
nxtu <- STArray s Int (IntMap Int) -> Int -> ST s (IntMap Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Int (IntMap Int)
nxt Int
u
                case IntMap Int
nxtu IntMap Int -> Int -> Maybe Int
forall a. IntMap a -> Int -> Maybe a
IM.!? Int -> Int
at Int
pos of
                    Maybe Int
Nothing -> Int -> ST s (Int, Int)
insLeafGo Int
u
                    Just Int
v  -> ST s (ST s (Int, Int)) -> ST s (Int, Int)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (ST s (ST s (Int, Int)) -> ST s (Int, Int))
-> ST s (ST s (Int, Int)) -> ST s (Int, Int)
forall a b. (a -> b) -> a -> b
$ IntMap Int -> Int -> Int -> Int -> ST s (Int, Int)
tryEdge IntMap Int
nxtu Int
v (Int -> Int -> ST s (Int, Int))
-> ST s Int -> ST s (Int -> ST s (Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
left Int
v ST s (Int -> ST s (Int, Int)) -> ST s Int -> ST s (ST s (Int, Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
len Int
v
              where
                tryEdge :: IntMap Int -> Int -> Int -> Int -> ST s (Int, Int)
tryEdge IntMap Int
nxtu Int
v Int
leftv Int
lenv
                    | Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
i = Int -> Int -> Int -> ST s (Int, Int)
go Int
prv (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenv) Int
v
                    | Int -> Int
at Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> Int
at Int
j    = ST s Int
doSplit ST s Int -> (Int -> ST s (Int, Int)) -> ST s (Int, Int)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> ST s (Int, Int)
insLeafGo
                    | Int
pos Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
leftv    = Int -> ST s (Int, Int)
goSuf Int
u
                    | Bool
otherwise       = (Int
u, Int
pos) (Int, Int) -> ST s () -> ST s (Int, Int)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> ST s ()
setSufPrv Int
u
                  where
                    lenw :: Int
lenw = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
pos
                    j :: Int
j = Int
leftv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenw
                    doSplit :: ST s Int
doSplit = do
                        Int
w <- ST s Int
nxtId
                        STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
left Int
w Int
leftv
                        STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
len Int
w Int
lenw
                        STArray s Int (IntMap Int) -> Int -> IntMap Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int (IntMap Int)
nxt Int
w (IntMap Int -> ST s ()) -> IntMap Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> IntMap Int
forall a. Int -> a -> IntMap a
IM.singleton (Int -> Int
at Int
j) Int
v
                        STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
left Int
v (Int
leftv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lenw)
                        STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
len Int
v (Int
lenv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
lenw)
                        STArray s Int (IntMap Int) -> Int -> IntMap Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int (IntMap Int)
nxt Int
u (IntMap Int -> ST s ()) -> IntMap Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Int -> Int
at Int
pos) Int
w IntMap Int
nxtu
                        Int -> ST s Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
w
                insLeafGo :: Int -> ST s (Int, Int)
insLeafGo Int
v = Int -> ST s ()
setSufPrv Int
v ST s () -> ST s () -> ST s ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
insLeaf Int
v ST s () -> ST s (Int, Int) -> ST s (Int, Int)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s (Int, Int)
goSuf Int
v
                setSufPrv :: Int -> ST s ()
setSufPrv Int
v = STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
suf Int
prv Int
v :: ST s ()
                insLeaf :: Int -> ST s ()
insLeaf Int
v = do
                    Int
w <- ST s Int
nxtId
                    STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
left Int
w Int
i
                    STUArray s Int Int -> Int -> Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STUArray s Int Int
len Int
w (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)
                    IntMap Int
nxtv <- STArray s Int (IntMap Int) -> Int -> ST s (IntMap Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Int (IntMap Int)
nxt Int
v
                    STArray s Int (IntMap Int) -> Int -> IntMap Int -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int (IntMap Int)
nxt Int
v (IntMap Int -> ST s ()) -> IntMap Int -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> Int -> IntMap Int -> IntMap Int
forall a. Int -> a -> IntMap a -> IntMap a
IM.insert (Int -> Int
at Int
i) Int
w IntMap Int
nxtv
                goSuf :: Int -> ST s (Int, Int)
goSuf Int
v
                    | Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
root = STUArray s Int Int -> Int -> ST s Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STUArray s Int Int
suf Int
u ST s Int -> (Int -> ST s (Int, Int)) -> ST s (Int, Int)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Int -> Int -> ST s (Int, Int)
go Int
v Int
pos
                    | Int
pos Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i   = Int -> Int -> Int -> ST s (Int, Int)
go Int
v (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
root
                    | Bool
otherwise = (Int, Int) -> ST s (Int, Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
root, Int
pos)
    ((Int, Int) -> Int -> ST s (Int, Int))
-> (Int, Int) -> [Int] -> ST s ()
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m ()
foldM_ (\(Int
u, Int
pos) Int
i -> Int -> Int -> Int -> ST s (Int, Int)
step Int
i Int
pos Int
u) (Int
root, Int
0) [Int
0 .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
    (,,) (Array Int (IntMap Int)
 -> UArray Int Int
 -> UArray Int Int
 -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> ST s (Array Int (IntMap Int))
-> ST
     s
     (UArray Int Int
      -> UArray Int Int
      -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STArray s Int (IntMap Int) -> ST s (Array Int (IntMap Int))
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STArray s Int (IntMap Int)
nxt ST
  s
  (UArray Int Int
   -> UArray Int Int
   -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> ST s (UArray Int Int)
-> ST
     s
     (UArray Int Int
      -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STUArray s Int Int -> ST s (UArray Int Int)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STUArray s Int Int
left ST
  s
  (UArray Int Int
   -> (Array Int (IntMap Int), UArray Int Int, UArray Int Int))
-> ST s (UArray Int Int)
-> ST s (Array Int (IntMap Int), UArray Int Int, UArray Int Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> STUArray s Int Int -> ST s (UArray Int Int)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze STUArray s Int Int
len

-- | Matches a pattern on the suffix tree. updEdge and at are as defined for buildSufT. m is the length
-- of the pattern. at' is a 0-based indexing function into the pattern. If the pattern is not present,
-- returns Nothing. Otherwise returns the accumulated value for the subtree where the pattern ends.
-- O(m * f(k)) plus at most one call of updEdge.
matchSufT :: (a -> Int -> a) -> (Int -> Chr) -> SufTNode a
          -> Int -> (Int -> Chr)
          -> Maybe a
matchSufT :: (a -> Int -> a)
-> (Int -> Int) -> SufTNode a -> Int -> (Int -> Int) -> Maybe a
matchSufT a -> Int -> a
updEdge Int -> Int
at (SufTNode a
a IntMap (SufTEdge a)
nxt) Int
m Int -> Int
at' = if Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then a -> Maybe a
forall a. a -> Maybe a
Just a
a else IntMap (SufTEdge a) -> Int -> Maybe a
go IntMap (SufTEdge a)
nxt Int
0 where
    go :: IntMap (SufTEdge a) -> Int -> Maybe a
go IntMap (SufTEdge a)
nxt Int
i = Int -> IntMap (SufTEdge a) -> Maybe (SufTEdge a)
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Int -> Int
at' Int
i) IntMap (SufTEdge a)
nxt Maybe (SufTEdge a) -> (SufTEdge a -> Maybe a) -> Maybe a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SufTEdge a -> Maybe a
go' where
        go' :: SufTEdge a -> Maybe a
go' (SufTEdge Int
left Int
len (SufTNode a
a IntMap (SufTEdge a)
nxt'))
            | Int
i' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = a -> Maybe a
forall a. a -> Maybe a
Just a
a
            | Int
i' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m             = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Int -> a
updEdge a
a (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
d))
            | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len            = IntMap (SufTEdge a) -> Int -> Maybe a
go IntMap (SufTEdge a)
nxt' Int
i'
            | Bool
otherwise           = Maybe a
forall a. Maybe a
Nothing
          where
            d :: Int
d = Int -> (Int -> Int) -> (Int -> Int) -> Int
forall t a. (Eq t, Eq a, Num t) => t -> (t -> a) -> (t -> a) -> t
commonPrefix (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
len (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i)) (Int -> Int
at (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
left)) (Int -> Int
at' (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i))
            i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d
    commonPrefix :: t -> (t -> a) -> (t -> a) -> t
commonPrefix t
n t -> a
f t -> a
g = (t -> Bool) -> (t -> t) -> t -> t
forall a. (a -> Bool) -> (a -> a) -> a -> a
until (\t
i -> t
i t -> t -> Bool
forall a. Eq a => a -> a -> Bool
== t
n Bool -> Bool -> Bool
|| t -> a
f t
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= t -> a
g t
i) (t -> t -> t
forall a. Num a => a -> a -> a
+t
1) t
0

-- | buildSufT together with matchSufT to avoid having to repeat arguments. Apply partially for
-- multiple queries.
buildMatchSufT :: (Int -> a) -> (a -> Int -> a) -> (a -> a -> a)
               -> Int -> (Int -> Chr)
               -> Int -> (Int -> Chr)
               -> Maybe a
buildMatchSufT :: (Int -> a)
-> (a -> Int -> a)
-> (a -> a -> a)
-> Int
-> (Int -> Int)
-> Int
-> (Int -> Int)
-> Maybe a
buildMatchSufT Int -> a
fromLeaf a -> Int -> a
updEdge a -> a -> a
merge Int
n Int -> Int
at = (a -> Int -> a)
-> (Int -> Int) -> SufTNode a -> Int -> (Int -> Int) -> Maybe a
forall a.
(a -> Int -> a)
-> (Int -> Int) -> SufTNode a -> Int -> (Int -> Int) -> Maybe a
matchSufT a -> Int -> a
updEdge Int -> Int
at SufTNode a
st where
    st :: SufTNode a
st = (Int -> a)
-> (a -> Int -> a)
-> (a -> a -> a)
-> Int
-> (Int -> Int)
-> SufTNode a
forall a.
(Int -> a)
-> (a -> Int -> a)
-> (a -> a -> a)
-> Int
-> (Int -> Int)
-> SufTNode a
buildSufT Int -> a
fromLeaf a -> Int -> a
updEdge a -> a -> a
merge Int
n Int -> Int
at

--------------------------------------------------------------------------------
-- For tests

instance NFData a => NFData (SufTEdge a) where
    rnf :: SufTEdge a -> ()
rnf (SufTEdge Int
_ Int
_ SufTNode a
u) = SufTNode a -> ()
forall a. NFData a => a -> ()
rnf SufTNode a
u

instance NFData a => NFData (SufTNode a) where
    rnf :: SufTNode a -> ()
rnf (SufTNode a
a IntMap (SufTEdge a)
nxt) = a -> ()
forall a. NFData a => a -> ()
rnf a
a () -> () -> ()
`seq` IntMap (SufTEdge a) -> ()
forall a. NFData a => a -> ()
rnf IntMap (SufTEdge a)
nxt

-- | Draws a suffix tree. Can be used for debugging.
drawSufT :: Show a => SufTNode a -> String
drawSufT :: SufTNode a -> String
drawSufT = (SufTNode a -> String)
-> (SufTNode a -> [(Maybe String, SufTNode a)])
-> SufTNode a
-> String
forall a.
(a -> String) -> (a -> [(Maybe String, a)]) -> a -> String
draw SufTNode a -> String
forall a. Show a => SufTNode a -> String
showa SufTNode a -> [(Maybe String, SufTNode a)]
forall a. SufTNode a -> [(Maybe String, SufTNode a)]
nbs where
    showa :: SufTNode a -> String
showa (SufTNode a
a IntMap (SufTEdge a)
_) = a -> String
forall a. Show a => a -> String
show a
a
    nbs :: SufTNode a -> [(Maybe String, SufTNode a)]
nbs (SufTNode a
_ IntMap (SufTEdge a)
nxt) = [(String -> Maybe String
forall a. a -> Maybe a
Just ((Int, Int) -> String
forall a. Show a => a -> String
show (Int
left, Int
len)), SufTNode a
v)| SufTEdge Int
left Int
len SufTNode a
v <- IntMap (SufTEdge a) -> [SufTEdge a]
forall a. IntMap a -> [a]
IM.elems IntMap (SufTEdge a)
nxt]