module Data.BAByNF.Core.Tree
    ( Tree (..)
    , Node (..)
    , nodes
    , empty
    , singleton
    , asSingleton
    , stringify
    , stringifyNode
    , mergeStrings
    , mergeStringsInNode
    , dropRefs
    , getChildrenWithRef
    , getChildWithRef
    , tryGetChildWithRef
    , getDescendantsWithPath
    , getFirstDescendantWithPath
    , tryGetFirstPath
    , getSubtreeIfRef
    , isStringEq
    , isRefOf
    ) where

import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString

import Data.BAByNF.Core.Ref (Ref)
import Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.BAByNF.Core.Ref as Ref

data Tree a where Tree :: Ref a => [Node a] -> Tree a

deriving instance (Eq a) => Eq (Tree a)
deriving instance (Show a) => Show (Tree a)

data Node a where
    StringNode :: ByteString -> Node a
    RefNode :: (Ref a) => a -> Tree a -> Node a

deriving instance (Eq a) => Eq (Node a)
deriving instance (Show a) => Show (Node a)

instance Semigroup (Tree a) where
    (<>) :: Tree a -> Tree a -> Tree a
    <> :: Tree a -> Tree a -> Tree a
(<>) (Tree [Node a]
a) (Tree [Node a]
b) = [Node a] -> Tree a
forall a. Ref a => [Node a] -> Tree a
Tree ([Node a] -> Tree a) -> [Node a] -> Tree a
forall a b. (a -> b) -> a -> b
$ [Node a]
a [Node a] -> [Node a] -> [Node a]
forall a. [a] -> [a] -> [a]
++ [Node a]
b

nodes :: (Ref a) => Tree a -> [Node a]
nodes :: forall a. Ref a => Tree a -> [Node a]
nodes (Tree [Node a]
ns) = [Node a]
ns

empty :: (Ref a) => Tree a
empty :: forall a. Ref a => Tree a
empty = [Node a] -> Tree a
forall a. Ref a => [Node a] -> Tree a
Tree []

singleton :: (Ref a) => Node a -> Tree a
singleton :: forall a. Ref a => Node a -> Tree a
singleton Node a
node = [Node a] -> Tree a
forall a. Ref a => [Node a] -> Tree a
Tree [Node a
node]

asSingleton :: (Ref a) => Tree a -> Maybe (Node a)
asSingleton :: forall a. Ref a => Tree a -> Maybe (Node a)
asSingleton (Tree [Node a
x]) = Node a -> Maybe (Node a)
forall a. a -> Maybe a
Just Node a
x
asSingleton Tree a
_ = Maybe (Node a)
forall a. Maybe a
Nothing

stringify :: Tree a -> ByteString
stringify :: forall a. Tree a -> ByteString
stringify (Tree [Node a]
ns) = [ByteString] -> ByteString
ByteString.concat ([ByteString] -> ByteString)
-> ([Node a] -> [ByteString]) -> [Node a] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Node a -> ByteString) -> [Node a] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Node a -> ByteString
forall a. Node a -> ByteString
stringifyNode ([Node a] -> ByteString) -> [Node a] -> ByteString
forall a b. (a -> b) -> a -> b
$ [Node a]
ns

stringifyNode :: Node a -> ByteString
stringifyNode :: forall a. Node a -> ByteString
stringifyNode (RefNode a
_ Tree a
tree) = Tree a -> ByteString
forall a. Tree a -> ByteString
stringify Tree a
tree
stringifyNode (StringNode ByteString
bs) = ByteString
bs

mergeStrings :: Tree a -> Tree a
mergeStrings :: forall a. Tree a -> Tree a
mergeStrings (Tree [Node a]
ns) = [Node a] -> Tree a
forall a. Ref a => [Node a] -> Tree a
Tree ([Node a] -> Tree a) -> [Node a] -> Tree a
forall a b. (a -> b) -> a -> b
$ [Node a] -> [Node a]
forall {a}. [Node a] -> [Node a]
merge [Node a]
ns
    where merge :: [Node a] -> [Node a]
