-- File created: 2008-12-28 17:20:14

{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies
           , FlexibleContexts, ScopedTypeVariables, Rank2Types
           , KindSignatures, ViewPatterns #-}

module Data.ListTrie.Patricia.Vector.Base
   ( Trie(..)
   , split, splitLookup
   , mapKeysWith, mapInKeysWith, mapInKeysWith'
   , foldrWithKey,  foldrAscWithKey,  foldrDescWithKey
   , foldlWithKey,  foldlAscWithKey,  foldlDescWithKey
   , foldlWithKey', foldlAscWithKey', foldlDescWithKey'
   , toList, toAscList, toDescList
   , fromList, fromListWith, fromListWith', fromListWithKey, fromListWithKey'
   , findMin, findMax, deleteMin, deleteMax, minView, maxView
   , findPredecessor, findSuccessor
   , lookupPrefix, addPrefix, deletePrefix, deleteSuffixes
   , splitPrefix, children, children1
   , showTrieWith
   , eqComparePrefixes, ordComparePrefixes
   ) where

import Control.Applicative (Applicative(..), (<$>))
import Control.Arrow       ((***), first)
import Control.Exception   (assert)
import qualified Data.DList as DL
import Data.DList          (DList)
import Data.Foldable       (foldr, foldl')
import Data.List           (partition)
import Data.Maybe          (fromJust, isJust)
import Data.Monoid         ((<>))
import Prelude hiding      (lookup, filter, foldr, null)
import qualified Prelude

import qualified Data.ListTrie.Base.Map.Internal as Map
import Data.ListTrie.Base.Classes
   ( Boolable(..)
   , Unwrappable(..)
   , Unionable(..), Differentiable(..), Intersectable(..)
   , Alt(..)
   , fmap', (<$!>)
   )

import qualified Data.Vector.Generic as VG
import Data.ListTrie.Base.Map (Map, OrdMap)
import Data.ListTrie.Util     ((.:), both, over_3)

type CMap trie map (v :: * -> *) k a = map k (trie map v k a)

class (VG.Vector v k, Monoid (v k), Map map k, Functor st, Unwrappable st, VG.Vector v k)
  => Trie trie st map v k | trie -> st where

  mkTrie :: st a -> v k -> CMap trie map v k a -> trie map v k a
  tParts :: trie map v k a -> (st a, v k, CMap trie map v k a)

  foldTrie :: Boolable (st a)
           => (st a -> v k -> CMap trie map v k a -> (st a, v k, CMap trie map v k a))
           -> trie map v k a -> trie map v k a
  foldTrie f =
    tryCompress . (uncurry2 mkTrie) . (uncurry2 f) . (over_3 (Map.map (foldTrie f))) . tParts
    where
      uncurry2 :: (a -> b -> c -> d) -> (a, b, c) -> d
      uncurry2 f' (a, b, c) = f' a b c

  -----------------------

  -- * Construction

  -- O(1)
  empty :: (Alt st a) => trie map v k a
  empty = mkTrie altEmpty mempty Map.empty

  -- O(1)
  singleton :: (Alt st a) => v k -> a -> trie map v k a
  singleton k v = mkTrie (pure v) k Map.empty

  -- O(min(m,s))
  insert :: (Alt st a, Boolable (st a))
         => v k -> a -> trie map v k a -> trie map v k a
  insert = insertWith const

  -- O(min(m,s))
  insert' :: (Alt st a, Boolable (st a))
          => v k -> a -> trie map v k a -> trie map v k a
  insert' = insertWith' const

  -- O(min(m,s))
  insertWith :: (Alt st a, Boolable (st a))
             => (a -> a -> a) -> v k -> a -> trie map v k a -> trie map v k a
  insertWith = genericInsertWith ($) (<$>)

  -- O(min(m,s))
  insertWith' :: (Alt st a, Boolable (st a))
              => (a -> a -> a) -> v k -> a -> trie map v k a -> trie map v k a
  insertWith' = (seq <*>) .: genericInsertWith ($!) (<$!>)

  -- O(min(m,s))
  delete :: (Alt st a, Boolable (st a))
         => v k -> trie map v k a -> trie map v k a
  delete = alter (const altEmpty)

  -- O(min(m,s))
  adjust :: (a -> a) -> v k -> trie map v k a -> trie map v k a
  adjust = genericAdjust ($) fmap

  -- O(min(m,s))
  adjust' :: (Alt st a, Boolable (st a))
          => (a -> a) -> v k -> trie map v k a -> trie map v k a
  adjust' = genericAdjust ($!) fmap'

  -- O(min(m,s))
  updateLookup :: (Alt st a, Boolable (st a))
               => (a -> st a) -> v k -> trie map v k a -> (st a, trie map v k a)
  updateLookup f = go
    where
      go k tr =
        let (v,prefix,m) = tParts tr
        in case comparePrefixes (Map.eqCmp m) prefix k of
             Same                   -> let v' = if hasValue v
                                               then f (unwrap v)
                                               else v
                                      in (v, safeMkTrie v' prefix m)
             PostFix (Right (unconsVector -> (Just x, xs))) ->
               case Map.lookup x m of
                 Nothing  -> (altEmpty, tr)
                 Just tr' ->
                   let (ret, upd) = go xs tr'
                   in ( ret
                      , safeMkTrie v prefix $
                        if null upd
                        then Map.delete x m
                        else Map.adjust (const upd) x m
                      )
             _ -> (altEmpty, tr)

  -- O(min(m,s))
  --
  -- This can be lazy in exactly one case: the key is a prefix of more than one
  -- key in the trie. In that case, we know that the resulting trie continues to
  -- contain those children.
  --
  -- In all other cases we have to check whether the function removed a key or
  -- not, in order to be able to keep the trie in an internally valid state.

  -- (I.e. we need to try to compress it.)
  alter :: (Alt st a, Boolable (st a))
        => (st a -> st a) -> v k -> trie map v k a -> trie map v k a
  alter = genericAlter (flip const)

  -- O(min(m,s))
  alter' :: (Alt st a, Boolable (st a))
         => (st a -> st a) -> v k -> trie map v k a -> trie map v k a
  alter' = genericAlter seq

  -- * Querying

  -- O(1)
  --
  -- Test the strict field last for maximal laziness
  null :: (Boolable (st a)) => trie map v k a -> Bool
  null tr = let (v,p,m) = tParts tr
            in Map.null m && noValue v && assert (VG.null p) True

  -- O(n m)
  size :: (Boolable (st a), Num n) => trie map v k a -> n
  size  tr = foldr  ((+) . size)  (if hasValue (tVal tr) then 1 else 0) (tMap tr)

  -- O(n m)
  size' :: (Boolable (st a), Num n) => trie map v k a -> n
  size' tr = foldl' (flip $ (+) . size')
             (if hasValue (tVal tr) then 1 else 0)
             (tMap tr)

  -- O(min(m,s))
  member :: (Alt st a, Boolable (st a))
         => v k -> trie map v k a -> Bool
  member = hasValue .: lookup

  -- O(min(m,s))
  notMember :: (Alt st a, Boolable (st a))
            => v k -> trie map v k a -> Bool
  notMember = not .: member

  -- O(min(m,s))
  lookup :: (Alt st a) => v k -> trie map v k a -> st a
  lookup k tr =
    let (v,prefix,m) = tParts tr
    in case comparePrefixes (Map.eqCmp m) prefix k of
         Same                   -> v
         PostFix (Right (unconsVector -> (Just x, xs))) ->
           maybe altEmpty (lookup xs) (Map.lookup x m)
         _                      -> altEmpty

  -- O(min(m,s))
  lookupWithDefault :: (Alt st a)
                    => a -> v k -> trie map v k a -> a
  lookupWithDefault def k tr = unwrap $ lookup k tr <|> pure def

  -- O(min(n1 m1,n2 m2))
  isSubmapOfBy :: (Boolable (st a), Boolable (st b))
               => (a -> b -> Bool)
               -> trie map v k a
               -> trie map v k b
               -> Bool
  isSubmapOfBy f = go0
    where
      go0 trl trr =
        let (vl,prel,ml) = tParts trl
            (vr,prer,mr) = tParts trr
        in case comparePrefixes (Map.eqCmp ml) prel prer of
             DifferedAt _ _ _  -> False
             -- Special case here: if the left trie is empty we return True.
             PostFix (Right _) -> null trl
             PostFix (Left xs) -> go mr vl ml xs
             Same              -> same vl vr ml mr

      go mr vl ml (unconsVector -> (Just x, xs)) =
        case Map.lookup x mr of
          Nothing -> False
          Just tr ->
            let (vr,pre,mr') = tParts tr
            in case comparePrefixes (Map.eqCmp mr) xs pre of
                 DifferedAt _ _ _  -> False
                 PostFix (Right _) -> False
                 PostFix (Left ys) -> go mr' vl ml ys
                 Same              -> same vl vr ml mr'

      go _ _ _ _ =
        error "Data.ListTrie.Patricia.Base.isSubmapOfBy :: internal error"

      same vl vr ml mr =
        let hvl = hasValue vl
            hvr = hasValue vr
        in and [ not (hvl && not hvr)
               , (not hvl && not hvr) || f (unwrap vl) (unwrap vr)
               , Map.isSubmapOfBy go0 ml mr
               ]

  -- O(min(n1 m1,n2 m2))
  isProperSubmapOfBy :: (Boolable (st a), Boolable (st b))
                     => (a -> b -> Bool)
                     -> trie map v k a
                     -> trie map v k b
                     -> Bool
  isProperSubmapOfBy g = f False
    where
      f proper trl trr =
        let (vl,prel,ml) = tParts trl
            (vr,prer,mr) = tParts trr
        in case comparePrefixes (Map.eqCmp ml) prel prer of
             DifferedAt _ _ _  -> False
             -- Special case, as in isSubsetOf.
             --
             -- Note that properness does not affect this: if we hit this
             -- case, we already know that the right trie is nonempty.
             PostFix (Right _) -> null trl
             PostFix (Left xs) -> go proper mr vl ml xs
             Same              -> same proper vl vr ml mr

      go proper mr vl ml (unconsVector -> (Just x, xs)) =
        case Map.lookup x mr of
          Nothing -> False
          Just tr ->
            let (vr,pre,mr') = tParts tr
            in case comparePrefixes (Map.eqCmp mr) xs pre of
                 DifferedAt _ _ _  -> False
                 PostFix (Right _) -> False
                 PostFix (Left ys) -> go proper mr' vl ml ys
                 Same              -> same proper vl vr ml mr'

      go _ _ _ _ _ =
        error "Data.ListTrie.Patricia.Base.isProperSubmapOfBy :: internal error"

      same proper vl vr ml mr =
        let hvl = hasValue vl
            hvr = hasValue vr

            -- As the non-Patricia version, so does this seem suboptimal.
            proper' = or [ proper
                         , not hvl && hvr
                         , not (Map.null $ Map.difference mr ml)
                         ]

        in and [ not (hvl && not hvr)
               , (not hvl && not hvr) || g (unwrap vl) (unwrap vr)
               , if Map.null ml
                 then proper'
                 else Map.isSubmapOfBy (f proper') ml mr
               ]

  -- * Combination

  -- The *Key versions are mostly rewritten from the basic ones: they have an
  -- additional O(m) cost from keeping track of the key, which is why the basic
  -- ones can't just call them.

  -- O(min(n1 m1,n2 m2))
  unionWith :: (Alt st a, Boolable (st a), Unionable st a)
            => (a -> a -> a) -> trie map v k a -> trie map v k a -> trie map v k a
  unionWith f = genericUnionWith (flip const) (unionVals f)

  -- O(min(n1 m1,n2 m2))
  unionWith' :: (Alt st a, Boolable (st a), Unionable st a)
            => (a -> a -> a) -> trie map v k a -> trie map v k a -> trie map v k a
  unionWith' f = genericUnionWith seq (unionVals' f)

  -- O(min(n1 m1,n2 m2))
  unionWithKey :: (Alt st a, Boolable (st a), Unionable st a)
               => (v k -> a -> a -> a)
               -> trie map v k a
               -> trie map v k a
               -> trie map v k a
  unionWithKey = genericUnionWithKey (flip const) unionVals

  -- O(min(n1 m1,n2 m2))
  unionWithKey' :: ( Alt st a, Boolable (st a), Unionable st a)
                => (v k -> a -> a -> a)
                -> trie map v k a
                -> trie map v k a
                -> trie map v k a
  unionWithKey' = genericUnionWithKey seq unionVals'

  -- O(sum(n))
  unionsWith :: (Alt st a, Boolable (st a), Unionable st a)
             => (a -> a -> a) -> [trie map v k a] -> trie map v k a
  unionsWith j = foldl' (unionWith j) empty

  -- O(sum(n))
  unionsWith' :: (Alt st a, Boolable (st a), Unionable st a)
              => (a -> a -> a) -> [trie map v k a] -> trie map v k a
  unionsWith' j = foldl' (unionWith' j) empty

  -- O(sum(n))
  unionsWithKey :: ( Alt st a, Boolable (st a), Unionable st a)
                => (v k -> a -> a -> a) -> [trie map v k a] -> trie map v k a
  unionsWithKey j = foldl' (unionWithKey j) empty

  -- O(sum(n))
  unionsWithKey' :: ( Alt st a, Boolable (st a)
                   , Unionable st a, Trie trie st map v k
                   )
                 => (v k -> a -> a -> a) -> [trie map v k a] -> trie map v k a
  unionsWithKey' j = foldl' (unionWithKey' j) empty

  -- O(min(n1 m1,n2 m2))
  differenceWith :: forall a b
                 . (Boolable (st a), Differentiable st a b)
                 => (a -> b -> Maybe a)
                 -> trie map v k a
                 -> trie map v k b
                 -> trie map v k a
  differenceWith j = go
    where
      go :: trie map v k a -> trie map v k b -> trie map v k a
      go tr1 tr2 =
        let (v1,pre1,m1) = tParts tr1
            (v2,pre2,m2) = tParts tr2
        in case comparePrefixes (Map.eqCmp m1) pre1 pre2 of
             DifferedAt _ _ _   -> tr1
             Same               -> mk v1 v2 pre1 m1 m2
             PostFix (Left  xs) -> goRight tr1 m2  xs
             PostFix (Right xs) -> goLeft  tr1 tr2 xs

      dw :: trie map v k a -> trie map v k b -> Maybe (trie map v k a)
      dw a b =
        let c = differenceWith j a b
        in if null c then Nothing else Just c

      mk v v' p m m' =
        let vd = differenceVals j v v'
        in tryCompress.mkTrie vd p $ Map.differenceWith dw m m'

      -- See the comment in 'intersection' for a longish example of the idea
      -- behind this, which is basically that if we see two prefixes like "foo"
      -- and "foobar", we traverse the "foo" trie looking for "bar". Then if we
      -- find "barbaz", we traverse the "foobar" trie looking for "baz", and so
      -- on.
      --
      -- We have two functions for the two tries because set difference is a
      -- noncommutative operation.
      goRight left rightMap (unconsVector -> (Just x, xs)) =
        let (v,pre,m) = tParts left
        in case Map.lookup x rightMap of
          Nothing     -> left
          Just right' ->
            let (v',pre',m') = tParts right'
            in case comparePrefixes (Map.eqCmp m) xs pre' of
                 DifferedAt _ _ _   -> left
                 Same               -> mk v v' pre m m'
                 PostFix (Left  ys) -> goRight left m'     ys
                 PostFix (Right ys) -> goLeft  left right' ys

      goRight _ _ _ = can'tHappen

      goLeft left right (unconsVector -> (Just x, xs)) =
        tryCompress . mkTrie vl prel $ Map.update f x ml
        where
          (vl,prel,ml) = tParts left
          (vr,   _,mr) = tParts right

          f left' =
            let (v,pre,m) = tParts left'
            in case comparePrefixes (Map.eqCmp m) pre xs of
                 DifferedAt _ _ _   -> Just left'
                 Same               -> tryNull $ mk v vr pre m mr
                 PostFix (Left  ys) -> tryNull $ goRight left' mr    ys
                 PostFix (Right ys) -> tryNull $ goLeft  left' right ys
      goLeft _ _ _ = can'tHappen

      tryNull t = if null t then Nothing else Just t

      can'tHappen = error "Data.ListTrie.Patricia.Base.differenceWith :: internal error"

  -- O(min(n1 m1,n2 m2))
  differenceWithKey :: ( Boolable (st a), Differentiable st a b)
                    => (v k -> a -> b -> Maybe a)
                    -> trie map v k a
                    -> trie map v k b
                    -> trie map v k a
  differenceWithKey j = go DL.empty
    where
      go k tr1 tr2 =
        let (v1,pre1,m1) = tParts tr1
            (v2,pre2,m2) = tParts tr2
        in case comparePrefixes (Map.eqCmp m1) pre1 pre2 of
             DifferedAt _ _ _   -> tr1
             Same               -> mk k v1 v2 pre1 m1 m2
             PostFix (Left  xs) -> goRight (key k pre2) tr1 m2  xs
             PostFix (Right xs) -> goLeft  (key k pre1) tr1 tr2 xs

      key k p = k `DL.append` DL.fromList (VG.toList p)

      dw k a b =
        let c = go k a b
        in if null c then Nothing else Just c

      mk k v v' p m m' =
        let k' = k `DL.append` DL.fromList (VG.toList p)
            vd = differenceVals (j $ (VG.fromList $ DL.toList k')) v v'
        in tryCompress.mkTrie vd p $
           Map.differenceWithKey (dw . (k' `DL.snoc`)) m m'

      goRight k left rightMap (unconsVector -> (Just x, xs)) =
        let (vl,_,ml) = tParts left
        in case Map.lookup x rightMap of
             Nothing    -> left
             Just right ->
               let (vr,pre,mr) = tParts right
                   k'          = k `DL.snoc` x
               in case comparePrefixes (Map.eqCmp ml) xs pre of
                    DifferedAt _ _ _   -> left
                    Same               -> mk k' vl vr pre ml mr
                    PostFix (Left  ys) -> goRight (key k' pre)
                                         left mr    ys
                    PostFix (Right ys) -> goLeft  (key k' xs)
                                         left right ys

      goRight _ _ _ _ = can'tHappen

      goLeft k left right (unconsVector -> (Just x, xs)) =
        tryCompress . mkTrie vl prel $ Map.update f x ml
        where
          (vl,prel,ml) = tParts left
          (vr,   _,mr) = tParts right

          k' = k `DL.snoc` x

          f left' =
            let (v,pre,m) = tParts left'
            in case comparePrefixes (Map.eqCmp m) pre xs of
                 DifferedAt _ _ _   -> Just left'
                 Same               -> tryNull $ mk k' v vr pre m mr
                 PostFix (Left  ys) -> tryNull $ goRight (key k' xs)
                                      left' mr    ys
                 PostFix (Right ys) -> tryNull $ goLeft  (key k' pre)
                                      left' right ys
      goLeft _ _ _ _ = can'tHappen

      tryNull t = if null t then Nothing else Just t

      can'tHappen = error "Data.ListTrie.Patricia.Base.differenceWithKey :: internal error"

  -- O(min(n1 m1,n2 m2))
  intersectionWith :: ( Alt st c, Boolable (st c)
                     , Intersectable st a b c, Intersectable st b a c)
                   => (a -> b -> c)
                   -> trie map v k a
                   -> trie map v k b
                   -> trie map v k c
  intersectionWith f = genericIntersectionWith (flip const) (intersectionVals f)

  -- O(min(n1 m1,n2 m2))
  intersectionWith' :: ( Alt st c, Boolable (st c)
                       , Intersectable st a b c, Intersectable st b a c
                       , Trie trie st map v k
                       )
                    => (a -> b -> c)
                    -> trie map v k a
                    -> trie map v k b
                    -> trie map v k c
  intersectionWith' f = genericIntersectionWith seq (intersectionVals' f)

  -- O(min(n1 m1,n2 m2))
  intersectionWithKey :: ( Alt st c, Boolable (st c)
                         , Intersectable st a b c, Intersectable st b a c)
                      => (v k -> a -> b -> c)
                      -> trie map v k a
                      -> trie map v k b
                      -> trie map v k c
  intersectionWithKey = genericIntersectionWithKey (flip const) intersectionVals

  -- O(min(n1 m1,n2 m2))
  intersectionWithKey' :: ( Alt st c, Boolable (st c)
                         , Intersectable st a b c, Intersectable st b a c)
                       => (v k -> a -> b -> c)
                       -> trie map v k a
                       -> trie map v k b
                       -> trie map v k c
  intersectionWithKey' = genericIntersectionWithKey seq intersectionVals'


  -- * Filtering

  -- O(n m)
  filterWithKey :: (Alt st a, Boolable (st a))
                => (v k -> a -> Bool) -> trie map v k a -> trie map v k a
  filterWithKey p = fromList . Prelude.filter (uncurry p) . toList

  -- O(n m)
  partitionWithKey :: (Alt st a, Boolable (st a))
                   => (v k -> a -> Bool)
                   -> trie map v k a
                   -> (trie map v k a, trie map v k a)
  partitionWithKey p = both fromList . partition (uncurry p) . toList

  -- * Mapping

-- O(n m)
mapKeysWith :: (Boolable (st a), Trie trie st map v k1, Trie trie st map v k2)
            => ([(v k2,a)] -> trie map v k2 a)
            -> (v k1 -> v k2)
            -> trie map v k1 a
            -> trie map v k2 a
mapKeysWith fromlist f = fromlist . map (first f) . toList

-- O(n m)
mapInKeysWith :: ( Alt st a, Boolable (st a), Unionable st a
                , Trie trie st map v k1, Trie trie st map v k2
                )
              => (a -> a -> a)
              -> (k1 -> k2)
              -> trie map v k1 a
              -> trie map v k2 a
mapInKeysWith = genericMapInKeysWith (flip const) (const ()) unionWith

-- O(n m)
mapInKeysWith' :: ( Alt st a, Boolable (st a), Unionable st a
                  , Trie trie st map v k1, Trie trie st map v k2
                  )
               => (a -> a -> a)
               -> (k1 -> k2)
               -> trie map v k1 a
               -> trie map v k2 a
mapInKeysWith' =
   genericMapInKeysWith
      seq
      (`seq` ())
      unionWith'

genericMapInKeysWith :: ( Alt st a, Boolable (st a), Unionable st a
                        , Trie trie st map v k1, Trie trie st map v k2
                        )
                     => (() -> trie map v k2 a -> trie map v k2 a)
                     -> (v k2 -> ())
                     -> (f -> trie map v k2 a -> trie map v k2 a -> trie map v k2 a)
                     -> f
                     -> (k1 -> k2)
                     -> trie map v k1 a
                     -> trie map v k2 a
genericMapInKeysWith seeq listSeq unionW j f = go
 where
   go tr =
      let (v,p,m) = tParts tr
          p'      = VG.map f p
       in listSeq p' `seeq`
             (mkTrie v p' $
                 Map.fromListKVWith (unionW j) . map (f *** go) . Map.toListKV $ m)

-- * Folding

-- O(n m)
foldrWithKey :: (Boolable (st a), Trie trie st map v k)
            => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldrWithKey f x = foldr (uncurry f) x . toList

-- O(n m)
foldrAscWithKey :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
               => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldrAscWithKey f x = foldr (uncurry f) x . toAscList

-- O(n m)
foldrDescWithKey :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
                => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldrDescWithKey f x = foldr (uncurry f) x . toDescList

-- O(n m)
foldlWithKey :: (Boolable (st a), Trie trie st map v k)
             => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlWithKey f x = foldl (flip $ uncurry f) x . toList

-- O(n m)
foldlAscWithKey :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
                => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlAscWithKey f x = foldl (flip $ uncurry f) x . toAscList

-- O(n m)
foldlDescWithKey :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
                 => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlDescWithKey f x = foldl (flip $ uncurry f) x . toDescList

-- O(n m)
foldlWithKey' :: (Boolable (st a), Trie trie st map v k)
            => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlWithKey' f x = foldl' (flip $ uncurry f) x . toList

-- O(n m)
foldlAscWithKey' :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
               => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlAscWithKey' f x = foldl' (flip $ uncurry f) x . toAscList

-- O(n m)
foldlDescWithKey' :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
                => (v k -> a -> b -> b) -> b -> trie map v k a -> b
foldlDescWithKey' f x = foldl' (flip $ uncurry f) x . toDescList

-- * Conversion between lists

-- O(n m)
toList :: (Boolable (st a), Trie trie st map v k) => trie map v k a -> [(v k,a)]
toList = genericToList Map.toListKV DL.cons

-- O(n m)
toAscList :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
          => trie map v k a -> [(v k,a)]
toAscList = genericToList Map.toAscList DL.cons

-- O(n m)
toDescList :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
           => trie map v k a -> [(v k,a)]
toDescList = genericToList (reverse . Map.toAscList) (flip DL.snoc)

genericToList :: (Boolable (st a), Trie trie st map v k)
              => (CMap trie map v k a -> [(k, trie map v k a)])
              -> ((v k,a) -> DList (v k,a) -> DList (v k,a))
              -> trie map v k a
              -> [(v k,a)]
genericToList tolist add = DL.toList . go DL.empty
 where
   go l tr =
      let (v,p,m) = tParts tr
          l'      = l `DL.append` DL.fromList (VG.toList p)
          xs      =
             DL.concat .
             map (\(x,t) -> go (l' `DL.snoc` x) t) .
             tolist $ m
       in if hasValue v
             then add (VG.fromList $ DL.toList l', unwrap v) xs
             else                              xs

-- O(n m)
fromList :: (Alt st a, Boolable (st a), Trie trie st map v k)
         => [(v k,a)] -> trie map v k a
fromList = fromListWith const

-- O(n m)
fromListWith :: (Alt st a, Boolable (st a), Trie trie st map v k)
             => (a -> a -> a) -> [(v k,a)] -> trie map v k a
fromListWith f = foldl' (flip . uncurry $ insertWith f) empty

-- O(n m)
fromListWith' :: (Alt st a, Boolable (st a), Trie trie st map v k)
             => (a -> a -> a) -> [(v k,a)] -> trie map v k a
fromListWith' f = foldl' (flip . uncurry $ insertWith' f) empty

-- O(n m)
fromListWithKey :: (Alt st a, Boolable (st a), Trie trie st map v k)
                => (v k -> a -> a -> a) -> [(v k,a)] -> trie map v k a
fromListWithKey f = foldl' (\tr (k,v) -> insertWith (f k) k v tr) empty

-- O(n m)
fromListWithKey' :: (Alt st a, Boolable (st a), Trie trie st map v k)
                => (v k -> a -> a -> a) -> [(v k,a)] -> trie map v k a
fromListWithKey' f = foldl' (\tr (k,v) -> insertWith' (f k) k v tr) empty

-- * Min/max

-- O(m)
minView :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
        => trie map v k a -> (Maybe (v k, a), trie map v k a)
minView = minMaxView (hasValue.tVal) (fst . Map.minViewWithKey)

-- O(m)
maxView :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
        => trie map v k a -> (Maybe (v k, a), trie map v k a)
maxView = minMaxView (Map.null.tMap) (fst . Map.maxViewWithKey)

minMaxView :: (Alt st a, Boolable (st a), Trie trie st map v k)
           => (trie map v k a -> Bool)
           -> (CMap trie map v k a -> Maybe (k, trie map v k a))
           -> trie map v k a
           -> (Maybe (v k, a), trie map v k a)
minMaxView _        _       tr_ | null tr_ = (Nothing, tr_)
minMaxView isWanted mapView tr_ = first Just (go tr_)
 where
   go tr =
      let (v,pre,m) = tParts tr
       in if isWanted tr
             then ((pre, unwrap v), safeMkTrie altEmpty pre m)

             else let (k,      tr')  = fromJust (mapView m)
                      (minMax, tr'') = go tr'
                   in ( first (prepend pre k) minMax
                      , mkTrie v pre $ if null tr''
                                          then Map.delete              k m
                                          else Map.adjust (const tr'') k m
                      )

-- O(m)
findMin :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
        => trie map v k a -> Maybe (v k, a)
findMin = findMinMax (hasValue . tVal) (fst . Map.minViewWithKey)

-- O(m)
findMax :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
        => trie map v k a -> Maybe (v k, a)
findMax = findMinMax (Map.null . tMap) (fst . Map.maxViewWithKey)

findMinMax :: (Boolable (st a), Trie trie st map v k)
           => (trie map v k a -> Bool)
           -> (CMap trie map v k a -> Maybe (k, trie map v k a))
           -> trie map v k a
           -> Maybe (v k, a)
findMinMax _        _       tr_ | null tr_ = Nothing
findMinMax isWanted mapView tr_ = fmap (first VG.fromList) $ Just (go DL.empty tr_)
 where
   go xs tr =
      let (v,pre,m) = tParts tr
          xs'       = xs `DL.append` DL.fromList (VG.toList pre)
       in if isWanted tr
             then (DL.toList xs', unwrap v)
             else let (k, tr') = fromJust . mapView $ m
                   in go (xs' `DL.snoc` k) tr'

-- O(m)
deleteMin :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
          => trie map v k a -> trie map v k a
deleteMin = snd . minView

-- O(m)
deleteMax :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
          => trie map v k a -> trie map v k a
deleteMax = snd . maxView

-- O(min(m,s))
split :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
      => v k -> trie map v k a -> (trie map v k a, trie map v k a)
split xs tr = let (l,_,g) = splitLookup xs tr in (l,g)

-- O(min(m,s))
splitLookup :: (Alt st a, Boolable (st a), Trie trie st map v k, OrdMap map k)
            => v k
            -> trie map v k a
            -> (trie map v k a, st a, trie map v k a)
splitLookup xs tr =
   let (v,pre,m) = tParts tr
    in case comparePrefixes (Map.eqCmp m) pre xs of
            Same                     -> (empty, v, mk altEmpty pre m)
            DifferedAt _ (unconsVector -> (Just p, _)) (unconsVector -> (Just x, _)) ->
               case Map.ordCmp m p x of
                    LT -> (tr, altEmpty, empty)
                    GT -> (empty, altEmpty, tr)
                    EQ -> can'tHappen

            PostFix (Left  _)      -> (empty, altEmpty, tr)
            PostFix (Right (unconsVector -> (Just y, ys))) ->
               let (ml, maybeTr, mg) = Map.splitLookup y m
                in case maybeTr of
                        -- Prefix goes in left side of split since it's shorter
                        -- than the given key and thus lesser
                        Nothing  -> (mk v pre ml, altEmpty, mk altEmpty pre mg)
                        Just tr' ->
                           let (tl, v', tg) = splitLookup ys tr'
                               ml' = if null tl then ml else Map.insert y tl ml
                               mg' = if null tg then mg else Map.insert y tg mg
                            in (mk v pre ml', v', mk altEmpty pre mg')
            _ -> can'tHappen
 where
   mk v pre = tryCompress . mkTrie v pre
   can'tHappen =
      error "Data.ListTrie.Patricia.Base.splitLookup :: internal error"

-- O(m)
findPredecessor :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
                => v k -> trie map v k a -> Maybe (v k, a)
findPredecessor _   tr | null tr = Nothing
findPredecessor xs_ tr_          = go xs_ tr_
 where
   go xs tr =
      let (v,pre,m) = tParts tr
       in case comparePrefixes (Map.eqCmp m) pre xs of
               Same             -> Nothing
               PostFix (Left _) -> Nothing

               DifferedAt _ (unconsVector -> (Just p, _)) (unconsVector -> (Just x, _)) ->
                  case Map.ordCmp m p x of
                       LT -> findMax tr
                       GT -> Nothing
                       EQ -> can'tHappen

               -- See comment in non-Patricia version for explanation of
               -- algorithm
               PostFix (Right (unconsVector -> (Just y, ys))) ->
                  let predecessor = Map.findPredecessor y m
                   in (first (prepend pre y)<$>(Map.lookup y m >>= go ys))
                      <|>
                      case predecessor of
                           Nothing         ->
                              if hasValue v
                                 then Just (pre, unwrap v)
                                 else Nothing
                           Just (best,btr) ->
                              first (prepend pre best) <$> findMax btr
               _ -> can'tHappen

   can'tHappen =
      error "Data.ListTrie.Patricia.Base.findPredecessor :: internal error"

-- O(m)
findSuccessor :: forall trie map st v k a
               . (Boolable (st a), Trie trie st map v k, OrdMap map k)
              => v k -> trie map v k a -> Maybe (v k, a)
findSuccessor _   tr | null tr = Nothing
findSuccessor xs_ tr_          = go xs_ tr_
 where
   go :: (Boolable (st a), Trie trie st map v k, OrdMap map k)
      => v k -> trie map v k a -> Maybe (v k, a)
   go xs tr =
      let (_,pre,m) = tParts tr
       in case comparePrefixes (Map.eqCmp m) pre xs of
               Same -> do (k,t) <- fst $ Map.minViewWithKey m
                          first (prepend pre k) <$> findMin t

               DifferedAt _ (unconsVector -> (Just p, _)) (unconsVector -> (Just x, _)) ->
                  case Map.ordCmp m p x of
                       LT -> Nothing
                       GT -> findMin tr
                       EQ -> can'tHappen

               PostFix (Left _)       -> findMin tr
               PostFix (Right (unconsVector -> (Just y, ys))) ->
                  let successor = Map.findSuccessor y m
                   in (first (prepend pre y)<$>(Map.lookup y m >>= go ys))
                      <|>
                      (successor >>= \(best,btr) ->
                         first (prepend pre best) <$> findMin btr)

               _ -> can'tHappen

   can'tHappen =
      error "Data.ListTrie.Patricia.Base.findSuccessor :: internal error"

-- * Trie-only operations

-- O(s)
lookupPrefix :: (Alt st a, Boolable (st a), Trie trie st map v k)
             => v k -> trie map v k a -> trie map v k a
lookupPrefix xs tr =
   let (_,pre,m) = tParts tr
    in case comparePrefixes (Map.eqCmp m) pre xs of
            DifferedAt _ _ _       -> empty
            Same                   -> tr
            PostFix (Left _)       -> tr
            PostFix (Right (unconsVector -> (Just y, ys))) ->
               case Map.lookup y m of
                    Nothing  -> empty
                    Just tr' -> let tr''         = lookupPrefix ys tr'
                                    (v',pre',m') = tParts tr''
                                 in if null tr''
                                       then tr''
                                       else mkTrie v' (pre <> (VG.cons y pre')) m'
            _ ->
               error
                  "Data.ListTrie.Patricia.Base.lookupPrefix :: internal error"

-- O(s)
addPrefix :: (Alt st a, Trie trie st map v k)
          => v k -> trie map v k a -> trie map v k a
addPrefix xs tr =
   let (v,pre,m) = tParts tr
    in mkTrie v (xs <> pre) m

-- O(s)
deletePrefix :: (Alt st a, Boolable (st a), Trie trie st map v k)
             => v k -> trie map v k a -> trie map v k a
deletePrefix xs tr =
   let (v,pre,m) = tParts tr
    in case comparePrefixes (Map.eqCmp m) pre xs of
            Same                   -> tryCompress (mkTrie v mempty m)
            PostFix (Left ys)      -> mkTrie v ys m
            DifferedAt _ _ _       -> empty
            PostFix (Right (unconsVector -> (Just y, ys))) ->
               case Map.lookup y m of
                    Nothing  -> empty
                    Just tr' -> deletePrefix ys tr'

            _ ->
               error
                  "Data.ListTrie.Patricia.Base.deletePrefix :: internal error"

-- O(s)
deleteSuffixes :: (Alt st a, Boolable (st a), Trie trie st map v k)
               => v k -> trie map v k a -> trie map v k a
deleteSuffixes xs tr =
   let (v,pre,m) = tParts tr
    in case comparePrefixes (Map.eqCmp m) pre xs of
            DifferedAt _ _ _       -> tr
            Same                   -> empty
            PostFix (Left _)       -> empty
            PostFix (Right (unconsVector -> (Just y, ys))) ->
               case Map.lookup y m of
                    Nothing  -> tr
                    Just tr' ->
                       let tr'' = deleteSuffixes ys tr'
                        in if null tr''
                              then tryCompress$ mkTrie v pre (Map.delete y m)
                              else mkTrie v pre (Map.insert y tr'' m)

            _ -> error "Data.ListTrie.Patricia.Base.deleteSuffixes \
                       \:: internal error"

-- O(1)
splitPrefix :: (Alt st a, Boolable (st a), Trie trie st map v k)
            => trie map v k a -> (v k, st a, trie map v k a)
splitPrefix tr =
   let (v,pre,m) = tParts tr
    in (pre, v, tryCompress $ mkTrie altEmpty mempty m)

-- O(1)
children :: Trie trie st map v k => trie map v k a -> CMap trie map v k a
children = tMap

-- O(1)
children1 :: Trie trie st map v k => trie map v k a -> CMap trie map v k a
children1 tr =
   let (v,pre,m) = tParts tr
    in case unconsVector pre of
            (Nothing, _)   -> m
            (Just p, ps) -> Map.singleton p (mkTrie v ps m)

-- * Visualization

-- O(n m)
showTrieWith :: (Show k, Show (v k), Trie trie st map v k)
             => (st a -> ShowS) -> trie map v k a -> ShowS
showTrieWith = go 0
 where
   go indent f tr =
      let (v,pre,m) = tParts tr
          spre      = shows pre
          lpre      = length (spre mempty)
          sv        = f v
          lv        = length (sv mempty)
       in spre . showChar ' '
        . sv . showChar ' '
        . (foldr (.) id . zipWith (flip ($)) (False : repeat True) $
              map (\(k,t) -> \b -> let sk = shows k
                                       lk = length (sk mempty)
                                       i  = indent + lpre + lv + 2
                                    in (if b
                                           then showChar '\n'
                                              . showString (replicate i ' ')
                                           else id)
                                     . showString "-> "
                                     . sk . showChar ' '
                                     . go (i + lk + 4) f t)
                  (Map.toListKV m))

-- helpers

-- mkTrie, but makes sure that empty tries don't have nonempty prefixes
-- intentionally strict in the value: gives update its semantics
safeMkTrie :: (Alt st a, Boolable (st a), Trie trie st map v k)
           => st a -> v k -> CMap trie map v k a -> trie map v k a
safeMkTrie v p m =
   if noValue v && Map.null m
      then empty
      else mkTrie v p m

prepend :: (VG.Vector v a, Monoid (v a)) => v a -> a -> v a -> v a
prepend prefix key = (prefix <>) . (VG.cons key)

data PrefixOrdering v a
   = Same
   | PostFix (Either (v a) (v a))
   | DifferedAt (v a) (v a) (v a)

-- Same                  If they're equal.
-- PostFix (Left  xs)    If the first argument was longer: xs is the remainder.
-- PostFix (Right xs)    Likewise, but for the second argument.
-- DifferedAt pre xs ys  Otherwise. pre is the part that was the same and
--                       xs and ys are the remainders for the first and second
--                       arguments respectively.
--
--                       all (pre `isPrefixOf`) [xs,ys] --> True.
comparePrefixes :: (VG.Vector v a, Monoid (v a))
                => (a -> a -> Bool) -> v a -> v a -> PrefixOrdering v a
comparePrefixes = go mempty
 where
   go _ _ lxs rxs | VG.null lxs && VG.null rxs = Same
   go _ _ lxs rxs | VG.null lxs = PostFix (Right rxs)
   go _ _ lxs rxs | VG.null rxs = PostFix (Left  lxs)

   go samePart (===) lxs rxs =
      if (VG.head lxs) === (VG.head rxs)
         then go (VG.cons (VG.head lxs) samePart) (===) (VG.tail lxs) (VG.tail rxs)
         else DifferedAt (VG.reverse samePart) lxs rxs

-- Exported for Eq/Ord instances
eqComparePrefixes :: (VG.Vector v a, Monoid (v a)) => (a -> a -> Bool) -> v a -> v a -> Bool
eqComparePrefixes eq xs ys = case comparePrefixes eq xs ys of
                                  Same -> True
                                  _    -> False

ordComparePrefixes :: (VG.Vector v a, Monoid (v a))
                   => (a -> a -> Ordering) -> v a -> v a -> Ordering
ordComparePrefixes ord xs ys =
   case comparePrefixes (\x y -> ord x y == EQ) xs ys of
        Same                     -> EQ
        PostFix r                -> either (const GT) (const LT) r
        DifferedAt _ lxs rxs | not (VG.null lxs || VG.null rxs) ->
                                 ord (VG.head lxs) (VG.head rxs)
        _                        -> error
           "Data.ListTrie.Patricia.Base.ordComparePrefixes :: internal error"

-- After modifying the trie, compress a trie node into the prefix if possible.
--
-- Doesn't recurse into children, only checks if this node and its child can be
-- joined into one. Does it repeatedly, though, until it can't compress any
-- more.
--
-- Note that this is a sledgehammer: for optimization, instead of using this in
-- every function, we could write a separate tryCompress for each function,
-- checking only for those cases that we know can arise. This has been done in
-- 'insert', at least, but not in many places.
tryCompress :: (Boolable (st a), Trie trie st map v k)
            => trie map v k a -> trie map v k a
tryCompress tr =
   let (v,pre,m) = tParts tr
    in case Map.singletonView m of

          -- We can compress the trie if there is only one child
          Just (x, tr')
             -- If the parent is empty, we can collapse it into the child
             | noValue v -> tryCompress $ mkTrie v' (prepend pre x pre') subM

             -- If the parent is full and the child is empty and childless, the
             -- child is irrelevant
             | noValue v' && Map.null subM -> mkTrie v pre subM
           where
             (v',pre',subM) = tParts tr'

          -- If the trie is empty, make sure the prefix is as well.
          --
          -- This case can arise in 'intersectionWith', at least.
          Nothing | noValue v && Map.null m -> mkTrie v mempty m

          -- Otherwise, leave it unchanged.
          _ -> tr

unconsVector :: VG.Vector v a => v a -> (Maybe a, v a)
unconsVector = first VG.headM . VG.splitAt 1

-----------------------
-- Support functions --
-----------------------

hasValue, noValue :: Boolable b => b -> Bool
hasValue = toBool
noValue  = not . hasValue

tVal :: Trie trie st map v k => trie map v k a -> st a
tVal = (\(a,_,_) -> a) . tParts

tMap :: Trie trie st map v k => trie map v k a -> CMap trie map v k a
tMap = (\(_,_,c) -> c) . tParts

genericInsertWith :: (Alt st a, Boolable (st a), Trie trie st map v k)
                  => (forall x y. (x -> y) -> x -> y)
                  -> ((a -> a) -> st a -> st a)
                  -> (a -> a -> a) -> v k -> a -> trie map v k a -> trie map v k a
genericInsertWith ($$) (<$$>) f = go
  where
    mkTrie' = ($$) mkTrie
    go k new tr =
      let (old,prefix,m) = tParts tr
      in case comparePrefixes (Map.eqCmp m) prefix k of
           Same -> mkTrie' ((f new <$$> old) <|> pure new) prefix m

           PostFix (Left (unconsVector -> (Just p, pr))) ->
             mkTrie' (pure new) k
             (Map.singleton p (mkTrie old pr m))
           PostFix (Right (unconsVector -> (Just x, xs))) ->
             -- Minor optimization: instead of tryCompress we just check
             -- for the case of an empty trie
             if null tr
             then singleton k new
             else mkTrie old prefix $
                  Map.insertWith (\_ oldt -> go xs new oldt)
                  x (singleton xs new) m

           DifferedAt pr' (unconsVector -> (Just p, pr)) (unconsVector -> (Just x, xs)) ->
             mkTrie altEmpty pr' $ Map.doubleton x (singleton xs new)
             p (mkTrie old pr m)

           _ -> error "Data.ListTrie.Patricia.Base.insertWith :: internal error"

genericAdjust :: Trie trie st map v k
              => (forall x y. (x -> y) -> x -> y)
              -> ((a -> a) -> st a -> st a)
              -> (a -> a) -> v k -> trie map v k a -> trie map v k a
genericAdjust ($$) myFmap f = go
 where
   go k tr =
      let (v,prefix,m) = tParts tr
       in case comparePrefixes (Map.eqCmp m) prefix k of
               Same                   -> (mkTrie $$ myFmap f v) prefix m
               PostFix (Right (unconsVector -> (Just x, xs))) ->
                  mkTrie v prefix $ Map.adjust (go xs) x m
               _                      -> tr

genericAlter :: (Alt st a, Boolable (st a), Trie trie st map v k)
             => (st a -> trie map v k a -> trie map v k a)
             -> (st a -> st a) -> v k -> trie map v k a -> trie map v k a
genericAlter seeq f = go
 where
   go k tr =
      let (v,prefix,m) = tParts tr
       in case comparePrefixes (Map.eqCmp m) prefix k of
               Same                   ->
                  let v' = f v
                   in -- We need to compress if the map was empty or a
                      -- singleton and the value was removed
                      if    (Map.null m || isJust (Map.singletonView m))
                         && not (hasValue v')
                         then tryCompress (mkTrie v' prefix m)
                         else v' `seeq` mkTrie v' prefix m

               PostFix (Right (unconsVector -> (Just x, xs))) ->
                  mkTrie v prefix $
                     Map.alter
                        (\mt ->
                           case mt of
                                Nothing ->
                                   let v' = f altEmpty
                                    in if hasValue v'
                                          then Just (singleton xs (unwrap v'))
                                          else Nothing
                                Just t ->
                                   let new = go xs t
                                    in if null new then Nothing else Just new)
                        x m

               PostFix (Left (unconsVector -> (Just p, ps))) ->
                  let v' = f altEmpty
                   in if hasValue v'
                         then mkTrie v' k $ Map.singleton p (mkTrie v ps m)
                         else tr

               DifferedAt pr (unconsVector -> (Just p, ps)) (unconsVector -> (Just x, xs)) ->
                  let v' = f altEmpty
                   in if hasValue v'
                         then mkTrie altEmpty pr $
                                 Map.doubleton p (mkTrie v  ps m)
                                               x (mkTrie v' xs Map.empty)
                         else tr

               _ -> error
                  "Data.ListTrie.Patricia.Base.genericAlter :: internal error"

genericUnionWith :: (Alt st a, Boolable (st a), Trie trie st map v k)
                 => (st a -> trie map v k a -> trie map v k a)
                 -> (st a -> st a -> st a)
                 -> trie map v k a
                 -> trie map v k a
                 -> trie map v k a
genericUnionWith seeq = go
 where
   go valUnion tr1 tr2 =
      let (v1,pre1,m1) = tParts tr1
          (v2,pre2,m2) = tParts tr2
       in case comparePrefixes (Map.eqCmp m1) pre1 pre2 of
               Same ->
                  let v = valUnion v1 v2

                      -- safeMkTrie not needed: if pre1 is not null then m1 or
                      -- v won't be and hence the union won't be.
                   in v `seeq` (tryCompress.mkTrie v pre1 $
                                               mapUnion valUnion m1 m2)

               PostFix remainder ->
                  -- As above, mkTrie is fine
                  --
                  -- The flip is important to retain left-biasedness
                  tryCompress $
                     either
                        (mkTrie v2 pre2 . mapUnion (flip valUnion) m2 .
                           decompress m1 v1)
                        (mkTrie v1 pre1 . mapUnion       valUnion  m1 .
                           decompress m2 v2)
                        remainder

               DifferedAt pr (unconsVector -> (Just x, xs)) (unconsVector -> (Just y, ys)) ->
                  -- As above, mkTrie is fine
                  mkTrie altEmpty pr $ Map.doubleton x (mkTrie v1 xs m1)
                                                     y (mkTrie v2 ys m2)

               _ -> can'tHappen

   mapUnion = Map.unionWith . go

   decompress m v (unconsVector -> (Just x, xs)) = Map.singleton x (mkTrie v xs m)
   decompress _ _ _                              = can'tHappen

   can'tHappen =
      error "Data.ListTrie.Patricia.Base.unionWith :: internal error"

genericUnionWithKey :: (Alt st a, Boolable (st a), Trie trie st map v k)
                    => (st a -> trie map v k a -> trie map v k a)
                    -> ((a -> a -> a) -> st a -> st a -> st a)
                    -> (v k -> a -> a -> a)
                    -> trie map v k a
                    -> trie map v k a
                    -> trie map v k a
genericUnionWithKey seeq = go mempty
 where
   go k valUnion j tr1 tr2 =
      let (v1,pre1,m1) = tParts tr1
          (v2,pre2,m2) = tParts tr2
       in case comparePrefixes (Map.eqCmp m1) pre1 pre2 of
               Same ->
                  let k' = k <> pre1
                      v  = valUnion (j $ k') v1 v2
                   in v `seeq`
                         (tryCompress.mkTrie v pre1 $
                            mapUnion valUnion j k pre1 m1 m2)

               PostFix remainder ->
                  tryCompress $
                     either
                        (mk v2 pre2 . mapUnion (flip.valUnion) j k pre2 m2
                           . decompress m1 v1)
                        (mk v1 pre1 . mapUnion       valUnion  j k pre1 m1
                           . decompress m2 v2)
                        remainder

               DifferedAt pr (unconsVector -> (Just x, xs)) (unconsVector -> (Just y, ys)) ->
                  mkTrie altEmpty pr $ Map.doubleton x (mkTrie v1 xs m1)
                                                     y (mkTrie v2 ys m2)

               _ -> can'tHappen

   mk = mkTrie

   mapUnion v j k p =
      Map.unionWithKey $
         \x -> go (k <> p `VG.snoc` x) v j

   decompress m v (unconsVector -> (Just x, xs)) = Map.singleton x (mkTrie v xs m)
   decompress _ _ _                              = can'tHappen

   can'tHappen =
      error "Data.ListTrie.Patricia.Base.unionWithKey :: internal error"

genericIntersectionWith :: forall a b c k v map st trie.
                           ( Alt st c, Boolable (st c)
                           , Trie trie st map v k
                           )
                        => (forall x. st x -> trie map v k x -> trie map v k x)
                        -> (st a -> st b -> st c)
                        -> trie map v k a
                        -> trie map v k b
                        -> trie map v k c
genericIntersectionWith seeq = go0
 where
   go0 valIsect trl trr =
      let (vl,prel,ml) = tParts trl
          (vr,prer,mr) = tParts trr
       in case comparePrefixes (Map.eqCmp ml) prel prer of
               DifferedAt _ _ _  -> empty
               Same              -> mk valIsect vl vr prel ml mr
               PostFix remainder ->
                  -- use the one with a longer prefix as the base for the
                  -- intersection, and descend into the map of the one with a
                  -- shorter prefix
                  either (go       valIsect  mr vl ml (DL.fromList $ VG.toList prel))
                         (go (flip valIsect) ml vr mr (DL.fromList $ VG.toList prer))
                         remainder

   mapIntersect valIsect =
      Map.filter (not.null) .:
         Map.intersectionWith (go0 valIsect)

   mk valIsect v v' p m m' =
      let vi = valIsect v v'
       in vi `seeq` (tryCompress.mkTrie vi p $ mapIntersect valIsect m m')

   -- Polymorphic recursion in 'go' (valIsect :: st a -> st b -> st c ---> st b
   -- -> st a -> st c) means that it has to be explicitly typed in order to
   -- compile.
   --
   -- The repeated "Trie trie st map v k" constraint is for Hugs.

   -- Like goLeft and goRight in 'difference', but handles both cases (since
   -- this is a commutative operation).
   --
   -- Traverse the map given as the 1st argument, looking for anything that
   -- begins with the given key (x:xs).
   --
   -- If it's found, great: make an intersected trie out of the trie found in
   -- the map and the boolean, map, and prefix given.
   --
   -- If it's not found but might still be, there are two cases.
   --
   -- 1. Say we've got the following two TrieSets:
   --
   -- fromList ["car","cat"]
   -- fromList ["car","cot"]
   --
   -- i.e. (where <> is stuff we don't care about here)
   --
   -- Tr False "ca" (fromList [('r', Tr True ""  <>),<>])
   -- Tr False "c"  (fromList [('a', Tr True "r" <>),<>])
   --
   -- We came in here with (x:xs) = "a", the remainder of comparing "ca" and
   -- "c". We're looking for anything that begins with "ca" from the children
   -- of the "c".
   --
   -- We find the prefix pre' = "r", and comparePrefixes gives PostFix (Right
   -- "r"). So now we want anything beginning with "car" in the other trie. We
   -- switch to traversing the other trie, i.e. the other given map: the
   -- children of "ca".
   --
   -- 2. Say we have the following:
   --
   -- fromList ["cat"]
   -- fromList ["cat","cot","cap"]
   --
   -- i.e.
   --
   -- Tr True "cat" <>
   -- Tr False "c" (fromList [('a',Tr False "" (fromList [('t',<>)])),<>])
   --
   -- (x:xs) = "at" now, and we find pre' = "". We get PostFix (Left "t"). This
   -- means that we're staying in the same trie, just looking for "t" now
   -- instead of "at". So we jump into the m' map.
   --
   -- Note that the prefix and boolean don't change: we've already got "ca",
   -- and we'd still like "cat" so we keep the True from there.
   go :: (Alt st z, Boolable (st z), Trie trie st map v k)
      => (st x -> st y -> st z)
      -> CMap trie map v k y
      -> st x
      -> CMap trie map v k x
      -> DList k
      -> v k
      -> trie map v k z
   go valIsect ma v mb pre (unconsVector -> (Just x, xs)) =
      case Map.lookup x ma of
           Nothing -> empty
           Just tr ->
              let (v',pre',m') = tParts tr
               in case comparePrefixes (Map.eqCmp ma) xs pre' of
                       DifferedAt _ _ _   -> empty
                       Same               ->
                          mk valIsect v v' (VG.fromList $ DL.toList pre) mb m'
                       PostFix (Right ys) ->
                          let nextPre = pre `DL.append` DL.fromList (VG.toList ys)
                           in go (flip valIsect) mb v' m' nextPre ys
                       PostFix (Left  ys) ->
                              go       valIsect  m' v  mb pre     ys

   go _ _ _ _ _ _ =
      error "Data.ListTrie.Patricia.Map.intersectionWith :: internal error"

genericIntersectionWithKey :: forall a b c k v map st trie.
                              (Alt st c, Boolable (st c), Trie trie st map v k)
                           => (forall x. st x -> trie map v k x -> trie map v k x)
                           -> ((a -> b -> c) -> st a -> st b -> st c)
                           -> (v k -> a -> b -> c)
                           -> trie map v k a
                           -> trie map v k b
                           -> trie map v k c
genericIntersectionWithKey seeq = main DL.empty
 where
   main k valIsect j trl trr =
      let (vl,prel,ml) = tParts trl
          (vr,prer,mr) = tParts trr
       in case comparePrefixes (Map.eqCmp ml) prel prer of
               DifferedAt _ _ _ -> empty
               Same             -> mk k valIsect j vl vr prel ml mr
               PostFix remainder ->
                  let prel' = DL.fromList $ VG.toList prel
                      prer' = DL.fromList $ VG.toList prer
                   in either
                         (go k       valIsect        j  mr vl ml prel')
                         (go k (flop valIsect) (flip.j) ml vr mr prer')
                         remainder

   mk k valIsect j v v' p m m' =
      let k' = k `DL.append` DL.fromList (VG.toList p)
          vi = valIsect (j $ (VG.fromList $ DL.toList k')) v v'
       in vi `seeq` (tryCompress.mkTrie vi p $
                                    mapIntersect k' valIsect j m m')

   mapIntersect k valIsect j =
      Map.filter (not.null) .:
         Map.intersectionWithKey (\x -> main (k `DL.snoc` x) valIsect j)

   flop :: ((x -> y -> z) -> st x -> st y -> st z)
         -> ((y -> x -> z) -> st y -> st x -> st z)
   flop f = flip . f . flip

   -- See intersectionWith: this explicit type is necessary
   go :: (Alt st z, Boolable (st z), Trie trie st map v k)
      => DList k
      -> ((x -> y -> z) -> st x -> st y -> st z)
      -> (v k -> x -> y -> z)
      -> CMap trie map v k y
      -> st x
      -> CMap trie map v k x
      -> DList k
      -> v k
      -> trie map v k z
   go k valIsect j ma v mb pre (unconsVector -> (Just x, xs)) =
      case Map.lookup x ma of
           Nothing -> empty
           Just tr ->
              let (v',pre',m') = tParts tr
               in case comparePrefixes (Map.eqCmp ma) xs pre' of
                       DifferedAt _ _ _   -> empty
                       Same               ->
                          mk k valIsect j v v' (VG.fromList $ DL.toList pre) mb m'
                       PostFix (Right ys) ->
                          let nextPre = pre `DL.append` DL.fromList (VG.toList ys)
                           in go k (flop valIsect) (flip.j) mb v' m' nextPre ys
                       PostFix (Left  ys) ->
                              go k        valIsect       j  m' v  mb pre     ys

   go _ _ _ _ _ _ _ _ =
      error "Data.ListTrie.Patricia.Map.intersectionWithKey :: internal error"