module Data.Apart.Structures.Tree.Prefix (Prefix, Labeled (..), seek, insert) where
import Control.Applicative (Alternative (..))
import Control.Comonad (Comonad (..))
import Control.Comonad.Cofree (Cofree (..), unwrap)
import Control.Lens (Lens', (^.), (%~))
import Data.Maybe (isJust)
import Data.Function ((&))
import Data.Foldable (find)
import Data.Monoid (Monoid (..), (<>))
import Data.Apart.Structures.Stack (Stack)
type Prefix s t = Cofree (Labeled s t)
data Labeled s t a = Hop s (t a) deriving Show
symbol :: Lens' (Prefix s t a) s
symbol f (x :< Hop s ns) = (\new -> x :< Hop new ns) <$> f s
nodes :: Lens' (Prefix s t a) (t (Prefix s t a))
nodes f (x :< Hop s ns) = (\new -> x :< Hop s new) <$> f ns
instance Functor t => Functor (Labeled s t) where
fmap f (Hop s as) = Hop s $ f <$> as
instance Foldable t => Foldable (Labeled s t) where
foldr f acc (Hop s as) = foldr f acc as
instance Traversable t => Traversable (Labeled s t) where
traverse f (Hop s as) = Hop s <$> traverse f as
seek :: (Functor t, Foldable t, Eq s) => Stack s -> Prefix s t v -> Maybe v
seek (s :< Just ss) prefix@((==) s . flip (^.) symbol -> True) =
(<$>) extract $ find (isJust . seek ss) $ unwrap prefix
seek (s :< Nothing) prefix@((==) s . flip (^.) symbol -> True) = Just $ extract prefix
seek (s :< _) prefix@((==) s . flip (^.) symbol -> False) = Nothing
insert :: (Foldable t, Alternative t, Eq s) => Stack s -> v -> Prefix s t v -> Prefix s t v
insert (s :< _) x prefix@((==) s . flip (^.) symbol -> False) = prefix
insert (s :< Nothing) x prefix@((==) s . flip (^.) symbol -> True) = x :< unwrap prefix
insert (s :< Just ss@(s' :< Just _)) x prefix@((==) s . flip (^.) symbol -> True) =
prefix & nodes %~ (<$>) (insert ss x)
insert (s :< Just ss@(s' :< Nothing)) x prefix@((==) s . flip (^.) symbol -> True) =
prefix & nodes %~ (<|>) (pure $ x :< Hop s' empty)