module Yi.CompletionTree (
CompletionTree (CompletionTree),
fromList, toList,
complete, update,
pretty,
unCompletionTree
) where
import Control.Arrow (first)
import Data.Function (on)
import Data.List (partition, maximumBy, intercalate)
import qualified Data.Map.Strict as M
import Data.Map.Strict (Map)
import Data.Maybe (isJust, fromJust, listToMaybe, catMaybes)
import qualified Data.ListLike as LL
import Data.ListLike (ListLike)
import Lens.Micro.Platform (over, Lens', _2, (.~), (&))
import Data.Binary (Binary)
newtype CompletionTree a = CompletionTree {_unCompletionTree :: (Map a (CompletionTree a))}
deriving (Monoid, Eq, Binary)
unCompletionTree :: Lens' (CompletionTree a) (Map a (CompletionTree a))
unCompletionTree f ct = (\unCompletionTree' -> ct {_unCompletionTree = unCompletionTree'}) <$>
f (_unCompletionTree ct)
instance (Ord a, Show a, ListLike a i) => Show (CompletionTree a) where
show ct = "fromList " ++ show (toList ct)
fromList :: (Ord a, ListLike a i, Eq i) => [a] -> CompletionTree a
fromList [] = mempty
fromList (x:xs)
| x == mempty = over unCompletionTree (M.insert mempty mempty) (fromList xs)
| otherwise = case maximumBy' (compare `on` childrenIn xs) (tail $ LL.inits x) of
Nothing -> over unCompletionTree (M.insert x mempty) (fromList xs)
Just parent -> case first (x:) $ partition (parent `LL.isPrefixOf`) xs of
([_],rest) -> over unCompletionTree (M.insert parent mempty) $ fromList rest
(hasParent, rest) -> over unCompletionTree (M.insert parent (fromList $
map (fromJust . LL.stripPrefix parent) hasParent)) $ fromList rest
where childrenIn list parent = length $ filter (parent `LL.isPrefixOf`) list
maximumBy' :: Eq a => (a -> a -> Ordering) -> [a] -> Maybe a
maximumBy' cmp l | atleast 2 (== max') l = Nothing
| otherwise = Just max'
where max' = maximumBy cmp l
atleast :: Int -> (a -> Bool) -> [a] -> Bool
atleast 0 _ _ = True
atleast _ _ [] = False
atleast n cmp' (x:xs) | cmp' x = atleast (n 1) cmp' xs
| otherwise = atleast n cmp' xs
complete :: (Eq i, Ord a, ListLike a i) => CompletionTree a -> (a, CompletionTree a)
complete (CompletionTree ct)
| M.size ct == 1 = if snd (M.elemAt 0 ct) == mempty
then M.elemAt 0 ct & _2 .~ fromList [mempty]
else M.elemAt 0 ct
| otherwise = (mempty,CompletionTree ct)
update :: (Ord a, ListLike a i, Eq i) => CompletionTree a -> a -> CompletionTree a
update (CompletionTree ct) p
| mempty == p = error "Can't update a CompletionTree with a mempty"
| isJust one && mempty == fromJust one = CompletionTree $ M.singleton mempty mempty
| isJust one = fromJust one
| isJust remaining = uncurry update $ fromJust remaining
| otherwise = CompletionTree $ M.mapKeys fromJust
$ M.filterWithKey (const . isJust)
$ M.mapKeys (LL.stripPrefix p) ct
where
one = M.lookup p ct
remaining = listToMaybe . catMaybes $
map (\p' -> (,fromJust $ LL.stripPrefix p' p) <$> M.lookup p' ct) (tail $ LL.inits p)
toList :: (Ord a, ListLike a i) => CompletionTree a -> [a]
toList ct
| mempty == ct = []
| otherwise = toList' ct
where
toList' (CompletionTree ct')
| M.null ct' = [mempty]
| otherwise = concat $ M.elems $ M.mapWithKey (\k v -> map (k `LL.append`) $ toList' v) ct'
pretty :: Show a => CompletionTree a -> String
pretty (CompletionTree ct)
| M.null ct = ""
| otherwise = "[" ++ intercalate "|" (M.elems (M.mapWithKey (\k v -> shows k (pretty v)) ct)) ++ "]"