content: {-# LANGUAGE DeriveAnyClass  #-}
  {-# LANGUAGE PatternSynonyms #-}
  
  {- |
  Module                  : Toml.Type.PrefixTree
  Copyright               : (c) 2018-2022 Kowainik
  SPDX-License-Identifier : MPL-2.0
  Maintainer              : Kowainik <xrom.xkov@gmail.com>
  Stability               : Stable
  Portability             : Portable
  
  Implementation of prefix tree for TOML AST.
  
  @since 0.0.0
  -}
  
  module Toml.Type.PrefixTree
      (
        -- * Non-empty prefix tree
        PrefixTree (..)
      , singleT
      , insertT
      , lookupT
      , toListT
      , addPrefixT
      , differenceWithT
  
        -- * Prefix map that stores roots of 'PrefixTree'
      , PrefixMap
      , single
      , insert
      , lookup
      , fromList
      , toList
      , differenceWith
      ) where
  
  import Prelude hiding (lookup)
  
  import Control.DeepSeq (NFData)
  import Data.Bifunctor (first)
  import Data.Foldable (foldl')
  import Data.HashMap.Strict (HashMap)
  import GHC.Generics (Generic)
  
  import Toml.Type.Key (pattern (:||), Key, KeysDiff (..), Piece, Prefix, keysDiff, (<|))
  
  import qualified Data.HashMap.Strict as HashMap
  
  
  {- | Map of layer names and corresponding 'PrefixTree's.
  
  @since 0.0.0
  -}
  type PrefixMap a = HashMap Piece (PrefixTree a)
  
  {- | Data structure to represent table tree for @toml@.
  
  @since 0.0.0
  -}
  data PrefixTree a
      = Leaf             -- ^ End of a key.
          !Key           -- ^ End piece of the key.
          !a             -- ^ Value at the end.
      | Branch           -- ^ Values along pieces of a key.
          !Prefix        -- ^ Greatest common key prefix.
          !(Maybe a)     -- ^ Possible value at that point.
          !(PrefixMap a) -- ^ Values at suffixes of the prefix.
      deriving stock (Show, Eq, Generic)
      deriving anyclass (NFData)
  
  -- | @since 0.3
  instance Semigroup (PrefixTree a) where
      a <> b = foldl' (\tree (k, v) -> insertT k v tree) a (toListT b)
  
  {- | Push 'Prefix' inside the given 'PrefixTree'.
  
  @since 1.3.2.0
  -}
  addPrefixT :: Prefix -> PrefixTree a -> PrefixTree a
  addPrefixT pref = \case
      Leaf k a -> Leaf (pref <> k) a
      Branch k ma pma -> Branch (pref <> k) ma pma
  
  {- | Convert branches to 'Leaf' or remove them at all.
  
  @since 1.3.2.0
  -}
  compressTree :: PrefixTree a -> Maybe (PrefixTree a)
  compressTree = \case
      l@(Leaf _ _) -> Just l
      b@(Branch p ma pma) -> case HashMap.toList pma of
          [] -> ma >>= \a -> Just (Leaf p a)
          [(_, child)] -> case ma of
              Just _ -> Just b
              Nothing -> compressTree $ addPrefixT p child
          _ : _ : _ -> Just b
  
  {- | Creates a 'PrefixTree' of one key-value element.
  
  @since 0.0.0
  -}
  singleT :: Key -> a -> PrefixTree a
  singleT = Leaf
  {-# INLINE singleT #-}
  
  {- | Creates a 'PrefixMap' of one key-value element.
  
  @since 0.0.0
  -}
  single :: Key -> a -> PrefixMap a
  single k@(p :|| _) = HashMap.singleton p . singleT k
  
  {- | Inserts key-value element into the given 'PrefixTree'.
  
  @since 0.0.0
  -}
  insertT :: Key -> a -> PrefixTree a -> PrefixTree a
  insertT newK newV (Leaf k v) =
      case keysDiff k newK of
          Equal -> Leaf k newV
          NoPrefix -> error "Algorithm error: can't be equal prefixes in insertT:Leaf case"
          FstIsPref rK -> Branch k (Just v) $ single rK newV
          SndIsPref lK -> Branch newK (Just newV) $ single lK v
          Diff p k1@(f :|| _) k2@(s :|| _) ->
            Branch p Nothing $ HashMap.fromList [(f, Leaf k1 v), (s, Leaf k2 newV)]
  insertT newK newV (Branch pref mv prefMap) =
      case keysDiff pref newK of
          Equal    -> Branch pref (Just newV) prefMap
          NoPrefix -> error "Algorithm error: can't be equal prefixes in insertT:Branch case"
          FstIsPref rK -> Branch pref mv $ insert rK newV prefMap
          SndIsPref lK@(piece :|| _) ->
              Branch newK (Just newV) $ HashMap.singleton piece (Branch lK mv prefMap)
          Diff p k1@(f :|| _) k2@(s :|| _) ->
              Branch p Nothing $ HashMap.fromList [ (f, Branch k1 mv prefMap)
                                                  , (s, Leaf k2 newV)
                                                  ]
  
  {- | Inserts key-value element into the given 'PrefixMap'.
  
  @since 0.0.0
  -}
  insert :: Key -> a -> PrefixMap a -> PrefixMap a
  insert k@(p :|| _) v prefMap = case HashMap.lookup p prefMap of
      Just tree -> HashMap.insert p (insertT k v tree) prefMap
      Nothing   -> HashMap.insert p (singleT k v) prefMap
  
  {- | Looks up the value at a key in the 'PrefixTree'.
  
  @since 0.0.0
  -}
  lookupT :: Key -> PrefixTree a -> Maybe a
  lookupT lk (Leaf k v) = if lk == k then Just v else Nothing
  lookupT lk (Branch pref mv prefMap) =
      case keysDiff pref lk of
          Equal       -> mv
          NoPrefix    -> Nothing
          Diff _ _ _  -> Nothing
          SndIsPref _ -> Nothing
          FstIsPref k -> lookup k prefMap
  
  {- | Looks up the value at a key in the 'PrefixMap'.
  
  @since 0.0.0
  -}
  lookup :: Key -> PrefixMap a -> Maybe a
  lookup k@(p :|| _) prefMap = HashMap.lookup p prefMap >>= lookupT k
  
  {- | Constructs 'PrefixMap' structure from the given list of 'Key' and value pairs.
  
  @since 0.0.0
  -}
  fromList :: [(Key, a)] -> PrefixMap a
  fromList = foldl' insertPair mempty
    where
      insertPair :: PrefixMap a -> (Key, a) -> PrefixMap a
      insertPair prefMap (k, v) = insert k v prefMap
  
  {- | Converts 'PrefixTree' to the list of pairs.
  
  @since 0.0.0
  -}
  toListT :: PrefixTree a -> [(Key, a)]
  toListT (Leaf k v) = [(k, v)]
  toListT (Branch pref ma prefMap) = case ma of
      Just a  -> (:) (pref, a)
      Nothing -> id
      $ map (\(k, v) -> (pref <> k, v)) $ toList prefMap
  
  {- | Converts 'PrefixMap' to the list of pairs.
  
  @since 0.0.0
  -}
  toList :: PrefixMap a -> [(Key, a)]
  toList = concatMap (\(p, tr) -> first (p <|) <$> toListT tr) . HashMap.toList
  
  {- | Difference of two 'PrefixMap's. Returns elements of the first 'PrefixMap'
  that are not existing in the second one.
  
  @since 1.3.2.0
  -}
  differenceWith :: (a -> b -> Maybe a) -> PrefixMap a -> PrefixMap b -> PrefixMap a
  differenceWith f = HashMap.differenceWith (differenceWithT f)
  
  {- | Difference of two 'PrefixTree's. Returns elements of the first 'PrefixTree'
  that are not existing in the second one.
  
  @since 1.3.2.0
  -}
  differenceWithT :: (a -> b -> Maybe a) -> PrefixTree a -> PrefixTree b -> Maybe (PrefixTree a)
  differenceWithT f pt1 pt2 = case (pt1, pt2) of
      (Leaf k1 a, Leaf k2 b)
          | k1 == k2 -> f a b >>= \aNew -> Just (Leaf k1 aNew)
          | otherwise -> Just (Leaf k1 a)
  
      (l@(Leaf k a), Branch p mb pmb) -> case keysDiff k p of
          Equal -> mb >>= f a >>= \aNew -> Just (Leaf k aNew)
          NoPrefix -> Just l
          FstIsPref _ -> Just l
          SndIsPref kSuf -> case HashMap.toList $ differenceWith f (single kSuf a) pmb of
              -- zero elements
              [] -> Nothing
              -- our single key
              [(_, aNew)] -> Just $ addPrefixT k aNew
              -- shouldn't happen, but for some reasons
              _ : _ : _ -> Nothing
          Diff _ _ _ -> Just l
  
      (br@(Branch p ma pma), Leaf k b) -> case keysDiff p k of
          Equal -> compressTree $ Branch p (ma >>= \a -> f a b) pma
          NoPrefix -> Just br
          FstIsPref kSuf -> compressTree $ Branch p ma (differenceWith f pma $ single kSuf b)
          SndIsPref _ -> Just br
          Diff _ _ _ -> Just br
  
      (b1@(Branch p1 ma pma), Branch p2 mb pmb) -> case keysDiff p1 p2 of
          Equal -> compressTree $
              Branch p1 (ma >>= \a -> mb >>= \b -> f a b) (differenceWith f pma pmb)
          NoPrefix -> Just b1
          FstIsPref p2Suf@(p2Head :|| _) -> compressTree $
              Branch p1 ma (differenceWith f pma $ HashMap.singleton p2Head $ Branch p2Suf mb pmb)
          SndIsPref p1Suf@(p1Head :|| _) -> case HashMap.lookup p1Head pmb of
              Nothing -> Just b1
              Just ch -> addPrefixT p2 <$> differenceWithT f (Branch p1Suf ma pma) ch
          Diff _ _ _ -> Just b1
  