diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Joshua Simmons (c) 2017 + +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 Joshua Simmons 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/app/App.hs b/app/App.hs new file mode 100644 --- /dev/null +++ b/app/App.hs @@ -0,0 +1,20 @@ +-- | +-- Module : Main +-- Copyright : Joshua Simmons 2017 +-- License : BSD3 +-- +-- Maintainer : joshua.simmons@emptypath.com +-- +-- a single usage of the library for profiling use +-- +module Main +( main +) where + +import Data.SuffixArray + +main :: IO () +main = do + let input :: [Int] + input = take 400000 $ cycle [0 .. 199] + print . (\x -> (take 5 . justSuffixes $ x, take 5 . justLcp $ x)) . suffixArrayOne $ input diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,77 @@ +-- | +-- Module : Main +-- Copyright : Joshua Simmons 2017 +-- License : BSD3 +-- +-- Maintainer : joshua.simmons@emptypath.com +-- +-- suffix-array benchmarks using criterion +-- +module Main +( main +) where + +import Criterion.Main +import System.Random (newStdGen, randoms) + +import Data.SuffixArray +import Data.SuffixArray.Internal + +main :: IO () +main = do + g <- newStdGen + let rands, sorts, reps :: [Int] + rands = randoms g + sorts = [1..] + reps = concatMap (\x -> replicate x x) [1..] + allDists = [rands, sorts, reps] + defaultMain + [ + bgroup "lcp" + [ bench (unwords [show (sz, k), var']) + $ whnf (\(n,a) -> let n' = n `div` length allDists + in var (map (take n' . map (`mod` a)) allDists)) + (sz, k) + | k <- [5, 40, 1000] + , sz <- [5000, 35000 .. 215000] + , (var, var') <- [ (naiveLcp, "naiveLcp") + , (justLcp . suffixArray, "suffixArray(lcp)")] + , interesting var' sz k "lcp" + ] + ,bgroup "single_suffixes" + [ bench (unwords [show (sz, k), dist', var']) + $ whnf (\(n,a) -> var (take n (map (`mod` a) dist))) (sz, k) + | (dist, dist') <- [ (rands, "rands"), (sorts, "sorts") + , (reps, "reps")] + , k <- [5, 40, 1000] + , sz <- [5000, 25000 .. 105000] ++ [200000] + , (var, var') <- [ (naiveOne, "naiveOne") + , (justSuffixes . suffixArrayOne, "suffixArrayOne")] + , interesting var' sz k dist' + ] + ,bgroup "all_together" + [ bench (unwords [show (sz, k), var']) + $ whnf (\(n,a) -> let n' = n `div` length allDists + in var (map (take n' . map (`mod` a)) allDists)) + (sz, k) + | k <- [5, 40, 1000] + , sz <- [5000, 35000 .. 215000] + , (var, var') <- [ (naive, "naive") + , (justSuffixes . suffixArray, "suffixArray")] + , interesting var' sz k "all" + ] + ] + +interesting :: String -> Int -> Int -> String -> Bool +interesting "naiveOne" n _ v + | n > 50000 && v == "sorts" = False + | n > 100000 && v /= "rands" = False + | otherwise = True +interesting "naive" n k _ + | k < 40 && n > 60000 = False + | otherwise = True +interesting "naiveLcp" n k _ + | k < 40 && n > 65000 = False + | k < 500 && n > 125000 = False + | otherwise = True +interesting _ _ _ _ = True diff --git a/src/Data/SuffixArray.hs b/src/Data/SuffixArray.hs new file mode 100644 --- /dev/null +++ b/src/Data/SuffixArray.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE CPP #-} +-- | +-- Module : Data.SuffixArray +-- Copyright : Joshua Simmons 2017 +-- License : BSD3 +-- +-- Maintainer : joshua.simmons@emptypath.com +-- +-- Suffix array library main module +-- +module Data.SuffixArray +( SuffixArray(..) +, suffixArray +, suffixArrayOne +, Alpha(..) +, justAlphas +, justLcp +, justSuffixes +) where + +#ifdef __GLASGOW_HASKELL__ +#if __GLASGOW_HASKELL__ < 710 +import Control.Applicative +#endif +#endif + +import Control.Monad (forM_, when) +import Control.Monad.ST (ST) +import qualified Data.Array.IArray as A +import Data.Array.IArray (Array, (!)) +import Data.Array.MArray ( newListArray, newArray_ + , readArray, writeArray) +import Data.Array.ST (STUArray, runSTUArray) +import Data.Array.Unboxed (UArray) +import Data.Array.Unsafe (unsafeFreeze) +import Data.List (sortBy) +import Data.Ord (comparing) +import Data.STRef ( newSTRef, readSTRef, writeSTRef + , modifySTRef') + +import Data.SuffixArray.Internal + +-- | Holds the computed suffix array data +data SuffixArray a = SuffixArray + { toSuffixes :: UArray Int Int + -- ^ The actual array of suffixes in lexicographic + -- order. + , toAlphas :: Array Int (Alpha a) + -- ^ The original string(s) with `Sentinal` values + -- included after each string. + , toLcp :: UArray Int Int + -- ^ Longest Common Prefix of each suffix with the + -- previous one in lexicographic order + } + deriving (Eq, Ord, Show) + +type Arr s = STUArray s Int Int + +-- | Compute the suffix array of the given string(s) concatenated together +-- with `Sentinal`s after each. +-- +-- worst case O(n lg n) time +-- (where n is the sum of the string lengths + the number of strings) +suffixArray :: Ord a => [[a]] -> SuffixArray a +suffixArray xs = SuffixArray ss as lcp + where + n = snd $ A.bounds as + as = let ps = prepare xs + n' = length ps - 1 + in A.listArray (0, n') ps + -- we represent each suffix as the number of characters we have + -- to drop from the original string to get that suffix + -- + -- and then we order them by their first letter and convert those + -- first letters into `rank`s, which are `Int`s that preserve the + -- same `Ord`ering. This is useful so we can sort them more easily + -- (allows using counting sort), and will help code reuse in the + -- main body. + -- + -- Note: We actually don't care about the ordering of suffixes yet, + -- it's just necessary to use the `rank` function. + orderedByHead = sortBy (comparing snd) . zip [0 ..] $ A.elems as + ranked = let (is, js) = unzip orderedByHead + in zip is (rank js) + ss :: UArray Int Int + ss = runSTUArray $ do + s <- newListArray (0, n) (map fst ranked) -- the suffixes + r <- newArray_ (0, n) -- the rank of each suffix + forM_ ranked $ uncurry (writeArray r) + t <- newArray_ (0, n) -- scratch array + c <- newArray_ (0, n) -- counts array + go 1 s r t c + -- After each iteration of `go`, the suffixes are sorted by their + -- k*2 first characters. k doubles each time, and in each iteration + -- we do O(n) work and are then ready for the next iteration. + go :: forall s. Int -> Arr s -> Arr s -> Arr s -> Arr s -> ST s (Arr s) + go k s r t c + | k > n = return s + | otherwise = do + let getR 0 x = readArray r x + getR i x = let ix = i + x + in if ix > n then return 0 + else readArray r ix + -- counting sort of suffixes, from s into s' + -- ordered by the rank of suffix i + x, for suffix x + -- (that is, suffix x without its first i characters) + csort i src dest = do + forM_ [0 .. n] $ flip (writeArray c) 0 -- zero out the counts + let f = getR i + -- count how many of each rank there are + writeArray c 0 i -- takes care of all that would be automatically 0 + forM_ [i .. n] $ \x -> do -- count the appropriate values in r + x' <- readArray r x + v <- readArray c x' + writeArray c x' (v+1) + -- replace each element in c with the starting index of + -- elements with that value + soFar <- newSTRef 0 + forM_ [0 .. n] $ \x -> do + v <- readArray c x + readSTRef soFar >>= writeArray c x + modifySTRef' soFar (+v) + elemsS <- (A.elems :: UArray Int Int -> [Int]) <$> unsafeFreeze src + forM_ elemsS $ \x -> do + r' <- f x -- rank of it + idx <- readArray c r' -- where it goes, based on its rank + writeArray c r' (idx + 1) -- next suffix with this rank goes + -- one later + writeArray dest idx x + csort k s t -- these two counting sorts comprise a radix sort of the + csort 0 t s -- suffixes by their rank pairs + -- now re-rank the suffixes in order + fstSuffix <- readArray s 0 + prevVal <- ((,) <$> getR 0 fstSuffix <*> getR k fstSuffix) >>= newSTRef + nextRank <- newSTRef 0 + elemsS <- (A.elems :: UArray Int Int -> [Int]) <$> unsafeFreeze s + forM_ elemsS $ \x -> do + val <- (,) <$> getR 0 x <*> getR k x + val' <- readSTRef prevVal + -- if its old rank pair is the same as of the previous suffix + -- (in partially sorted order), it gets the same rank, otherwise + -- we increase by one + when (val /= val') $ modifySTRef' nextRank succ + readSTRef nextRank >>= writeArray t x + writeSTRef prevVal val + maxRank <- readSTRef nextRank + if maxRank < n + then go (k*2) s t r c -- double the size of the prefix we're sorting by + else return s -- ranks are already unique for all, stop early + -- LCP array in the same order as the suffix array + lcp = A.ixmap (0, n) (ss !) plcp + -- PLCP, permuted LCP array which is in order by position instead of + -- lexicographic order by the suffix being referred to. + -- + -- Algoritm is courtesy of the paper "Permuted Longest-Common-Prefix + -- Array" by Kärkkäinen, et al. + -- http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.186.2185 + -- (several PDFs available free online) + -- + -- This runs in worst case O(n) time + plcp = runSTUArray plcp' + plcp' :: forall s. ST s (Arr s) + plcp' = do + -- keep track of what suffix is before each one, in lexicographic + -- order (the `first` one has none before it, so we treat it special) + let first = ss ! 0 + (prev :: Arr s) <- newArray_ (0, n) + forM_ [1 .. n] $ \i -> writeArray prev (ss ! i) (ss ! (i-1)) + len <- newSTRef 0 + res <- newArray_ (0, n) + forM_ [0 .. n] $ \i -> + if i == first -- no previous prefix + then writeSTRef len 0 >> writeArray res i 0 + else do + -- See the PLCP Array paper for details, but the important + -- part is that PLCP[i] >= PLCP[i-1] - 1, which lets us + -- skip a *lot* of character comparisons in the worst-case + -- + -- This is otherwise essentially the same as the naive LCP + -- computation (see 'Data.SuffixArray.Internal.naiveLcp') + len' <- readSTRef len + prev' <- readArray prev i + let suffixOff x = map (as !) [x ..] + newMatching = length . takeWhile id + $ zipWith (==) (suffixOff (i + len')) + (suffixOff (prev' + len')) + writeArray res i (len' + newMatching) + writeSTRef len $ max 0 (len' + newMatching - 1) + return res + +-- | Convenience function to compute the suffix array of a single string. +-- (Still gets a `Sentinal` at the end) +-- +-- worst case O(n lg n) time +-- (where n is the length of the string) +suffixArrayOne :: Ord a => [a] -> SuffixArray a +suffixArrayOne = suffixArray . (:[]) + +-- | Convenience function to just give a list of the suffixes in +-- lexicographic order. +justSuffixes :: SuffixArray a -> [Int] +justSuffixes = A.elems . toSuffixes + +-- | Convenience function to just give a list characters in the +-- concatenated original strings. +justAlphas :: SuffixArray a -> [Alpha a] +justAlphas = A.elems . toAlphas + +-- | Convenience function to just give a list of the longest common +-- prefix of every suffix with the previous suffix in lexicographic +-- order. +justLcp :: SuffixArray a -> [Int] +justLcp = A.elems . toLcp diff --git a/src/Data/SuffixArray/Internal.hs b/src/Data/SuffixArray/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/SuffixArray/Internal.hs @@ -0,0 +1,114 @@ +-- | +-- Module : Data.SuffixArray.Internal +-- Copyright : Joshua Simmons 2017 +-- License : BSD3 +-- +-- Maintainer : joshua.simmons@emptypath.com +-- +-- Stability : unstable +-- +-- Internal implementation details, unstable and not +-- to be relied upon for any reason. +-- +module Data.SuffixArray.Internal +( Alpha(..) +, naive +, naiveOne +, naiveLcp +, naiveLcpOne +, prepare +, prepareOne +, rank +, suffixes +) where + +import Data.List (group, sort, sortBy) +import Data.Ord (comparing) + +-- | Yields the non-empty suffixes of a list in order of decreasing length. +-- +-- This differs from `Data.List.tails` in that it does not include the +-- empty list at the end. +suffixes :: [a] -> [[a]] +suffixes xxs@(_:xs) = xxs : suffixes xs +suffixes [] = [] + +-- | A character in a string (or set of strings) we're going to compute the +-- suffix array of. +-- Includes `Sentinal` markers for the end of strings. +data Alpha a = Sentinal Int -- ^ Used to mark the end of a string. + -- The `Int` parameter is used to encode + -- which string this is the end of, in cases + -- where there are multiple. + | Alpha a -- ^ An actual character in the string. + deriving (Eq, Ord, Show) + +-- | Convenience value containing `Sentinal`s in order. +sentinals :: [Alpha a] +sentinals = map Sentinal [0..] + +-- | Prepare a list of strings to compute the suffix array of them. +-- Just wraps every character in `Alpha` and adds `Sentinal`s to the end of +-- each string, and concatenates it together. +prepare :: [[a]] -> [Alpha a] +prepare = concat . zipWith (\a b -> b ++ [a]) sentinals . map (map Alpha) + +-- | Convenience function to `prepare` a single string. +prepareOne :: [a] -> [Alpha a] +prepareOne = prepare . (:[]) + +-- | A naively implemented suffix array implementation which will be used +-- for correctness checking and possibly to benchmark against. Shouldn't +-- usually be used in production code, as it is quite slow in the worst +-- case. In cases with few identical suffixes, it can actually perform +-- quite well. See benchmarks for some details. +-- +-- worst case O(n^2 lg n) time +-- (where n is the sum of the string lengths + the number of strings) +naive :: Ord a => [[a]] -> [Int] +naive = + map fst . sortBy (comparing snd) . zip [0 ..] . suffixes . prepare + +-- | Convenience wrapper around `naive` for a single string. +-- +-- worst case O(n^2 lg n) time +-- (where n is the length of the string) +naiveOne :: Ord a => [a] -> [Int] +naiveOne = naive . (:[]) + +-- | A naively implemented LCP implementation, used for correctness +-- testing the actual algorithm. +-- +-- The Longest Common Prefix list gives the longest common prefix of +-- each suffix and the previous suffix, with the suffixes in lexicographic +-- order. +-- +-- worst case O(n^2 lg n) time +-- (where n is the sum of the string lengths + the number of strings) +-- +-- The LCP part is an extra O(n^2) in addition to the work of computing +-- the suffix array in the first place. +naiveLcp :: Ord a => [[a]] -> [Int] +naiveLcp = (\xs -> zipWith lcp xs ([] : xs)) . sort . suffixes . prepare + where lcp as bs = length . takeWhile id $ zipWith (==) as bs + +-- | Convenience wrapper around `naiveLcp` for a single string. +-- +-- worst case O(n^2 lg n) time +-- (where n is the length of the string) +-- +-- The LCP part is an extra O(n^2) in addition to the work of computing +-- the suffix array in the first place. +naiveLcpOne :: Ord a => [a] -> [Int] +naiveLcpOne = naiveLcp . (:[]) + +-- | Take a sorted list of elements and replace each value with an `Int` +-- such that any comparisons between elements in the original list would +-- yield exactly the same result in the output list. +-- +-- i.e.: let rs = rank xs +-- in all [ (xs!!i) `compare` (xs!!j) == (rs!!i) `compare` (rs!!j) +-- | let idx = [0 .. length xs - 1], i <- idx, j <- idx +-- ] +rank :: Ord a => [a] -> [Int] +rank = concat . zipWith (map . const) [0 ..] . group diff --git a/suffix-array.cabal b/suffix-array.cabal new file mode 100644 --- /dev/null +++ b/suffix-array.cabal @@ -0,0 +1,63 @@ +name: suffix-array +version: 0.3.0.0 +synopsis: Simple and moderately efficient suffix array implementation +description: A simple implementation of a suffix array, with + longest-common-prefix array. While not + asymptotically optimal, performs well in practice + for medium use. +homepage: https://github.com/kadoban/suffix-array#readme +bug-reports: https://github.com/kadoban/suffix-array/issues +license: BSD3 +license-file: LICENSE +author: Joshua Simmons +maintainer: joshua.simmons@emptypath.com +copyright: (c) 2017 Joshua Simmons +category: Data Structures +build-type: Simple +-- extra-source-files: README.md +cabal-version: >=1.10 +tested-with: GHC==7.8.4, GHC==7.10.2, GHC==7.10.3, GHC==8.0.1 + +library + hs-source-dirs: src + exposed-modules: Data.SuffixArray + , Data.SuffixArray.Internal + build-depends: array >= 0.5 && < 0.6 + , base >= 4.7 && < 5 + default-language: Haskell2010 + ghc-options: -O2 + +executable suffix-array-exe + hs-source-dirs: app + main-is: App.hs + build-depends: base + , suffix-array + default-language: Haskell2010 + ghc-options: -O2 -rtsopts + +benchmark suffix-array-bench + type: exitcode-stdio-1.0 + hs-source-dirs: bench + main-is: Bench.hs + build-depends: base + , criterion >= 1.0 && < 1.2 + , random >= 1.0 && < 1.2 + , suffix-array + default-language: Haskell2010 + +test-suite suffix-array-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Test.hs + build-depends: array + , base + , containers >= 0.5 && < 0.6 + , suffix-array + , tasty >= 0.10 && < 0.12 + , tasty-hunit >= 0.9 && < 0.10 + , tasty-quickcheck >= 0.8 && < 0.9 + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/kadoban/suffix-array diff --git a/test/Test.hs b/test/Test.hs new file mode 100644 --- /dev/null +++ b/test/Test.hs @@ -0,0 +1,87 @@ +-- | +-- Module : Main +-- Copyright : Joshua Simmons 2017 +-- License : BSD3 +-- +-- Maintainer : joshua.simmons@emptypath.com +-- +-- suffix-array test-suite using tasty and etc. +-- +module Main +( main +) where + +import qualified Data.Array.IArray as A +import qualified Data.Set as S +import Test.Tasty +import Test.Tasty.HUnit +import qualified Test.Tasty.QuickCheck as QC + +import Data.List (sort, tails) + +import Data.SuffixArray +import Data.SuffixArray.Internal + +main :: IO () +main = defaultMain tests + +tests :: TestTree +tests = testGroup "Tests" [basics, naives, actual, actualLcp] + +basics :: TestTree +basics = testGroup "tests of basic underlying utils and etc." + [ QC.testProperty "suffixes == filter (not . null) tails" $ + \xs -> suffixes (xs :: [Int]) == filter (not . null) (tails xs) + , QC.testProperty "(length . suffixes) == length" $ + \xs -> length (suffixes xs) == length (xs :: [Int]) + ] + +naives :: TestTree +naives = testGroup "test the naive implementation to make sure it works" + [ testCase "banana" $ + naiveOne "banana" @?= [6, 5, 3, 1, 0, 4, 2] -- wikipedia example + , testCase "banana hammock" $ -- hand calculated + naive ["banana", "hammock"] @?= [6,14,5,8,3,1,0,12,7,13,9,10,4,2,11] + , QC.testProperty "length" $ + \xs -> length (xs :: String) + 1 == length (naiveOne xs) + , QC.testProperty "distinct" $ + \xs -> let xs' = naiveOne (xs :: String) + in length xs' == length (distinct xs') + , QC.testProperty "distinct2" $ + \xs -> let xs' = naiveOne (xs :: [Int]) + in sort xs' == distinct xs' + , QC.testProperty "length of many" $ + \xs -> sum (map length xs) + length (xs :: [String]) == length (naive xs) + , QC.testProperty "empty first" $ + \xs -> head (naiveOne xs) == length (xs :: [Integer]) + ] + +actual :: TestTree +actual = testGroup "test the actual implementation to make sure it works" + [ QC.testProperty "against naiveOne" $ + \xs -> naiveOne (xs :: [Int]) == A.elems (toSuffixes (suffixArrayOne xs)) + , QC.testProperty "against naive" $ + \xs -> naive (xs :: [[Int]]) == A.elems (toSuffixes (suffixArray xs)) + , testCase "[0]" $ + A.elems (toSuffixes $ suffixArrayOne [0 :: Integer]) @?= [1, 0] + , testCase "bunch of reps" $ + naive [take 1000 reps, take 2000 reps] + @?= justSuffixes (suffixArray [take 1000 reps, take 2000 reps]) + ] + +reps :: [Int] +reps = concatMap (\x -> replicate x x) [1..] + +actualLcp :: TestTree +actualLcp = testGroup "test the actual implementation of LCP stuff" + [ QC.testProperty "against naiveLcpOne" $ + \xs -> naiveLcpOne (xs :: [Int]) == justLcp (suffixArrayOne xs) + , QC.testProperty "against naiveLcp" $ + \xs -> naiveLcp (xs :: [[Int]]) == justLcp (suffixArray xs) + , testCase "bunch of reps" $ + naiveLcp [take 1000 reps, take 2000 reps, take 3000 reps] + @?= justLcp (suffixArray [take 1000 reps, take 2000 reps, take 3000 reps]) + ] + +distinct :: Ord a => [a] -> [a] +distinct = S.toList . S.fromList