module Data.SuffixArray
( SuffixArray(..)
, suffixArray
, simpleEquator
, fancyEquator
, shiftList
, composeLists
, populateClassesBy
, fromList
, toList
, elems
) where
import Data.Ix
import Data.List (foldl', (!!))
import System.IO.Unsafe
import Foreign.Storable
import Foreign.Marshal.Array
import Foreign.Ptr
import Data.Bits (shiftL)
import Debug.Trace
import qualified Data.Vector as V
import qualified Data.Vector.Generic.Mutable as MVector
import Data.CountingSort
data SuffixArray a = SuffixArray (V.Vector a) (V.Vector Int)
deriving Show
elems :: SuffixArray a -> V.Vector (V.Vector a)
elems (SuffixArray d i) = V.map vecAt i
where vecAt idx = V.drop idx d
fromList :: (Ix a, Ord a, Bounded a) => [a] -> SuffixArray a
fromList = suffixArray . V.fromList
toList :: SuffixArray a -> [[a]]
toList (SuffixArray d i) = V.foldr vecAt [] i
where vecAt idx l = V.toList (V.drop idx d) : l
suffixArray :: (Ix a, Ord a, Bounded a)
=> V.Vector a -> SuffixArray a
suffixArray s = let p = countingSort s (V.generate n id)
equator = simpleEquator s p
c = populateClassesBy equator p
in go 0 p c
where
n = V.length s
go h p c | (1 << h) >= n = SuffixArray s p
go h p c = let
pn = shiftList n h p
ck = V.toList $ composeLists c pn
p' = countingSort (V.fromList ck) pn
equator = fancyEquator c p' h n
c' = populateClassesBy equator p'
in go (h + 1) p' c'
(<<) = shiftL
type Equator = Int -> Int -> Bool
simpleEquator :: (Ix a, Ord a, Bounded a)
=> V.Vector a -> V.Vector Int -> Equator
simpleEquator s indexes i j = (s V.! (indexes V.! i)) == (s V.! (indexes V.! j))
fancyEquator :: (Ix a, Ord a, Bounded a)
=> V.Vector a -> V.Vector Int -> Int -> Int -> Equator
fancyEquator s indexes h n i j
= (s V.! i') == (s V.! j') && (s V.! mid1) == (s V.! mid2)
where mid1 = ((i' + (1 << h)) `mod` n)
mid2 = ((j' + (1 << h)) `mod` n)
i' = indexes V.! i
j' = indexes V.! j
shiftList :: Int -> Int -> V.Vector Int -> V.Vector Int
shiftList n h p = V.map step p
where step v = let x = (v (1 << h))
x' = if x < 0 then x + n else x
in x'
composeLists :: V.Vector Int -> V.Vector Int -> V.Vector Int
composeLists c indexes = V.map (c V.!) indexes
populateClassesBy :: Equator -> V.Vector Int -> V.Vector Int
populateClassesBy equals p = unsafePerformIO $ do
let n = V.length p
arr <- MVector.replicate n 0
let
go i classNum | i == n = return ()
go i classNum = do
let pcur = p V.! i
let newClassNum = if i `equals` (i 1)
then classNum
else classNum + 1
MVector.write arr pcur (newClassNum 1)
go (i + 1) newClassNum
go 1 1
V.unsafeFreeze arr