-- |
--   Module      :  Data.Edison.Assoc.TernaryTrie
--   Copyright   :  Copyright (c) 2002, 2008 Andrew Bromage
--   License     :  MIT; see COPYRIGHT file for terms and conditions
--
--   Maintainer  :  robdockins AT fastmail DOT fm
--   Stability   :  stable
--   Portability :  GHC, Hugs (MPTC and FD)
--
--   Finite maps indexed by lists or strings @[k]@, implemented as ternary
--   search tries

module Data.Edison.Assoc.TernaryTrie (
    -- * Type of ternary search tries
    FM,

    -- * AssocX operations
    empty,singleton,fromSeq,insert,insertSeq,union,unionSeq,delete,deleteAll,
    deleteSeq,null,size,member,count,lookup,lookupM,lookupAll,
    lookupAndDelete,lookupAndDeleteM,lookupAndDeleteAll,
    lookupWithDefault,adjust,adjustAll,adjustOrInsert,adjustAllOrInsert,
    adjustOrDelete,adjustOrDeleteAll,strict,strictWith,
    map,fold,fold',fold1,fold1',filter,partition,elements,structuralInvariant,

    -- * Assoc operations
    toSeq,keys,mapWithKey,foldWithKey,foldWithKey',filterWithKey,partitionWithKey,

    -- * FiniteMapX operations
    fromSeqWith,fromSeqWithKey,insertWith,insertWithKey,insertSeqWith,
    insertSeqWithKey,unionl,unionr,unionWith,unionSeqWith,intersectionWith,
    difference,properSubset,subset,properSubmapBy,submapBy,sameMapBy,
    properSubmap,submap,sameMap,

    -- * FiniteMap operations
    unionWithKey,unionSeqWithKey,intersectionWithKey,

    -- * OrdAssocX operations
    minView, minElem, deleteMin, unsafeInsertMin,
    maxView, maxElem, deleteMax, unsafeInsertMax,
    foldr, foldr', foldr1, foldr1', foldl, foldl', foldl1, foldl1',
    unsafeFromOrdSeq, unsafeAppend, filterLT, filterLE, filterGT, filterGE,
    partitionLT_GE, partitionLE_GT, partitionLT_GT,

    -- * OrdAssoc operations
    minViewWithKey, minElemWithKey, maxViewWithKey, maxElemWithKey,
    foldrWithKey, foldrWithKey', foldlWithKey, foldlWithKey',
    toOrdSeq,

    -- * Other supported operations
    mergeVFM, mergeKVFM,

    -- * Documentation
    moduleName
) where