merge [] = []
          merge [Node a
x] = [Node a -> Node a
forall a. Node a -> Node a
mergeStringsInNode Node a
x]
          merge (Node a
x:[Node a]
xs) = Node a -> Node a
forall a. Node a -> Node a
mergeStringsInNode Node a
x Node a -> [Node a] -> [Node a]
forall a. a -> [a] -> [a]
: [Node a] -> [Node a]
merge [Node a]
xs


mergeStringsInNode :: Node a -> Node a
mergeStringsInNode :: forall a. Node a -> Node a
mergeStringsInNode (RefNode a
ref Tree a
tree) = a -> Tree a -> Node a
forall a. Ref a => a -> Tree a -> Node a
RefNode a
ref (Tree a -> Tree a
forall a. Tree a -> Tree a
mergeStrings Tree a
tree)
mergeStringsInNode Node a
node = Node a
node

dropRefs :: [a] -> Tree a -> Tree a
dropRefs :: forall a. [a] -> Tree a -> Tree a
dropRefs [a]
refs (Tree [Node a]
ns) = [Node a] -> Tree a
forall a. Ref a => [Node a] -> Tree a
Tree ([Node a] -> Tree a) -> [Node a] -> Tree a
forall a b. (a -> b) -> a -> b
$ [Node a]
ns [Node a] -> (Node a -> [Node a]) -> [Node a]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Node a -> [Node a]
applyDrop
    where applyDrop :: Node a -> [Node a]
applyDrop node :: Node a
node@(StringNode ByteString
_) = [Node a
node]
          applyDrop (RefNode a
ref Tree a
tree) =
            let tree' :: Tree a
