content: {-# LANGUAGE DataKinds, GADTs, KindSignatures, StandaloneDeriving, RankNTypes,
      TypeSynonymInstances #-}
  
  module BigArray where
  
  import Data.Monoid
  import Control.Applicative
  import Data.Functor.Identity
  
  data Nat = Ze | Su Nat
  
  data Bnd k = Bot | Key k | Top deriving (Ord, Eq, Show)
  
  data T23 (n :: Nat) (k :: *) (v :: *) where
    Leaf   :: T23 Ze k v
    Node2  :: T23 n k v -> (k, v) -> T23 n k v -> T23 (Su n) k v
    Node3  :: T23 n k v -> (k, v) -> T23 n k v -> (k, v) -> T23 n k v -> T23 (Su n) k v
  deriving instance (Show k, Show v) => Show (T23 n k v)
  
  instance Functor (T23 n k) where
    fmap f Leaf = Leaf
    fmap f (Node2 l (k, v) r) = Node2 (fmap f l) (k, f v) (fmap f r)
    fmap f (Node3 l (j, u) m (k, v) r) =
      Node3 (fmap f l) (j, f u) (fmap f m) (k, f v) (fmap f r)
  
  data I23 n k v where
    Level :: T23 n k v -> I23 n k v
    Grow2 :: T23 n k v -> (k, v) -> T23 n k v -> I23 n k v
  
  find23 :: Ord k => k -> T23 n k v -> Maybe v
  find23 k Leaf = Nothing
  find23 k (Node2 lj (j, jv) ju) = case compare k j of
    LT -> find23 k lj
    EQ -> Just jv
    GT -> find23 k ju
  find23 k (Node3 li (i, iv) ij (j, jv) ju) = case (compare k i, compare k j) of
    (LT, _) -> find23 k li
    (EQ, _) -> Just iv
    (_, LT) -> find23 k ij
    (_, EQ) -> Just jv
    (_, GT) -> find23 k ju
  
  insert23 :: Ord k => (k, v) -> T23 n k v -> I23 n k v
  insert23 iv Leaf = Grow2 Leaf iv Leaf
  insert23 iv@(i, _) (Node2 lj jv@(j, _) ju) =
    case compare i j of
      LT -> case insert23 iv lj of
        Grow2 lh hv hj -> Level (Node3 lh hv hj jv ju)
        Level lj       -> Level (Node2 lj jv ju)
      EQ -> Level (Node2 lj iv ju)
      GT -> case insert23 iv ju of
        Grow2 jk kv ku -> Level (Node3 lj jv jk kv ku)
        Level ju       -> Level (Node2 lj jv ju)
  insert23 iv@(i, _) (Node3 lj jv@(j, _) jk kv@(k, _) ku) =
    case (compare i j, compare i k) of
      (LT, _) -> case insert23 iv lj of
        Grow2 lh hv hj -> Grow2 (Node2 lh hv hj) jv (Node2 jk kv ku)
        Level lj       -> Level (Node3 lj jv jk kv ku)
      (EQ, _) -> Level (Node3 lj iv jk kv ku)
      (_, LT) -> case insert23 iv jk of
        Grow2 jh hv hk -> Grow2 (Node2 lj jv jh) hv (Node2 hk kv ku)
        Level jk       -> Level (Node3 lj jv jk kv ku)
      (_, EQ) -> Level (Node3 lj jv jk iv ku)
      (_, GT) -> case insert23 iv ku of
        Grow2 kh hv hu -> Grow2 (Node2 lj jv jk) kv (Node2 kh hv hu)
        Level ku       -> Level (Node3 lj jv jk kv ku)
  
  anoin23 :: Ord k => T23 n k v -> T23 n k v -> I23 n k v
  anoin23 Leaf Leaf = Level Leaf
  anoin23 (Node2 ai i' im) (Node2 mr r' rz) = case anoin23 im mr of
    Level ir       -> Level (Node3 ai i' ir r' rz)
    Grow2 im m' mr -> Grow2 (Node2 ai i' im) m' (Node2 mr r' rz)
  anoin23 (Node2 ai i' im) (Node3 mp p' pu u' uz) = case anoin23 im mp of
    Level ip       -> Grow2 (Node2 ai i' ip) p' (Node2 pu u' uz)
    Grow2 im m' mp -> Grow2 (Node2 ai i' im) m' (Node3 mp p' pu u' uz)
  anoin23 (Node3 ag g' gj j' jm) (Node2 mr r' rz) = case anoin23 jm mr of
    Level jr       -> Grow2 (Node2 ag g' gj) j' (Node2 jr r' rz)
    Grow2 jm m' mr -> Grow2 (Node3 ag g' gj j' jm) m' (Node2 mr r' rz)
  anoin23 (Node3 ag g' gj j' jm) (Node3 mp p' pu u' uz) = case anoin23 jm mp of
    Level jp       -> Grow2 (Node2 ag g' gj) j' (Node3 jp p' pu u' uz)
    Grow2 jm m' mp -> Grow2 (Node3 ag g' gj j' jm) m' (Node3 mp p' pu u' uz)
  
  data D23 n k v where
    DSame :: T23 n k v -> D23 n k v
    DDrop :: T23 n k v -> D23 (Su n) k v
  
  i2dSu :: I23 n k v -> D23 (Su n) k v
  i2dSu (Level t) = DDrop t
  i2dSu (Grow2 li i' iu) = DSame (Node2 li i' iu)
  
  dtNode :: D23 n k v -> (k, v) -> T23 n k v -> D23 (Su n) k v
  dtNode (DSame am) m' mz                     = DSame (Node2 am m' mz)
  dtNode (DDrop am) m' (Node2 mr r' rz)       = DDrop (Node3 am m' mr r' rz)
  dtNode (DDrop am) m' (Node3 mp p' ps s' sz) =
    DSame (Node2 (Node2 am m' mp) p' (Node2 ps s' sz))
  
  tdNode :: T23 n k v -> (k, v) -> D23 n k v -> D23 (Su n) k v
  tdNode am m' (DSame mz) = DSame (Node2 am m' mz)
  tdNode (Node2 af f' fm) m' (DDrop mz) = DDrop (Node3 af f' fm m' mz)
  tdNode (Node3 ad d' dh h' hm) m' (DDrop mz) =
    DSame (Node2 (Node2 ad d' dh) h' (Node2 hm m' mz))
  
  delete23 :: Ord k => k -> T23 n k v -> D23 n k v
  delete23 k Leaf = DSame Leaf
  delete23 k (Node2 ai i' iz) = case compare k (fst i') of
    LT -> dtNode (delete23 k ai) i' iz
    EQ -> i2dSu (anoin23 ai iz)
    GT -> tdNode ai i' (delete23 k iz)
  delete23 k (Node3 ai i' ip p' pz) = case (compare k (fst i'), compare k (fst p')) of
    (LT, _) -> case delete23 k ai of
      DSame ai -> DSame (Node3 ai i' ip p' pz)
      DDrop ai -> case ip of
        Node2 im m' mp       ->
          DSame (Node2 (Node3 ai i' im m' mp) p' pz)
        Node3 il l' ln n' np ->
          DSame (Node3 (Node2 ai i' il) l' (Node2 ln n' np) p' pz)
    (EQ, _) -> case anoin23 ai ip of
      Level ap        -> DSame (Node2 ap p' pz)
      Grow2 ai i' ip  -> DSame (Node3 ai i' ip p' pz)
    (_, LT) -> case delete23 k ip of
      DSame ip -> DSame (Node3 ai i' ip p' pz)
      DDrop ip -> case pz of
        Node2 pt t' tz -> DSame (Node2 ai i' (Node3 ip p' pt t' tz))
        Node3 pr r' ru u' uz ->
          DSame (Node3 ai i' (Node2 ip p' pr) r' (Node2 ru u' uz))
    (_, EQ) -> case anoin23 ip pz of
      Level iz        -> DSame (Node2 ai i' iz)
      Grow2 ip p' pz  -> DSame (Node3 ai i' ip p' pz)
    (_, GT) -> case delete23 k pz of
      DSame pz -> DSame (Node3 ai i' ip p' pz)
      DDrop pz -> case ip of
        Node2 im m' mp       ->
          DSame (Node2 ai i' (Node3 im m' mp p' pz))
        Node3 il l' ln n' np ->
          DSame (Node3 ai i' (Node2 il l' ln) n' (Node2 np p' pz))
  
  data Arr (k :: *)(v :: *) where
    Arr :: T23 n k v -> Arr k v
  deriving instance (Show k, Show v) => Show (Arr k v)
  
  instance Functor (Arr k) where
    fmap f (Arr t) = Arr (fmap f t)
  
  emptyArr :: Arr k v
  emptyArr = Arr Leaf
  
  insertArr :: Ord k => (k, v) -> Arr k v -> Arr k v
  insertArr iv (Arr lu) = case insert23 iv lu of
    Level lu       -> Arr lu
    Grow2 lj jv ju -> Arr (Node2 lj jv ju)
  
  findArr :: Ord k => k -> Arr k v -> Maybe v
  findArr k (Arr lu) = find23 k lu
  
  deleteArr :: Ord k => k -> Arr k v -> Arr k v
  deleteArr k (Arr lu) = case delete23 k lu of
    DSame lu -> Arr lu
    DDrop lu -> Arr lu
  
  single :: Ord k => (k, v) -> Arr k v
  single x = insertArr x emptyArr
  
  isEmptyArr :: Arr k v -> Bool
  isEmptyArr (Arr Leaf) = True
  isEmptyArr _ = False
  
  travT23 :: Applicative f => ((k, v) -> f w) -> T23 n k v -> f (T23 n k w)
  travT23 f Leaf = pure Leaf
  travT23 f (Node2 l x@(k, _) r) = Node2 <$> travT23 f l <*> ((,) k <$> f x) <*> travT23 f r
  travT23 f (Node3 l x@(j, _) m y@(k, _) r) =
    Node3 <$> travT23 f l <*> ((,) j <$> f x) <*> travT23 f m <*> ((,) k <$> f y) <*> travT23 f r
  
  travArr :: Applicative f => ((k, v) -> f w) -> Arr k v -> f (Arr k w)
  travArr f (Arr t) = Arr <$> travT23 f t
  
  foldMapArr :: Monoid x => ((k, v) -> x) -> Arr k v -> x
  foldMapArr f = getConst . travArr (Const . f)
  
  foldMapSet :: Monoid y => (x -> y) -> Set x -> y
  foldMapSet f = foldMapArr (f . fst)
  
  instance (Ord k, Semigroup v) => Semigroup (Arr k v) where (<>) = mappend
  instance (Ord k, Semigroup v) => Monoid (Arr k v) where
    mempty = emptyArr
    mappend l r = appEndo (foldMapArr up l) r where
      up (k, v) = Endo $ \ r -> case findArr k r of
        Just w  -> insertArr (k, v <> w) r
        Nothing -> insertArr (k, v) r
  
  type Set x = Arr x ()
  
  domain :: Ord k => Arr k v -> Set k
  domain = runIdentity . travArr (\ (k, _) -> Identity ())
  
  inSet :: Ord x => x -> Set x -> Bool
  inSet x s = case findArr x s of
    Just _ -> True
    _      -> False
  
  singleton :: Ord x => x -> Set x
  singleton = single . flip (,) ()
  
  leftmostArr :: Arr k v -> Maybe k
  leftmostArr (Arr t) = go t where
    go :: T23 n k v -> Maybe k
    go Leaf = Nothing
    go (Node2 l (k, _) _) = go l <|> Just k
    go (Node3 l (k, _) _ _ _) = go l <|> Just k
  
  rightmostArr :: Arr k v -> Maybe k
  rightmostArr (Arr t) = go t where
    go :: T23 n k v -> Maybe k
    go Leaf = Nothing
    go (Node2 _ (k, _) r) = go r <|> Just k
    go (Node3 _ _ _ (k, _) r) = go r <|> Just k
  
  intersectSet :: Ord x => Set x -> Set x -> Set x
  intersectSet xs =
    foldMapArr (\ (y, ()) -> if inSet y xs then singleton y else mempty)
  
  subSet :: Ord x => Set x -> Set x -> Bool
  subSet xs ys = getAll (foldMapSet (All . (`inSet` ys)) xs)
  