import Prelude hiding (null,map,lookup,foldr,foldl,foldr1,foldl1,foldl',filter)
import qualified Prelude
import qualified Data.Edison.Assoc as A
import Data.Edison.Prelude ( runFail_ )
import qualified Data.Edison.Seq as S
import qualified Data.List as L
import qualified Control.Monad.Fail as Fail
import Control.Monad
import Data.Coerce (coerce)
import Data.Monoid
import Data.Semigroup as SG
import Data.Maybe (isJust, isNothing)

import Data.Edison.Assoc.Defaults
import Test.QuickCheck (Arbitrary(..), CoArbitrary(..), Gen(), NonNegative(..), variant, sized, resize, choose, oneof)


-- signatures for exported functions
moduleName    :: String
empty         :: Ord k => FM k a
singleton     :: Ord k => [k] -> a -> FM k a
fromSeq       :: (Ord k,S.Sequence seq) => seq ([k],a) -> FM k a
insert        :: Ord k => [k] -> a -> FM k a -> FM k a
insertSeq     :: (Ord k,S.Sequence seq) => seq ([k],a) -> FM k a -> FM k a
union         :: Ord k => FM k a -> FM k a -> FM k a
unionSeq      :: (Ord k,S.Sequence seq) => seq (FM k a) -> FM k a
delete        :: Ord k => [k] -> FM k a -> FM k a
deleteAll     :: Ord k => [k] -> FM k a -> FM k a
deleteSeq     :: (Ord k,S.Sequence seq) => seq [k] -> FM k a -> FM k a
null          :: Ord k => FM k a -> Bool
size          :: Ord k => FM k a -> Int
member        :: Ord k => [k] -> FM k a -> Bool
count         :: Ord k => [k] -> FM k a -> Int
lookup        :: Ord k => [k] -> FM k a -> a
lookupM       :: (Ord k, Fail.MonadFail rm) => [k] -> FM k a -> rm a
lookupAll     :: (Ord k,S.Sequence seq) => [k] -> FM k a -> seq a
lookupAndDelete    :: Ord k => [k] -> FM k a -> (a, FM k a)
lookupAndDeleteM   :: (Ord k, Fail.MonadFail rm) => [k] -> FM k a -> rm (a, FM k a)
lookupAndDeleteAll :: (Ord k, S.Sequence seq) => [k] -> FM k a -> (seq a,FM k a)
lookupWithDefault  :: Ord k => a -> [k] -> FM k a -> a
adjust        :: Ord k => (a -> a) -> [k] -> FM k a -> FM k a
adjustAll     :: Ord k => (a -> a) -> [k] -> FM k a -> FM k a
adjustOrInsert    :: Ord k => (a -> a) -> a -> [k] -> FM k a -> FM k a
adjustAllOrInsert :: Ord k => (a -> a) -> a -> [k] -> FM k a -> FM k a
adjustOrDelete    :: Ord k => (a -> Maybe a) -> [k] -> FM k a -> FM k a
adjustOrDeleteAll :: Ord k => (a -> Maybe a) -> [k] -> FM k a -> FM k a
strict            :: FM k a -> FM k a
strictWith        :: (a -> b) -> FM k a -> FM k a
map           :: Ord k => (a -> b) -> FM k a -> FM k b
fold          :: Ord k => (a -> b -> b) -> b -> FM k a -> b
fold1         :: Ord k => (a -> a -> a) -> FM k a -> a
fold'         :: Ord k => (a -> b -> b) -> b -> FM k a -> b
fold1'        :: Ord k => (a -> a -> a) -> FM k a -> a
filter        :: Ord k => (a -> Bool) -> FM k a -> FM k a
partition     :: Ord k => (a -> Bool) -> FM k a -> (FM k a, FM k a)
elements      :: (Ord k,S.Sequence seq) => FM k a -> seq a

fromSeqWith      :: (Ord k,S.Sequence seq) =>
                        (a -> a -> a) -> seq ([k],a) -> FM k a
fromSeqWithKey   :: (Ord k,S.Sequence seq) => ([k] -> a -> a -> a) -> seq ([k],a) -> FM k a
insertWith       :: Ord k => (a -> a -> a) -> [k] -> a -> FM k a -> FM k a
insertWithKey    :: Ord k => ([k] -> a -> a -> a) -> [k] -> a -> FM k a -> FM k a
insertSeqWith    :: (Ord k,S.Sequence seq) =>
                        (a -> a -> a) -> seq ([k],a) -> FM k a -> FM k a
insertSeqWithKey :: (Ord k,S.Sequence seq) =>
                        ([k] -> a -> a -> a) -> seq ([k],a) -> FM k a -> FM k a
unionl           :: Ord k => FM k a -> FM k a -> FM k a
unionr           :: Ord k => FM k a -> FM k a -> FM k a
unionWith        :: Ord k => (a -> a -> a) -> FM k a -> FM k a -> FM k a
unionSeqWith     :: (Ord k,S.Sequence seq) =>
                        (a -> a -> a) -> seq (FM k a) -> FM k a
intersectionWith :: Ord k => (a -> b -> c) -> FM k a -> FM k b -> FM k c
difference       :: Ord k => FM k a -> FM k b -> FM k a
properSubset     :: Ord k => FM k a -> FM k b -> Bool
subset           :: Ord k => FM k a -> FM k b -> Bool
properSubmapBy   :: Ord k => (a -> a -> Bool) -> FM k a -> FM k a -> Bool
submapBy         :: Ord k => (a -> a -> Bool) -> FM k a -> FM k a -> Bool
sameMapBy        :: Ord k => (a -> a -> Bool) -> FM k a -> FM k a -> Bool
properSubmap     :: (Ord k, Eq a) => FM k a -> FM k a -> Bool
submap           :: (Ord k, Eq a) => FM k a -> FM k a -> Bool
sameMap          :: (Ord k, Eq a) => FM k a -> FM k a -> Bool

toSeq            :: (Ord k,S.Sequence seq) => FM k a -> seq ([k],a)
keys             :: (Ord k,S.Sequence seq) => FM k a -> seq [k]
mapWithKey       :: Ord k => ([k] -> a -> b) -> FM k a -> FM k b
foldWithKey      :: Ord k => ([k] -> a -> b -> b) -> b -> FM k a -> b
foldWithKey'     :: Ord k => ([k] -> a -> b -> b) -> b -> FM k a -> b
filterWithKey    :: Ord k => ([k] -> a -> Bool) -> FM k a -> FM k a
partitionWithKey :: Ord k => ([k] -> a -> Bool) -> FM k a -> (FM k a, FM k a)
unionWithKey     :: Ord k => ([k] -> a -> a -> a) -> FM k a -> FM k a -> FM k a
unionSeqWithKey  :: (Ord k,S.Sequence seq) =>
                       ([k] -> a -> a -> a) -> seq (FM k a) -> FM k a
intersectionWithKey :: Ord k => ([k] -> a -> b -> c) -> FM k a -> FM k b -> FM k c

foldr          :: Ord k => (a -> b -> b) -> b -> FM k a -> b
foldr1         :: Ord k => (a -> a -> a) -> FM k a -> a
foldr'         :: Ord k => (a -> b -> b) -> b -> FM k a -> b
foldr1'        :: Ord k => (a -> a -> a) -> FM k a -> a

foldrWithKey   :: Ord k => ([k] -> a -> b -> b) -> b -> FM k a -> b
foldrWithKey'  :: Ord k => ([k] -> a -> b -> b) -> b -> FM k a -> b
foldlWithKey   :: Ord k => (b -> [k] -> a -> b) -> b -> FM k a -> b
foldlWithKey'  :: Ord k => (b -> [k] -> a -> b) -> b -> FM k a -> b
toOrdSeq       :: (Ord k,S.Sequence seq) => FM k a -> seq ([k],a)

moduleName = "Data.Edison.Assoc.TernaryTrie"


data FM k a
  = FM !(Maybe a) !(FMB k a)

-- | This is isomorphic to an iteration of binary trees with keys @k@.
--
-- @
-- data BT k v = E | I k v (BT k v) (BT k v)
-- data Layer k v x = Layer (Maybe v) (BT k x)
--
-- FMB k v = Fix (Layer k v)
-- @
--
-- The trees are weight-balanced trees, ensuring that the sizes of the
-- two subtrees of any node are bounded by each other up to a constant factor.
--
-- @
-- size l + size r <= 1
--
-- -- or --
--
-- size l <= 6 * size r
-- size r <= 6 * size l
-- @
--
-- Source: <https://yoichihirai.com/bst.pdf Balancing Weight-Balanced Trees>
-- by Hirai and Yamamoto, 2011 (Section 4)
data FMB k v
  = E
  | I !Int !k !(Maybe v) !(FMB k v) !(FMB' k v) !(FMB k v)
  deriving Show

newtype FMB' k v
  = FMB' (FMB k v)
  deriving Show

-- | The balance factor must be either 3 or 4.
-- With other factors, the invariant gets broken by delete, minViewWithKey and maxViewWithKey.
-- (cf. Section 4 of the paper linked above)
balance :: Int
balance = 4

sizeFMB :: FMB k v -> Int
sizeFMB E = 0
sizeFMB (I size _ _ _ _ _) = size

mkFMB :: k -> Maybe v -> FMB k v -> FMB' k v -> FMB k v -> FMB k v
mkFMB k v l m r
  = I (1 + sizeFMB l + sizeFMB r) k v l m r

lookupFMB :: (Ord k) => [k] -> FMB k v -> Maybe v
lookupFMB []        _
  = Nothing
lookupFMB (_:_) E
  = Nothing
lookupFMB nk@(x:xs) (I _ k v l (FMB' fmbm) r)
  = case compare x k of
        LT -> lookupFMB nk l
        GT -> lookupFMB nk r
        EQ -> if L.null xs then v else lookupFMB xs fmbm

listToFMB :: [k] -> v -> FMB k v
listToFMB [x]    v = mkFMB x (Just v) E (FMB' E)                E
listToFMB (x:xs) v = mkFMB x Nothing  E (FMB' $ listToFMB xs v) E
listToFMB _ _ = error "TernaryTrie.listToFMB: bug!"

addToFMB :: (Ord k) => [k] -> (Maybe v -> Maybe v) -> FMB k v -> FMB k v
addToFMB xs combiner E
  = case combiner Nothing of
      Just v -> listToFMB xs v
      Nothing -> E
addToFMB nk@(x:xs) combiner (I size k v l m@(FMB' fmbm) r)
  = case compare x k of
        LT -> mkBalancedFMB k v (addToFMB nk combiner l) m r
        GT -> mkBalancedFMB k v l m (addToFMB nk combiner r)
        EQ -> case xs of
                [] -> case combiner v of
                  Nothing | FMB' E <- m -> appendFMB l r
                  v' -> I size k v' l m r
                _  -> case addToFMB xs combiner fmbm of
                  E | Nothing <- v -> appendFMB l r
                  m' -> I size k v l (FMB' m') r
addToFMB _ _ _ = error "TernaryTrie.addToFMB: bug!"

addToFM :: (Ord k) => [k] -> (Maybe v -> Maybe v) -> FM k v -> FM k v
addToFM [] combiner (FM n fmb)
  = FM (combiner n) fmb
addToFM xs combiner (FM n fmb)
  = FM n (addToFMB xs combiner fmb)

lookupAndDelFromFMB :: (Ord k) => z -> (v -> FMB k v -> z) -> [k] -> FMB k v -> z
lookupAndDelFromFMB onFail _ _ E = onFail
lookupAndDelFromFMB onFail cont nk@(x:xs) (I size k v l m@(FMB' fmbm) r)
  = case compare x k of
        LT -> lookupAndDelFromFMB onFail (\w l' -> cont w (mkBalancedFMB k v l' m r)) nk l
        GT -> lookupAndDelFromFMB onFail (\w r' -> cont w (mkBalancedFMB k v l m r')) nk r
        EQ -> case xs of
                [] -> case v of
                        Nothing -> onFail
                        Just w  -> case fmbm of
                                      E -> cont w (appendFMB l r)
                                      _ -> cont w (I size k Nothing l m r)
                _  -> lookupAndDelFromFMB onFail (\w m' -> case m' of
                  E | Nothing <- v -> cont w (appendFMB l r)
                  _ -> cont w (I size k v l (FMB' m') r)) xs fmbm
lookupAndDelFromFMB _ _ _ _ = error "TernaryTrie.lookupAndDelFromFMB: bug!"

lookupAndDelFromFM :: (Ord k) => z -> (v -> FM k v -> z) -> [k] -> FM k v -> z
lookupAndDelFromFM onFail _ [] (FM Nothing _)  = onFail
lookupAndDelFromFM _ cont [] (FM (Just v) fmb) = cont v (FM Nothing fmb)
lookupAndDelFromFM onFail cont xs (FM n fmb) =
   lookupAndDelFromFMB onFail (\w fmb' -> cont w (FM n fmb')) xs fmb


delFromFMB :: (Ord k) => [k] -> FMB k v -> FMB k v
delFromFMB _ E
  = E
delFromFMB nk@(x:xs) (I size k v l m@(FMB' fmbm) r)
  = case compare x k of
        LT -> mkBalancedFMB k v (delFromFMB nk l) m r
        GT -> mkBalancedFMB k v l m (delFromFMB nk r)
        EQ -> case xs of
                [] -> case fmbm of
                        E -> appendFMB l r
                        _ -> I size k Nothing l m r
                _  -> case delFromFMB xs fmbm of
                  E | Nothing <- v -> appendFMB l r
                  m' -> I size k v l (FMB' m') r
delFromFMB _ _ = error "TernaryTrie.delFromFMB: bug!"


delFromFM :: (Ord k) => [k] -> FM k v -> FM k v
delFromFM [] (FM _ fmb)
  = FM Nothing fmb
delFromFM xs (FM n fmb)
  = FM n (delFromFMB xs fmb)


mkBalancedFMB :: k -> Maybe v -> FMB k v -> FMB' k v -> FMB k v -> FMB k v
mkBalancedFMB k v l m r
  | size_l + size_r < 2
    = mkFMB k v l m r
  | size_r > balance * size_l        -- Right tree too big
    = case r of
        I _ _ _ rl _ rr
            | sizeFMB rl < 2 * sizeFMB rr
                -> single_L l m r
            | otherwise
                -> double_L l m r
        _ -> error "TernaryTrie.mkBalancedFMB: bug!"

  | size_l > balance * size_r   -- Left tree too big
    = case l of
        I _ _ _ ll _ lr
            | sizeFMB lr < 2 * sizeFMB ll
                -> single_R l m r
            | otherwise
                -> double_R l m r
        _ -> error "TernaryTrie.mkBalancedFMB: bug!"

  | otherwise                           -- No imbalance
    = mkFMB k v l m r
  where
        size_l   = sizeFMB l
        size_r   = sizeFMB r

        single_L l m (I _ k_r v_r rl rm rr)
          = mkFMB k_r v_r (mkFMB k v l m rl) rm rr
        single_L _ _ _ = error "TernaryTrie:mkBalancedFMB: bug!"

        double_L l m (I _ k_r v_r (I _ k_rl v_rl rll rlm rlr) rm rr)
          = mkFMB k_rl v_rl (mkFMB k v l m rll) rlm (mkFMB k_r v_r rlr rm rr)
        double_L _ _ _ = error "TernaryTrie:mkBalancedFMB: bug!"

        single_R (I _ k_l v_l ll lm lr) m r
          = mkFMB k_l v_l ll lm (mkFMB k v lr m r)
        single_R _ _ _ = error "TernaryTrie:mkBalancedFMB: bug!"

        double_R (I _ k_l v_l ll lm (I _ k_lr v_lr lrl lrm lrr)) m r
          = mkFMB k_lr v_lr (mkFMB k_l v_l ll lm lrl) lrm (mkFMB k v lrr m r)
        double_R _ _ _ = error "TernaryTrie:mkBalancedFMB: bug!"


mkVBalancedFMB :: k -> Maybe v -> FMB k v -> FMB' k v -> FMB k v -> FMB k v
mkVBalancedFMB k Nothing l (FMB' E) r
  = appendFMB l r
mkVBalancedFMB k v E m E
  = mkFMB k v E m E
mkVBalancedFMB k v l@E m (I _ kr vr rl rm rr)
  = mkBalancedFMB kr vr (mkVBalancedFMB k v l m rl) rm rr
mkVBalancedFMB k v (I _ kl vl ll lm lr) m r@E
  = mkBalancedFMB kl vl ll lm (mkVBalancedFMB k v lr m r)
mkVBalancedFMB k v l@(I _ kl vl ll lm lr) m r@(I _ kr vr rl rm rr)
  | balance * size_l < size_r
    = mkBalancedFMB kr vr (mkVBalancedFMB k v l m rl) rm rr
  | balance * size_r < size_l
    = mkBalancedFMB kl vl ll lm (mkVBalancedFMB k v lr m r)
  | otherwise
    = mkFMB k v l m r
  where
        size_l = sizeFMB l
        size_r = sizeFMB r

    -- Constraint: All keys in the first FMB are less than
    -- that in the second FMB.
appendFMB :: FMB k v -> FMB k v -> FMB k v
appendFMB E m2 = m2
appendFMB m1 E = m1
appendFMB fmb1@(I size1 k1 v1 l1 m1 r1) fmb2@(I size2 k2 v2 l2 m2 r2)
  | size1 > size2
    = mkVBalancedFMB k1 v1 l1 m1 (appendFMB r1 fmb2)
  | otherwise
    = mkVBalancedFMB k2 v2 (appendFMB fmb1 l2) m2 r2

mapVFM :: (Maybe a -> Maybe b) -> FM k a -> FM k b
mapVFM f (FM n fmb)
  = FM (f n) (mapVFMB f fmb)

mapVFMB :: (Maybe a -> Maybe b) -> FMB k a -> FMB k b
mapVFMB f m
  = mapVFMB' m
  where
        mapVFMB' E = E
        mapVFMB' (I _ k v l (FMB' m) r)
          = case (mapVFMB' m, f v) of
                (E,Nothing) -> appendFMB (mapVFMB' l) (mapVFMB' r)
                (m',v')     -> mkVBalancedFMB k v'
                                    (mapVFMB' l) (FMB' m') (mapVFMB' r)

mapKVFM :: ([k] -> Maybe a -> Maybe b) -> FM k a -> FM k b
mapKVFM f (FM n fmb)
  = FM (f [] n) (mapKVFMB [] fmb)
  where
        mapKVFMB _ E = E
        mapKVFMB ks (I _ k v l (FMB' m) r)
          = mkVBalancedFMB k (f (reverse (k:ks)) v)
              (mapKVFMB ks l)
              (FMB' (mapKVFMB (k:ks) m))
              (mapKVFMB ks r)

nullFMB :: FMB k v -> Bool
nullFMB E = True
nullFMB (I _ _ v l (FMB' m) r)
  = case v of
      Just _  -> False
      Nothing -> nullFMB l && nullFMB m && nullFMB r

nullFM :: FM k v -> Bool
nullFM (FM (Just _) _)  = False
nullFM (FM Nothing fmb) = nullFMB fmb

data FMBCtx k v
  = T
  | L !k !(Maybe v) !(FMBCtx k v) !(FMB' k v) !(FMB k v)
  | R !k !(Maybe v) !(FMB k v) !(FMB' k v) !(FMBCtx k v)

splayFMB :: (Ord k) => k -> FMB k a -> (Maybe a, FMB k a, FMB' k a, FMB k a)
splayFMB key fmb
  = splaydown T fmb
  where
    splaydown ctx E
      = splayup ctx Nothing E (FMB' E) E
    splaydown ctx (I _ k v l m r)
      = case compare key k of
            LT -> splaydown (L k v ctx m r) l
            GT -> splaydown (R k v l m ctx) r
            EQ -> splayup ctx v l m r

    splayup ctx v l m r
      = splayup' ctx l r
      where
          splayup' T l r
            = (v, l, m, r)
          splayup' (L ck cv ctx cm cr) tl tr
            = splayup' ctx tl (mkVBalancedFMB ck cv tr cm cr)
          splayup' (R ck cv cl cm ctx) tl tr
            = splayup' ctx (mkVBalancedFMB ck cv cl cm tl) tr

mergeVFMB :: (Ord k) => (Maybe a -> Maybe b -> Maybe c) ->
                FMB k a -> FMB k b -> FMB k c
mergeVFMB f fmbx fmby
  = mergeVFMB' fmbx fmby
  where
    mergeVFMB' E E
      = E
    mergeVFMB' E fmby@(I _ _ _ _ (FMB' _) _)
      = mapVFMB (\v -> f Nothing v) fmby
    mergeVFMB' fmbx@(I _ _ _ _ (FMB' _) _) E
      = mapVFMB (\v -> f v Nothing) fmbx
    mergeVFMB' fmbx@(I sizex kx vx lx (FMB' mx) rx)
               fmby@(I sizey ky vy ly (FMB' my) ry)
      | sizex >= sizey
        = let (vy, ly, FMB' my, ry) = splayFMB kx fmby
          in case (mergeVFMB' mx my, f vx vy) of
                (E,Nothing) -> appendFMB (mergeVFMB' lx ly) (mergeVFMB' rx ry)
                (m',v)      -> mkVBalancedFMB kx v
                                   (mergeVFMB' lx ly)
                                   (FMB' m')
                                   (mergeVFMB' rx ry)
      | otherwise
        = let (vx, lx, FMB' mx, rx) = splayFMB ky fmbx
          in case (mergeVFMB' mx my, f vx vy) of
                (E,Nothing) -> appendFMB (mergeVFMB' lx ly) (mergeVFMB' rx ry)
                (m',v)      -> mkVBalancedFMB ky v
                                   (mergeVFMB' lx ly)
                                   (FMB' m')
                                   (mergeVFMB' rx ry)

mergeVFM :: (Ord k) => (Maybe a -> Maybe b -> Maybe c) ->
                FM k a -> FM k b -> FM k c
mergeVFM f (FM vx fmbx) (FM vy fmby)
  = FM (f vx vy) (mergeVFMB f fmbx fmby)


mergeKVFMB :: (Ord k) => ([k] -> Maybe a -> Maybe b -> Maybe c) ->
                FMB k a -> FMB k b -> FMB k c
mergeKVFMB f fmbx fmby
  = mergeKVFMB' [] fmbx fmby
  where
    mergeKVFMB' _ E E
      = E
    mergeKVFMB' ks E fmby
      = mergeKVFMBs (\k v -> f k Nothing v) ks fmby
    mergeKVFMB' ks fmbx E
      = mergeKVFMBs (\k v -> f k v Nothing) ks fmbx
    mergeKVFMB' ks fmbx@(I sizex kx vx lx (FMB' mx) rx)
                   fmby@(I sizey ky vy ly (FMB' my) ry)
      | sizex >= sizey
        = let (vy, ly, FMB' my, ry) = splayFMB kx fmby
              ks' = reverse (kx:ks)
          in case (mergeKVFMB' ks' mx my, f ks' vx vy) of
                (E,Nothing) -> appendFMB
                                    (mergeKVFMB' ks lx ly)
                                    (mergeKVFMB' ks rx ry)
                (m',v)      -> mkVBalancedFMB kx v
                                    (mergeKVFMB' ks lx ly)
                                    (FMB' m')
                                    (mergeKVFMB' ks rx ry)
      | otherwise
        = let (vx, lx, FMB' mx, rx) = splayFMB ky fmbx
              ks' = reverse (ky:ks)
          in case (mergeKVFMB' ks' mx my, f ks' vx vy) of
                (E,Nothing) -> appendFMB
                                    (mergeKVFMB' ks lx ly)
                                    (mergeKVFMB' ks rx ry)
                (m',v)      -> mkVBalancedFMB ky v
                                    (mergeKVFMB' ks lx ly)
                                    (FMB' m')
                                    (mergeKVFMB' ks rx ry)

    mergeKVFMBs f ks fmb
      = mergeKVFMBs' ks fmb
      where
          mergeKVFMBs' _ E
            = E
          mergeKVFMBs' ks (I _ k v l (FMB' m) r)
            = case (mergeKVFMBs' (k:ks) m, f (reverse (k:ks)) v) of
                (E, Nothing) -> appendFMB
                                    (mergeKVFMBs' ks l)
                                    (mergeKVFMBs' ks r)
                (m,v)        -> mkVBalancedFMB k v
                                    (mergeKVFMBs' ks l)
                                    (FMB' m)
                                    (mergeKVFMBs' ks r)

mergeKVFM :: (Ord k) => ([k] -> Maybe a -> Maybe b -> Maybe c) ->
                FM k a -> FM k b -> FM k c
mergeKVFM f (FM vx fmbx) (FM vy fmby)
  = FM (f [] vx vy) (mergeKVFMB f fmbx fmby)


-- The public interface.
--

-- AssocX

empty = FM Nothing E

singleton [] v = FM (Just v) E
singleton xs v = FM Nothing (listToFMB xs v)

fromSeq = fromSeqUsingInsertSeq

insert k v fm = addToFM k (\_ -> Just v) fm

insertSeq = insertSeqUsingFoldr

union = mergeVFM mplus

unionSeq = unionSeqUsingReduce

delete k fm = delFromFM k fm

deleteAll = delete

deleteSeq = deleteSeqUsingFoldr

null = nullFM

size (FM k fmb)
    | isNothing k = fmb_size fmb 0
    | otherwise   = fmb_size fmb 1
    where fmb_size E k = k
          fmb_size (I _ _ Nothing l (FMB' m) r) k = fmb_size l $ fmb_size m $ fmb_size r k
          fmb_size (I _ _ _ l (FMB' m) r ) k      = fmb_size l $ fmb_size m $ fmb_size r $! k+1


member = memberUsingLookupM

count = countUsingMember

lookup m k = runFail_ (lookupM m k)

lookupM [] (FM Nothing _)
  = fail "TernaryTrie.lookup: lookup failed"
lookupM [] (FM (Just v) _)
  = return v
lookupM xs (FM _ fmb)
  = case  lookupFMB xs fmb  of
        Nothing -> fail "TernaryTrie.lookup: lookup failed"
        Just v  -> return v

lookupAll = lookupAllUsingLookupM

lookupAndDelete =
    lookupAndDelFromFM
      (error "TernaryTrie.lookupAndDelete: lookup failed")
      (,)

lookupAndDeleteM =
    lookupAndDelFromFM
      (fail  "TernaryTrie.lookupAndDeleteM: lookup failed")
      (\w m -> return (w,m))

lookupAndDeleteAll k m =
    lookupAndDelFromFM
      (S.empty,m)
      (\w m' -> (S.singleton w,m'))
      k m

lookupWithDefault = lookupWithDefaultUsingLookupM

adjust f k
  = addToFM k (\mv -> case mv of
                        Nothing -> mv
                        Just v  -> Just (f v))

adjustAll = adjust

adjustOrInsert f z k
  = addToFM k (\mv -> case mv of
                        Nothing -> Just z
                        Just v  -> Just (f v))

adjustAllOrInsert = adjustOrInsert

adjustOrDelete f k
  = addToFM k (\mv -> case mv of
                        Nothing -> mv
                        Just v  -> f v)

adjustOrDeleteAll = adjustOrDelete

map f
  = mapVFM (\mv -> case mv of
                        Nothing -> Nothing
                        Just v  -> Just (f v))

fold = foldr
fold' = foldr'

foldr op z (FM n fmb)
  = foldMV n . foldFMB fmb $ z
  where
    foldMV Nothing  = id
    foldMV (Just v) = op v

    foldFMB E
      = id
    foldFMB (I _ _ v l (FMB' m) r)
      = foldFMB l . foldMV v . foldFMB m . foldFMB r

foldrWithKey f z (FM n fmb)
  = foldMV [] n . foldFMB id fmb $ z
  where
     foldMV _ Nothing  = id
     foldMV ks (Just v) = f ks v

     foldFMB _ E = id
     foldFMB kf (I _ k mv l (FMB' m) r)
       = foldFMB kf l . foldMV (kf [k]) mv . foldFMB (kf . (k:)) m . foldFMB kf r

foldlWithKey f z (FM n fmb)
  = foldFMB id fmb . foldMV [] n $ z
  where
     g k x a = f a k x

     foldMV _ Nothing  = id
     foldMV ks (Just v) = g ks v

     foldFMB _ E = id
     foldFMB kf (I _ k mv l (FMB' m) r)
       = foldFMB kf r . foldFMB (kf . (k:)) m . foldMV (kf [k]) mv . foldFMB kf l

foldrWithKey' = foldrWithKey
foldlWithKey' = foldlWithKey

foldl :: (a -> b -> a) -> a -> FM t b -> a
foldl op z (FM n fmb)
  = foldFMB fmb . foldMV n $ z
  where
    foldMV Nothing  = id
    foldMV (Just v) = (flip op) v

    foldFMB E = id
    foldFMB (I _ _ v l (FMB' m) r)
      = foldFMB r . foldFMB m . foldMV v . foldFMB l


-- FIXME, understand this code to strictify it
foldr' = foldr
foldl' :: (a -> b -> a) -> a -> FM t b -> a
foldl' = foldl

foldr1 f fm =
  case maxView fm of
     Just (z,fm') -> foldr f z fm'
     Nothing      -> error $ moduleName++".foldr1: empty map"

foldl1 :: (b -> b -> b) -> FM k b -> b
foldl1 f fm =
  case minView fm of
     Just (z,fm') -> foldl f z fm'
     Nothing      -> error $ moduleName++".foldl1: empty map"


basecase :: Maybe t1 -> (t1 -> t) -> t -> t
basecase Nothing  = \_ n -> n
basecase (Just x) = \j _ -> j x

comb ::                                (t1 -> t1 -> t1)
                                    -> ((t1 -> t2) -> t2 -> t3)
                                    -> ((t1 -> t) -> t -> t2)
                                    -> (t1 -> t)
                                    -> t
                                    -> t3
comb f p1 p2
   = \j n -> p1 (\x -> p2 (\y -> j (f x y)) (j x)) (p2 j n)

fold1 f (FM mv fmb)
  = comb f (basecase mv) (fold1FMB fmb) id (error $ moduleName++".fold1: empty map")
  where
      fold1FMB E
        = \_ n -> n
      fold1FMB (I _ _ mv l (FMB' m) r)
        = comb f (basecase mv) $ comb f (fold1FMB l) $ comb f (fold1FMB m) $ (fold1FMB r)

fold1' = fold1

{-
FIXME -- can these be somehow fixed to have the right order...

foldr1 f (FM v fmb)
  = comb f (basecase v) (fold1FMB fmb) id (error $ moduleName++".foldr1: empty map")
  where
      fold1FMB E
        = \j n -> n
      fold1FMB (I _ _ v l (FMB' m) r)
        = comb f (fold1FMB l) $ comb f (basecase v) $ comb f (fold1FMB m) $ (fold1FMB r)


foldl1 f (FM v fmb)
  = comb f (fold1FMB fmb) (basecase v) id (error $ moduleName++".foldl1: empty map")
  where
      fold1FMB E
        = \j n -> n
      fold1FMB (I _ _ v l (FMB' m) r)
        = comb f (fold1FMB r) $ comb f (fold1FMB m) $ comb f (basecase v) $ (fold1FMB l)
-}



-- FIXME, understand this code to strictify it
foldr1' = foldr1
foldl1' :: (b -> b -> b) -> FM k b -> b
foldl1' = foldl1


filter p = mapVFM (\mv -> case mv of
                            Nothing -> mv
                            Just v  -> if p v then mv else Nothing)

partition = partitionUsingFilter

elements = elementsUsingFold

strict z@(FM _ fmb) = strictFMB fmb `seq` z
 where strictFMB n@E = n
       strictFMB n@(I _ _ _ l (FMB' m) r) =
           strictFMB l `seq` strictFMB m `seq` strictFMB r `seq` n

strictWith f z@(FM v fmb) = f' v `seq` strictWithFMB fmb `seq` z
   where f' v@Nothing  = v
         f' v@(Just x) = f x `seq` v

         strictWithFMB n@E = n
         strictWithFMB n@(I _ _ v l (FMB' m) r) =
           f' v `seq` strictWithFMB l `seq` strictWithFMB m `seq` strictWithFMB r `seq` n


-- FiniteMapX

fromSeqWith = fromSeqWithUsingInsertSeqWith

fromSeqWithKey = fromSeqWithKeyUsingInsertSeqWithKey

insertWith f k v
  = addToFM k (\vem ->
      case vem of
          Nothing -> Just v
          Just ve -> Just (f ve v))

insertWithKey = insertWithKeyUsingInsertWith

insertSeqWith = insertSeqWithUsingInsertWith

insertSeqWithKey = insertSeqWithKeyUsingInsertWithKey

unionl = union
unionr = flip union

unionWith f = unionWithKey (const f)

unionSeqWith = unionSeqWithUsingReduce

intersectionWith f = intersectionWithKey (const f)

difference mx my
  = mergeVFM (\v1 v2 -> case v2 of
              Nothing -> v1
              Just _  -> Nothing) mx my

properSubset = properSubsetUsingSubset

subset (FM nx fmbx) (FM ny fmby)
  = subsetEqM nx ny && subsetEqFMB fmbx fmby
  where
    subsetEqM Nothing _ = True
    subsetEqM (Just _) Nothing = False
    subsetEqM (Just _) (Just _) = True

    subsetEqFMB E _ = True
    subsetEqFMB fmbx@(I _ _ _ _ _ _) E
      = nullFMB fmbx
    subsetEqFMB fmbx@(I sizex kx vx lx (FMB' mx) rx)
            fmby@(I sizey ky vy ly (FMB' my) ry)
      | sizex >= sizey
        = let (vy, ly, FMB' my, ry) = splayFMB kx fmby
          in    subsetEqM vx vy
             && subsetEqFMB lx ly
             && subsetEqFMB mx my
             && subsetEqFMB rx ry
      | otherwise
        = let (vx, lx, FMB' mx, rx) = splayFMB ky fmbx
          in    subsetEqM vx vy
             && subsetEqFMB lx ly
             && subsetEqFMB mx my
             && subsetEqFMB rx ry


submapBy = submapByUsingLookupM
properSubmapBy = properSubmapByUsingSubmapBy
sameMapBy = sameMapByUsingSubmapBy
properSubmap = A.properSubmap
submap = A.submap
sameMap = A.sameMap

-- Assoc

toSeq = toSeqUsingFoldWithKey

keys = keysUsingFoldWithKey

mapWithKey f
  = mapKVFM (\k mv -> case mv of
          Nothing -> Nothing
          Just v  -> Just (f k v))

foldWithKey op r (FM n fmb)
  = foldWithKeyB [] n . foldWithKeyFM [] fmb $ r
  where
      foldWithKeyB _ Nothing = id
      foldWithKeyB k (Just v) = op k v

      foldWithKeyFM _ E = id
      foldWithKeyFM ks (I _ k v l (FMB' m) r)
        = foldWithKeyFM ks l
        . foldWithKeyB (reverse (k:ks)) v
        . foldWithKeyFM (k:ks) m
        . foldWithKeyFM ks r


-- FIXME, make this strict
foldWithKey' = foldWithKey


filterWithKey f
  = mapKVFM (\k mv -> case mv of
          Nothing -> mv
          Just v  -> if f k v then mv else Nothing)

partitionWithKey f m
  = (filterWithKey f m, filterWithKey (\k v -> not (f k v)) m)

-- FiniteMap

unionWithKey f
  = mergeKVFM (\k v1m v2m ->
    case v1m of
        Nothing -> v2m
        Just v1 ->
            case v2m of
            Nothing -> v1m
            Just v2 -> Just (f k v1 v2))

unionSeqWithKey = unionSeqWithKeyUsingReduce

intersectionWithKey f
  = mergeKVFM (\k v1m v2m ->
    case v1m of
        Nothing -> Nothing
        Just v1 ->
            case v2m of
            Nothing -> Nothing
            Just v2 -> Just (f k v1 v2))

-- OrdAssocX

minViewFMB :: Fail.MonadFail m => FMB k a -> (FMB k a -> FM k a) -> m (a, FM k a)
minViewFMB E _ = fail $ moduleName++".minView: empty map"
minViewFMB (I i k (Just v) E m r)        f = return (v, f t)
  where
    t = case m of
      FMB' E -> r
      _ -> I i k Nothing E m r
minViewFMB (I _ _ Nothing  E (FMB' E) _) _ = error $ moduleName++".minView: bug!"
minViewFMB (I _ k Nothing  E (FMB' m) r) f = minViewFMB m (\m' -> f (mkVBalancedFMB k Nothing E (FMB' m') r))
minViewFMB (I _ k mv l m r)              f = minViewFMB l (\l' -> f (mkVBalancedFMB k mv l' m r))

minView :: Fail.MonadFail m => FM k a -> m (a,FM k a)
minView (FM (Just v) fmb) = return (v, FM Nothing fmb)
minView (FM Nothing fmb)  = minViewFMB fmb (FM Nothing)

minViewWithKeyFMB :: Fail.MonadFail m => FMB k a -> ([k] -> [k]) -> (FMB k a -> FM k a) -> m (([k],a),FM k a)
minViewWithKeyFMB E _ _ = fail $ moduleName++".minView: empty map"
minViewWithKeyFMB (I i k (Just v) E m r)        kf f = return ((kf [k],v),f t)
  where
    t = case m of
      FMB' E -> r
      _ -> I i k Nothing E m r
minViewWithKeyFMB (I _ _ Nothing  E (FMB' E) _) _ _ = error $ moduleName++".minViewWithKey: bug!"
minViewWithKeyFMB (I _ k Nothing  E (FMB' m) r) kf f = minViewWithKeyFMB m (kf . (k:))
                                                        (\m' -> f (mkVBalancedFMB k Nothing E (FMB' m') r))
minViewWithKeyFMB (I _ k mv l m r)              kf f = minViewWithKeyFMB l kf
                                                        (\l' -> f (mkVBalancedFMB k mv l' m r))

minViewWithKey :: Fail.MonadFail m => FM k a -> m (([k],a),FM k a)
minViewWithKey (FM (Just v) fmb) = return (([],v),FM Nothing fmb)
minViewWithKey (FM Nothing fmb)  = minViewWithKeyFMB fmb id (FM Nothing)


minElemFMB :: FMB k a -> a
minElemFMB E = error $ moduleName++".minElem: empty map"
minElemFMB (I _ _ (Just v) E _ _)        = v
minElemFMB (I _ _ Nothing  E (FMB' m) _) = minElemFMB m
minElemFMB (I _ _ _ l _ _)              = minElemFMB l

minElem :: FM t1 t -> t
minElem (FM (Just v) _) = v
minElem (FM Nothing  fmb) = minElemFMB fmb


minElemWithKeyFMB :: ([k] -> [k]) -> FMB k a -> ([k],a)
minElemWithKeyFMB _ E = error $ moduleName++".minElemWithKey: empty map"
minElemWithKeyFMB kf (I _ k (Just v) E _ _)        = (kf [k],v)
minElemWithKeyFMB kf (I _ k Nothing  E (FMB' m) _) = minElemWithKeyFMB (kf . (k:)) m
minElemWithKeyFMB kf (I _ _ _ l _ _)              = minElemWithKeyFMB kf l

minElemWithKey :: FM k a -> ([k],a)
minElemWithKey (FM (Just v) _) = ([],v)
minElemWithKey (FM Nothing  fmb) = minElemWithKeyFMB id fmb

deleteMin :: Ord k => FM k a -> FM k a
deleteMin = deleteMinUsingMinView

unsafeInsertMin :: Ord k => [k] -> a -> FM k a -> FM k a
unsafeInsertMin = insert

maxViewFMB :: Fail.MonadFail m => FMB k a -> (FMB k a -> FM k a) -> m (a, FM k a)
maxViewFMB (I _ _ (Just v) l (FMB' E) E) f = return (v, f l)
--maxViewFMB (I i k (Just v) l (FMB' E) E) f = return (v, f (I i k Nothing l (FMB' E) E))
maxViewFMB (I _ _ Nothing  _ (FMB' E) E) _ = error $ moduleName++".maxView: bug!"
maxViewFMB (I i k mv l (FMB' m) E)       f = maxViewFMB m (\m' -> f (t m'))
  where
    t m' = case m' of
      E | Nothing <- mv -> l
      _ -> I i k mv l (FMB' m') E
maxViewFMB (I _ k mv l m r)              f = maxViewFMB r (\r' -> f (mkVBalancedFMB k mv l m r'))
maxViewFMB E                             _ = error $ moduleName++".maxView: bug!"

maxView :: Fail.MonadFail m => FM k a -> m (a, FM k a)
maxView (FM Nothing E)  = fail $ moduleName++".maxView: empty map"
maxView (FM (Just v) E) = return (v,FM Nothing E)
maxView (FM mv fmb)     = maxViewFMB fmb (FM mv)


maxViewWithKeyFMB :: Monad m => FMB k a -> ([k] -> [k]) -> (FMB k a -> FM k a) -> m (([k],a),FM k a)
maxViewWithKeyFMB (I _ k (Just v) l (FMB' E) E) kf f = return ((kf [k],v),f l)
maxViewWithKeyFMB (I _ _ Nothing  _ (FMB' E) E) _ _ = error $ moduleName++".maxViewWithKey: bug!"
maxViewWithKeyFMB (I i k mv l (FMB' m) E)       kf f = maxViewWithKeyFMB m (kf . (k:))
                                                        (\m' -> f (t m'))
  where
    t m' = case m' of
      E | Nothing <- mv -> l
      _ -> I i k mv l (FMB' m') E
maxViewWithKeyFMB (I _ k mv l m r)              kf f = maxViewWithKeyFMB r kf
                                                        (\r' -> f (mkVBalancedFMB k mv l m r'))
maxViewWithKeyFMB E                             _ _ = error $ moduleName++".maxViewWithKey: bug!"


maxViewWithKey :: Fail.MonadFail m => FM k a -> m (([k],a), FM k a)
maxViewWithKey (FM Nothing E)  = fail $ moduleName++".maxViewWithKey: empty map"
maxViewWithKey (FM (Just v) E) = return (([],v),FM Nothing E)
maxViewWithKey (FM mv fmb)     = maxViewWithKeyFMB fmb id (FM mv)



maxElemFMB :: FMB k a -> a
maxElemFMB (I _ _ (Just v) _ (FMB' E) E) = v
maxElemFMB (I _ _ Nothing  _ (FMB' E) E) = error $ moduleName++".maxElem: bug!"
maxElemFMB (I _ _ _ _ (FMB' m) E)       = maxElemFMB m
maxElemFMB (I _ _ _ _ _ r)              = maxElemFMB r
maxElemFMB E                             = error $ moduleName++".maxElem: bug!"

maxElem :: FM k a -> a
maxElem (FM (Just v) E) = v
maxElem (FM Nothing  E) = error $ moduleName++".maxElem: empty map"
maxElem (FM _ fmb)      = maxElemFMB fmb

maxElemWithKeyFMB :: FMB k a -> ([k] -> [k]) -> ([k],a)
maxElemWithKeyFMB (I _ k (Just v) _ (FMB' E) E) kf = (kf [k],v)
maxElemWithKeyFMB (I _ _ Nothing  _ (FMB' E) E) _ = error $ moduleName++".maxElemWithKey: bug!"
maxElemWithKeyFMB (I _ k _ _ (FMB' m) E)       kf = maxElemWithKeyFMB m (kf . (k:))
maxElemWithKeyFMB (I _ _ _ _ _ r)              kf = maxElemWithKeyFMB r kf
maxElemWithKeyFMB E                             _ = error $ moduleName++".maxElemWithKey: bug!"


maxElemWithKey :: FM k a -> ([k],a)
maxElemWithKey (FM (Just v) E) = ([],v)
maxElemWithKey (FM Nothing E)  = error $ moduleName++".maxElemWithKey: empty map"
maxElemWithKey (FM _ fmb)      = maxElemWithKeyFMB fmb id


deleteMax :: Ord k => FM k a -> FM k a
deleteMax = deleteMaxUsingMaxView

unsafeInsertMax :: Ord k => [k] -> a -> FM k a -> FM k a
unsafeInsertMax = insert

unsafeFromOrdSeq :: (Ord k,S.Sequence seq) => seq ([k],a) -> FM k a
unsafeFromOrdSeq = fromSeq

unsafeAppend :: Ord k => FM k a -> FM k a -> FM k a
unsafeAppend = union

-- FIXME this doesn't respect the structural invariant... why??
{-
unsafeAppend (FM (Just v) fmb1) (FM Nothing fmb2) = FM (Just v) (appendFMB fmb1 fmb2)
unsafeAppend (FM Nothing  fmb1) (FM mv fmb2)      = FM mv       (appendFMB fmb1 fmb2)
unsafeAppend (FM (Just _) _) (FM (Just _) _)      = error $ moduleName++".unsafeAppend: bug!"
-}

filterL_FMB :: Ord k => (k -> Maybe a -> FMB k a -> FMB k a) -> k -> [k] -> FMB k a -> FMB k a
filterL_FMB _ _ _ E = E
filterL_FMB f k ks (I _ key mv l (FMB' m) r)
    | key < k   = mkVBalancedFMB key mv l (FMB' m) (filterL_FMB f k ks r)
    | key > k   = filterL_FMB f k ks l
    | otherwise = case ks of
                    []       -> f k mv l
                    (k':ks') -> mkVBalancedFMB key mv l (FMB' (filterL_FMB f k' ks' m)) E

filterLT :: Ord k => [k] -> FM k a -> FM k a
filterLT [] _               = FM Nothing E
filterLT (k:ks) (FM mv fmb) = FM mv (filterL_FMB (\_ _ l -> l) k ks fmb)

filterLE :: Ord k => [k] -> FM k a -> FM k a
filterLE [] (FM mv _)       = FM mv E
filterLE (k:ks) (FM mv fmb) = FM mv (filterL_FMB (\k mv l -> mkVBalancedFMB k mv l (FMB' E) E) k ks fmb)



filterG_FMB :: Ord k => (k -> Maybe a -> FMB k a -> FMB k a -> FMB k a) -> k -> [k] -> FMB k a -> FMB k a
filterG_FMB _ _ _ E = E
filterG_FMB f k ks (I _ key mv l (FMB' m) r)
    | key < k   = filterG_FMB f k ks r
    | key > k   = mkVBalancedFMB key mv (filterG_FMB f k ks l) (FMB' m) r
    | otherwise = case ks of
                    []       -> f k mv m r
                    (k':ks') -> mkVBalancedFMB key Nothing E (FMB' (filterG_FMB f k' ks' m)) r

filterGT :: Ord k => [k] -> FM k a -> FM k a
filterGT []     (FM _  fmb) = FM Nothing fmb
filterGT (k:ks) (FM _ fmb) = FM Nothing (filterG_FMB (\k _ m r -> mkVBalancedFMB k Nothing E (FMB' m) r) k ks fmb)

filterGE :: Ord k => [k] -> FM k a -> FM k a
filterGE []     fm          = fm
filterGE (k:ks) (FM _ fmb) = FM Nothing (filterG_FMB (\k mv m r -> mkVBalancedFMB k mv E (FMB' m) r) k ks fmb)

--FIXME do better...
partitionLT_GE :: Ord k => [k] -> FM k a -> (FM k a,FM k a)
partitionLT_GE ks fm = (filterLT ks fm, filterGE ks fm)

partitionLE_GT :: Ord k => [k] -> FM k a -> (FM k a,FM k a)
partitionLE_GT ks fm = (filterLE ks fm, filterGT ks fm)

partitionLT_GT :: Ord k => [k] -> FM k a -> (FM k a,FM k a)
partitionLT_GT ks fm = (filterLT ks fm, filterGT ks fm)

toOrdSeq = toOrdSeqUsingFoldrWithKey

-- instance declarations

instance Ord k  => A.AssocX (FM k) [k] where
  {empty = empty; singleton = singleton; fromSeq = fromSeq; insert = insert;
   insertSeq = insertSeq; union = union; unionSeq = unionSeq;
   delete = delete; deleteAll = deleteAll; deleteSeq = deleteSeq;
   null = null; size = size; member = member; count = count;
   lookup = lookup; lookupM = lookupM; lookupAll = lookupAll;
   lookupAndDelete = lookupAndDelete; lookupAndDeleteM = lookupAndDeleteM;
   lookupAndDeleteAll = lookupAndDeleteAll;
   lookupWithDefault = lookupWithDefault; adjust = adjust;
   adjustAll = adjustAll; adjustOrInsert = adjustOrInsert;
   adjustAllOrInsert = adjustAllOrInsert;
   adjustOrDelete = adjustOrDelete; adjustOrDeleteAll = adjustOrDeleteAll;
   fold = fold; fold' = fold'; fold1 = fold1; fold1' = fold1';
   filter = filter; partition = partition; elements = elements;
   strict = strict; strictWith = strictWith;
   structuralInvariant = structuralInvariant; instanceName _ = moduleName}

instance Ord k  => A.Assoc (FM k) [k] where
  {toSeq = toSeq; keys = keys; mapWithKey = mapWithKey;
   foldWithKey = foldWithKey; foldWithKey' = foldWithKey';
   filterWithKey = filterWithKey;
   partitionWithKey = partitionWithKey}

instance Ord k => A.FiniteMapX (FM k) [k] where
  {fromSeqWith = fromSeqWith; fromSeqWithKey = fromSeqWithKey;
   insertWith  = insertWith; insertWithKey = insertWithKey;
   insertSeqWith = insertSeqWith; insertSeqWithKey = insertSeqWithKey;
   unionl = unionl; unionr = unionr; unionWith = unionWith;
   unionSeqWith = unionSeqWith; intersectionWith = intersectionWith;
   difference = difference; properSubset = properSubset; subset = subset;
   properSubmapBy = properSubmapBy; submapBy = submapBy;
   sameMapBy = sameMapBy}

instance Ord k => A.FiniteMap (FM k) [k] where
  {unionWithKey = unionWithKey; unionSeqWithKey = unionSeqWithKey;
   intersectionWithKey = intersectionWithKey}

instance Ord k => A.OrdAssocX (FM k) [k] where
  {minView = minView; minElem = minElem; deleteMin = deleteMin;
   unsafeInsertMin = unsafeInsertMin; maxView = maxView; maxElem = maxElem;
   deleteMax = deleteMax; unsafeInsertMax = unsafeInsertMax;
   foldr = foldr; foldr' = foldr'; foldl = foldl; foldl' = foldl';
   foldr1 = foldr1; foldr1' = foldr1'; foldl1 = foldl1; foldl1' = foldl1';
   unsafeFromOrdSeq = unsafeFromOrdSeq; unsafeAppend = unsafeAppend;
   filterLT = filterLT; filterLE = filterLE; filterGT = filterGT;
   filterGE = filterGE;  partitionLT_GE = partitionLT_GE;
   partitionLE_GT = partitionLE_GT; partitionLT_GT = partitionLT_GT}

instance Ord k => A.OrdAssoc (FM k) [k] where
  {minViewWithKey = minViewWithKey; minElemWithKey = minElemWithKey;
   maxViewWithKey = maxViewWithKey; maxElemWithKey = maxElemWithKey;
   foldrWithKey = foldrWithKey; foldrWithKey' = foldrWithKey';
   foldlWithKey = foldlWithKey; foldlWithKey' = foldlWithKey';
   toOrdSeq = toOrdSeq}

instance Ord k => A.OrdFiniteMapX (FM k) [k]
instance Ord k => A.OrdFiniteMap (FM k) [k]


instance Ord k => Functor (FM k) where
  fmap = map

instance (Ord k, Show k, Show a) => Show (FM k a) where
  showsPrec = showsPrecUsingToList

instance (Ord k, Read k, Read a) => Read (FM k a) where
  readsPrec = readsPrecUsingFromList

instance (Ord k, Eq a) => Eq (FM k a) where
  (==) = sameMap

instance (Ord k, Ord a) => Ord (FM k a) where
  compare = compareUsingToOrdList

--
-- Test code follows
--

keyInvariantFMB :: Ord k => (k -> Bool) -> FMB k a -> Bool
keyInvariantFMB _ E = True
keyInvariantFMB p (I _ k _ l _ r)
  =    p k
    && keyInvariantFMB p l
    && keyInvariantFMB p r

actualSizeFMB :: FMB k a -> Int
actualSizeFMB E = 0
actualSizeFMB (I _ _ _ l _ r) = 1 + actualSizeFMB l + actualSizeFMB r

structuralInvariantFMB :: Ord k => FMB k a -> Bool
structuralInvariantFMB E = True
structuralInvariantFMB fmb@(I size k v l (FMB' m) r)
  =    structuralInvariantFMB l
    && structuralInvariantFMB m
    && structuralInvariantFMB r
    && keyInvariantFMB (<k) l
    && keyInvariantFMB (>k) r
    && actualSizeFMB fmb == size
    && isBalanced l r
    && relevantRoot fmb

isBalanced :: FMB k a -> FMB k a -> Bool
isBalanced l r = sizel + sizer <= 1
  || (sizel <= balance * sizer && sizer <= balance * sizel)
  where
      sizel = sizeFMB l
      sizer = sizeFMB r

-- | This invariant is used by minView
relevantRoot :: FMB k a -> Bool
relevantRoot (I _ _ Nothing _ (FMB' E) _) = False
relevantRoot _ = True

structuralInvariant :: Ord k => FM k a -> Bool
structuralInvariant (FM _ fmb) = structuralInvariantFMB fmb

-- | Generate weight-balanced trees either by direct recursion or via
-- 'fromSeq'. The former is much more likely to hit counterexamples to wrong
-- @balance@ coefficients. We keep the latter generator around just in case,
-- because it generates a more realistic distribution.
instance (Integral k, Arbitrary k, Arbitrary a) => Arbitrary (FM k a) where
  arbitrary = oneof [genFM, fromSeq <$> (arbitrary :: Gen [([k], a)])]
  shrink (FM v m) = [FM v m | (v, FMB' m) <- shrinkTuple shrink shrinkFMB' (v, FMB' m)]

instance (Ord k,CoArbitrary k,CoArbitrary a) => CoArbitrary (FM k a) where
  coarbitrary (FM x fmb) = coarbitrary_maybe x . coarbitrary_fmb fmb


coarbitrary_maybe :: (CoArbitrary t) => Maybe t  -> Test.QuickCheck.Gen b
                                                 -> Test.QuickCheck.Gen b
coarbitrary_maybe Nothing = variant (0 :: Int)
coarbitrary_maybe (Just x) = variant (1 :: Int) . coarbitrary x

coarbitrary_fmb :: (CoArbitrary t1, CoArbitrary t) => FMB t t1 -> Gen a -> Gen a
coarbitrary_fmb E = variant (0 :: Int)
coarbitrary_fmb (I _ k x l (FMB' m) r) =
        variant (1 :: Int) . coarbitrary k . coarbitrary_maybe x .
        coarbitrary_fmb l . coarbitrary_fmb m . coarbitrary_fmb r

instance Ord k => Semigroup (FM k a) where
   (<>) = union
instance Ord k => Monoid (FM k a) where
   mempty  = empty
   mappend = (SG.<>)
   mconcat = unionSeq

-- Testing

genFM :: (Integral k, Arbitrary a) => Gen (FM k a)
genFM = do
  FM <$> arbitrary <*> genFMB_

-- Choose the number of elements in the top layer upfront,
-- and distribute it while recursing down.
genFMB_ :: (Integral k, Arbitrary a) => Gen (FMB k a)
genFMB_ = sized $ \sz -> do
  n <- choose (0, sz)
  resize (sz - n) (genFMB 0 n)

-- Distribute the size @sz@ to generate the middle children of the nodes in the
-- top layer.
genFMB :: (Integral k, Arbitrary a) => Int -> Int -> Gen (FMB k a)
genFMB i 0 = pure E
genFMB i n = sized $ \sz -> do
  let b = if n <= 2 then 0 else (n-1+balance) `div` (balance+1)
  l <- choose (b, n-1-b)
  z <- choose (0, sz)
  m <- resize (min z (sz-z)) genFMB_
  v <- case m of E -> Just <$> arbitrary ; _ -> arbitrary
  let k = fromIntegral (i+l)
  I n k v
    <$> resize z (genFMB i l)
    <*> pure (FMB' m)
    <*> resize (sz - z) (genFMB (i+l+1) (n-l-1))

-- Be careful to preserve balance during shrinking.
shrinkFMB :: Arbitrary a => FMB k a -> [FMB k a]
shrinkFMB E = []
shrinkFMB (I s k v l m r) = E : l : r : do
    let (*-) = shrinkTuple ; infixr 3 *-
    (v, (l, (m@(FMB' m'), r))) <- (shrinkJust *- shrinkFMB *- shrinkFMB' *- shrinkFMB) (v, (l, (m, r)))
    let s = sizeFMB l + sizeFMB r + 1
        t = I s k v l m r
    guard (isBalanced l r && (isJust v || not (nullFMB' m)))
    pure t

nullFMB' :: FMB' k v -> Bool
nullFMB' (FMB' E) = True
nullFMB' _ = False

shrinkFMB' :: Arbitrary a => FMB' k a -> [FMB' k a]
shrinkFMB' (FMB' m) = coerce $
  tailsFMB m ++ shrinkFMB m

-- List the middle children of the top layer.
tailsFMB :: FMB k a -> [FMB k a]
tailsFMB E = []
tailsFMB (I _ _ _ l (FMB' m) r) = m : tailsFMB l ++ tailsFMB r

-- Don't remove elements
shrinkJust :: Arbitrary a => Maybe a -> [Maybe a]
shrinkJust Nothing = []
shrinkJust (Just x) = Just <$> shrink x

shrinkTuple :: (a -> [a]) -> (b -> [b]) -> (a, b) -> [(a, b)]
shrinkTuple sa sb (a, b) = [(a', b) | a' <- sa a] ++ [(a, b') | b' <- sb b]