tree'@(Tree [Node a]
ns') = [a] -> Tree a -> Tree a
forall a. [a] -> Tree a -> Tree a
dropRefs [a]
refs Tree a
tree
             in if (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (a -> a -> Bool
forall a. Ref a => a -> a -> Bool
Ref.eq a
ref) [a]
refs
                then [Node a]
ns'
                else [a -> Tree a -> Node a
forall a. Ref a => a -> Tree a -> Node a
RefNode a
ref Tree a
tree']

getChildrenWithRef :: a -> Tree a -> [Tree a]
getChildrenWithRef :: forall a. a -> Tree a -> [Tree a]
getChildrenWithRef a
ref (Tree [Node a]
ns) = [Node a]
ns [Node a] -> (Node a -> [Tree a]) -> [Tree a]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Node a -> [Tree a]
filterOnRef
    where filterOnRef :: Node a -> [Tree a]
filterOnRef (RefNode a
ref' Tree a
subtree) = [Tree a
subtree | a -> a -> Bool
forall a. Ref a => a -> a -> Bool
Ref.eq a
ref a
ref']
          filterOnRef Node a
_ = []

getChildWithRef :: a -> Tree a -> Maybe (Tree a)
getChildWithRef :: forall a. a -> Tree a -> Maybe (Tree a)
getChildWithRef a
ref Tree a
tree =
    case a -> Tree a -> [Tree a]
forall a. a -> Tree a -> [Tree a]
getChildrenWithRef a
ref Tree a
tree of
        [] -> Maybe (Tree a)
forall a. Maybe a
Nothing
        Tree a
x : [Tree a]
_ -> Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
x

tryGetChildWithRef :: (Ref a) => a -> Tree a -> Either String (Tree a)
tryGetChildWithRef :: forall a. Ref a => a -> Tree a -> Either String (Tree a)
tryGetChildWithRef a
ref Tree a
tree = 
    case a -> Tree a -> Maybe (Tree a)
forall a. a -> Tree a -> Maybe (Tree a)
getChildWithRef a
ref Tree a
tree of
        Maybe (Tree a)
Nothing -> String -> Either String (Tree a)
forall a b. a -> Either a b
Left (String -> Either String (Tree a))
-> String -> Either String (Tree a)
forall a b. (a -> b) -> a -> b
$ String
"no subtree with ref <" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Ref a => a -> String
Ref.display a
ref String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"> defined"
        Just Tree a
subtree -> Tree a -> Either String (Tree a)
forall a b. b -> Either a b
Right Tree a
subtree

getDescendantsWithPath :: (Ref a) => NonEmpty a -> Tree a -> [Tree a]
getDescendantsWithPath :: forall a. Ref a => NonEmpty a -> Tree a -> [Tree a]
getDescendantsWithPath (a
r :| [a]
rs) Tree a
tree =
    let matching :: [Tree a]
matching = a -> Tree a -> [Tree a]
forall a. a -> Tree a -> [Tree a]
getChildrenWithRef a
r Tree a
tree
     in case [a]
rs of
        [] -> [Tree a]
matching
        a
r' : [a]
rs' -> [Tree a]
matching [Tree a] -> (Tree a -> [Tree a]) -> [Tree a]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= NonEmpty a -> Tree a -> [Tree a]
forall a. Ref a => NonEmpty a -> Tree a -> [Tree a]
getDescendantsWithPath (a
r' a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
rs')
getFirstDescendantWithPath :: (Ref a) => NonEmpty a -> Tree a -> Maybe (Tree a)
getFirstDescendantWithPath :: forall a. Ref a => NonEmpty a -> Tree a -> Maybe (Tree a)
getFirstDescendantWithPath NonEmpty a
refs Tree a
tree = case NonEmpty a -> Tree a -> [Tree a]
forall a. Ref a => NonEmpty a -> Tree a -> [Tree a]
getDescendantsWithPath NonEmpty a
refs Tree a
tree of
    [] -> Maybe (Tree a)
forall a. Maybe a
Nothing
    Tree a
x:[Tree a]
_ -> Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
x

tryGetFirstPath :: (Ref a) => NonEmpty a -> Tree a -> Either String (Tree a)
tryGetFirstPath :: forall a. Ref a => NonEmpty a -> Tree a -> Either String (Tree a)
tryGetFirstPath (a
r :| [a]
rs) Tree a
tree =
    let e :: Either String (Tree a)
e = a -> Tree a -> Either String (Tree a)
forall a. Ref a => a -> Tree a -> Either String (Tree a)
tryGetChildWithRef a
r Tree a
tree
     in case [a]
rs of
        [] -> Either String (Tree a)
e
        (a
r':[a]
rs') -> Either String (Tree a)
e Either String (Tree a)
-> (Tree a -> Either String (Tree a)) -> Either String (Tree a)
forall a b.
Either String a -> (a -> Either String b) -> Either String b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= NonEmpty a -> Tree a -> Either String (Tree a)
forall a. Ref a => NonEmpty a -> Tree a -> Either String (Tree a)
tryGetFirstPath (a
r' a -> [a] -> NonEmpty a
forall a. a -> [a] -> NonEmpty a
:| [a]
rs')   

getSubtreeIfRef :: (Ref a) => a -> Node a -> Maybe (Tree a)
getSubtreeIfRef :: forall a. Ref a => a -> Node a -> Maybe (Tree a)
getSubtreeIfRef a
ref (RefNode a
ref' Tree a
subtree) = if a -> a -> Bool
forall a. Ref a => a -> a -> Bool
Ref.eq a
ref a
ref' then Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just Tree a
subtree else Maybe (Tree a)
forall a. Maybe a
Nothing
getSubtreeIfRef a
_ Node a
_ = Maybe (Tree a)
forall a. Maybe a
Nothing

isStringEq :: (Ref a) => Node a -> ByteString -> Bool
isStringEq :: forall a. Ref a => Node a -> ByteString -> Bool
isStringEq (StringNode ByteString
bs) ByteString
bs' = ByteString
bs ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
bs' 
isStringEq Node a
_ ByteString
_ = Bool
False

isRefOf :: (Ref a) => Node a -> a -> Bool 
isRefOf :: forall a. Ref a => Node a -> a -> Bool
isRefOf (RefNode a
ref Tree a
_) a
ref' = a -> a -> Bool
forall a. Ref a => a -> a -> Bool
Ref.eq a
ref a
ref'
isRefOf Node a
_ a
_ = Bool
False