SuffixTree.hs

import           Data.List               (foldl', tails)
import           Data.Map.Internal.Debug
import qualified Data.Map.Strict         as M
import           Data.Maybe
import           Data.Monoid
import qualified Data.Set as S

data TrieValue v
  = TrieValue
  { next :: Trie v
  , value :: v
  } deriving (Show, Eq)

instance Semigroup v => Semigroup (TrieValue v) where
  TrieValue n1 tv1 <> TrieValue n2 tv2 =
    TrieValue (n1 <> n2) (tv1 <> tv2)

-- | suffix tree, a trie populated with suffixes
newtype Trie v = Trie (M.Map Char (TrieValue v))
  deriving (Eq, Show)

instance Monoid k => Monoid (Trie k) where
  mempty = empty

instance Semigroup k => Semigroup (Trie k) where
  Trie m <> Trie n = Trie (m <> n)

empty :: Trie v
empty = Trie mempty

fromList :: Monoid v => [(v, String)] -> Trie v
fromList xs = foldr insertWithTails empty xs
  where
    insertWithTails (k,v) m =
      foldr (Main.insert k) m (tails v)

fromList_ :: String -> Trie ()
fromList_ vs = foldr (Main.insert ()) mempty (tails vs)

insert
  :: Monoid v
  => v
  -> String
  -> Trie v
  -> Trie v
insert _ [] t            = t
insert v (x:xs) (Trie m) =
  case M.lookup x m of
    Nothing ->
      Trie (M.insertWith (<>) x
            (TrieValue (insert v xs mempty) v) m)
    Just n ->
      Trie (M.insertWith (<>) x
            (TrieValue (insert v xs (next n)) v) m)

search :: Monoid v => String -> Trie v -> v
search [] _ = mempty
search xs m = go xs m
  where
    go [x] (Trie n) =
      case M.lookup x n of
        Nothing -> mempty
        Just g -> value g

    go (x:xs) (Trie n) =
      case M.lookup x n of
        Nothing -> mempty
        Just g -> go xs (next g)

size :: Trie v -> Int
size (Trie m) =
  case M.size m of
    0 -> 0
    n -> n + sum (size . next <$> M.elems m)

deleteById key (Trie m) = Trie $
  M.fromList
  [ (c, TrieValue (deleteById key n) (S.delete key keys))
  | (c, TrieValue n keys) <- M.toList m
  , S.singleton key /= keys
  ]

main :: IO ()
main = do
  let things =
        [ ( S.singleton 1
          , "eat"
          ),
          ( S.singleton 2
          , "drink"
          ),
          ( S.singleton 3
          , "sleep"
          ),
          ( S.singleton 4
          , "yay"
          ),
          ( S.singleton 5
          , "foo"
          )
        ]
  let ts = fromList things
  print $ search "e" ts
  print $ search "s" ts
  print $ search "ink" ts
  print $ search "eat" ts
  print $ search "a" ts
  print $ search "at" (deleteById 4 ts)