{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_HADDOCK not-home #-}

-- |
-- This is an internal module. You probably don't need to import this. Use
-- "Data.Seqn.Seq" instead.
--
-- = WARNING
--
-- Definitions in this module allow violating invariants that would otherwise be
-- guaranteed by "Data.Seqn.Seq". Use at your own risk!
--
module Data.Seqn.Internal.Tree
  (
    -- * Tree
    Tree(..)

    -- * Basic
  , size
  , bin

    -- * Folds
  , foldl'
  , ifoldl'
  , foldr'
  , ifoldr'
  , traverse
  , itraverse

    -- * Construct
  , generateA

    -- * Index
  , adjustF
  , insertAt
  , deleteAt

    -- * Slice
  , cons
  , snoc
  , uncons
  , unsnoc
  , splitAtF

    -- * Transform
  , mapMaybeA
  , mapEitherA

    -- * Zip and unzip
  , zipWithStreamM
  , unzipWithA
  , unzipWith3A

    -- * Tree helpers
  , fold
  , foldSimple
  , link
  , glue
  , merge
  , balanceL
  , balanceR

    -- * Testing
  , valid
  , debugShowsPrec
  ) where

import Prelude hiding (concatMap, break, drop, dropWhile, filter, foldl', lookup, map, replicate, reverse, scanl, scanr, span, splitAt, take, takeWhile, traverse, unzip, unzip3, zip, zip3, zipWith, zipWith3)
import qualified Control.Applicative as Ap
import Control.DeepSeq (NFData(..), NFData1(..))
import Data.Bifunctor (Bifunctor(..))

import qualified Data.Seqn.Internal.Util as U
import Data.Seqn.Internal.Stream (Stream(..), Step(..))

data Tree a
  = Bin {-# UNPACK #-} !Int !a !(Tree a) !(Tree a)
  | Tip

--------------
-- Instances
--------------

instance NFData a => NFData (Tree a) where
  rnf :: Tree a -> ()
rnf = \case
    Bin Int
_ a
x Tree a
l Tree a
r -> a -> ()
forall a. NFData a => a -> ()
rnf a
x () -> () -> ()
forall a b. a -> b -> b
`seq` Tree a -> ()
forall a. NFData a => a -> ()
rnf Tree a
l () -> () -> ()
forall a b. a -> b -> b
`seq` Tree a -> ()
forall a. NFData a => a -> ()
rnf Tree a
r
    Tree a
Tip -> ()
  {-# INLINABLE rnf #-}

instance NFData1 Tree where
  liftRnf :: forall a. (a -> ()) -> Tree a -> ()
liftRnf a -> ()
f = Tree a -> ()
go
    where
      go :: Tree a -> ()
go (Bin Int
_ a
x Tree a
l Tree a
r) = a -> ()
f a
x () -> () -> ()
forall a b. a -> b -> b
`seq` Tree a -> ()
go Tree a
l () -> () -> ()
forall a b. a -> b -> b
`seq` Tree a -> ()
go Tree a
r
      go Tree a
Tip = ()
  {-# INLINE liftRnf #-}

-------------------
-- Basic Tree ops
-------------------

singleton :: a -> Tree a
singleton :: forall a. a -> Tree a
singleton a
x = Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip
{-# INLINE singleton #-}

size :: Tree a -> Int
size :: forall a. Tree a -> Int
size (Bin Int
n a
_ Tree a
_ Tree a
_) = Int
n
size Tree a
Tip = Int
0
{-# INLINE size #-}

-- O(1). Link two trees with a value in between. Precondition: The trees are
-- balanced wrt each other.
bin :: a -> Tree a -> Tree a -> Tree a
bin :: forall a. a -> Tree a -> Tree a -> Tree a
bin a
x Tree a
l Tree a
r = Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Tree a -> Int
forall a. Tree a -> Int
size Tree a
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Tree a -> Int
forall a. Tree a -> Int
size Tree a
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x Tree a
l Tree a
r
{-# INLINE bin #-}

----------
-- Folds
----------

-- Note [Folds]
-- ~~~~~~~~~~~~
-- Certain functions, such as folds, are implemented recursively on non-empty
-- trees, i.e. `go :: Int -> a -> Tree a -> Tree a -> b` instead of
-- `go :: Tree a -> b`. This is simply because benchmarks show this to be
-- faster.

foldl' :: (b -> a -> b) -> b -> Tree a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Tree a -> b
foldl' b -> a -> b
f !b
z0 = \case
  Bin Int
_ a
x Tree a
l Tree a
r -> b -> a -> Tree a -> Tree a -> b
go b
z0 a
x Tree a
l Tree a
r
  Tree a
Tip -> b
z0
  where
    go :: b -> a -> Tree a -> Tree a -> b
go !b
z !a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
_ a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr ->
          let !z' :: b
z' = b -> a -> Tree a -> Tree a -> b
go b
z a
lx Tree a
ll Tree a
lr
          in b -> a -> Tree a -> Tree a -> b
go (b -> a -> b
f b
z' a
x) a
rx Tree a
rl Tree a
rr
        Tree a
Tip ->
          let !z' :: b
z' = b -> a -> Tree a -> Tree a -> b
go b
z a
lx Tree a
ll Tree a
lr
          in b -> a -> b
f b
z' a
x
      Tree a
Tip -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr -> b -> a -> Tree a -> Tree a -> b
go (b -> a -> b
f b
z a
x) a
rx Tree a
rl Tree a
rr
        Tree a
Tip -> b -> a -> b
f b
z a
x
{-# INLINE foldl' #-}

ifoldl' :: (Int -> b -> a -> b) -> b -> Int -> Tree a -> b
ifoldl' :: forall b a. (Int -> b -> a -> b) -> b -> Int -> Tree a -> b
ifoldl' Int -> b -> a -> b
f !b
z0 !Int
i0 = \case
  Bin Int
_ a
x Tree a
l Tree a
r -> b -> Int -> a -> Tree a -> Tree a -> b
go b
z0 Int
i0 a
x Tree a
l Tree a
r
  Tree a
Tip -> b
z0
  where
    go :: b -> Int -> a -> Tree a -> Tree a -> b
go !b
z !Int
i !a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
lsz a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr ->
          let !z' :: b
z' = b -> Int -> a -> Tree a -> Tree a -> b
go b
z Int
i a
lx Tree a
ll Tree a
lr
          in b -> Int -> a -> Tree a -> Tree a -> b
go (Int -> b -> a -> b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsz) b
z' a
x) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lszInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) a
rx Tree a
rl Tree a
rr
        Tree a
Tip ->
          let !z' :: b
z' = b -> Int -> a -> Tree a -> Tree a -> b
go b
z Int
i a
lx Tree a
ll Tree a
lr
          in Int -> b -> a -> b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsz) b
z' a
x
      Tree a
Tip -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr -> b -> Int -> a -> Tree a -> Tree a -> b
go (Int -> b -> a -> b
f Int
i b
z a
x) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) a
rx Tree a
rl Tree a
rr
        Tree a
Tip -> Int -> b -> a -> b
f Int
i b
z a
x
{-# INLINE ifoldl' #-}

foldr' :: (a -> b -> b) -> b -> Tree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Tree a -> b
foldr' a -> b -> b
f !b
z0 = \case
  Bin Int
_ a
x Tree a
l Tree a
r -> b -> a -> Tree a -> Tree a -> b
go b
z0 a
x Tree a
l Tree a
r
  Tree a
Tip -> b
z0
  where
    go :: b -> a -> Tree a -> Tree a -> b
go !b
z !a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
_ a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr ->
          let !z' :: b
z' = b -> a -> Tree a -> Tree a -> b
go b
z a
rx Tree a
rl Tree a
rr
          in b -> a -> Tree a -> Tree a -> b
go (a -> b -> b
f a
x b
z') a
lx Tree a
ll Tree a
lr
        Tree a
Tip -> b -> a -> Tree a -> Tree a -> b
go (a -> b -> b
f a
x b
z) a
lx Tree a
ll Tree a
lr
      Tree a
Tip -> case Tree a
r of
        Bin Int
_ a
rx Tree a
rl Tree a
rr -> a -> b -> b
f a
x (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$! b -> a -> Tree a -> Tree a -> b
go b
z a
rx Tree a
rl Tree a
rr
        Tree a
Tip -> a -> b -> b
f a
x b
z
{-# INLINE foldr' #-}

ifoldr' :: (Int -> a -> b -> b) -> b -> Int -> Tree a -> b
ifoldr' :: forall a b. (Int -> a -> b -> b) -> b -> Int -> Tree a -> b
ifoldr' Int -> a -> b -> b
f !b
z0 !Int
i0 = \case
  Bin Int
_ a
x Tree a
l Tree a
r -> b -> Int -> a -> Tree a -> Tree a -> b
go b
z0 Int
i0 a
x Tree a
l Tree a
r
  Tree a
Tip -> b
z0
  where
    go :: b -> Int -> a -> Tree a -> Tree a -> b
go !b
z !Int
i !a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
_ a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr ->
          let !z' :: b
z' = b -> Int -> a -> Tree a -> Tree a -> b
go b
z Int
i a
rx Tree a
rl Tree a
rr
          in b -> Int -> a -> Tree a -> Tree a -> b
go (Int -> a -> b -> b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
rsz) a
x b
z') (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
rszInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
lx Tree a
ll Tree a
lr
        Tree a
Tip -> b -> Int -> a -> Tree a -> Tree a -> b
go (Int -> a -> b -> b
f Int
i a
x b
z) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
lx Tree a
ll Tree a
lr
      Tree a
Tip -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr -> Int -> a -> b -> b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
rsz) a
x (b -> b) -> b -> b
forall a b. (a -> b) -> a -> b
$! b -> Int -> a -> Tree a -> Tree a -> b
go b
z Int
i a
rx Tree a
rl Tree a
rr
        Tree a
Tip -> Int -> a -> b -> b
f Int
i a
x b
z
{-# INLINE ifoldr' #-}

fold
  :: b
  -> (Int -> a -> b -> b -> b)
  -> (Int -> a -> b -> b)
  -> (Int -> a -> b -> b)
  -> (a -> b)
  -> Tree a
  -> b
fold :: forall b a.
b
-> (Int -> a -> b -> b -> b)
-> (Int -> a -> b -> b)
-> (Int -> a -> b -> b)
-> (a -> b)
-> Tree a
-> b
fold b
tip Int -> a -> b -> b -> b
glr Int -> a -> b -> b
gl Int -> a -> b -> b
gr a -> b
g = \case
  Bin Int
sz a
x Tree a
l Tree a
r -> Int -> a -> Tree a -> Tree a -> b
go Int
sz a
x Tree a
l Tree a
r
  Tree a
Tip -> b
tip
  where
    go :: Int -> a -> Tree a -> Tree a -> b
go !Int
sz !a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
lsz a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr -> Int -> a -> b -> b -> b
glr Int
sz a
x (Int -> a -> Tree a -> Tree a -> b
go Int
lsz a
lx Tree a
ll Tree a
lr) (Int -> a -> Tree a -> Tree a -> b
go Int
rsz a
rx Tree a
rl Tree a
rr)
        Tree a
Tip -> Int -> a -> b -> b
gl Int
sz a
x (Int -> a -> Tree a -> Tree a -> b
go Int
lsz a
lx Tree a
ll Tree a
lr)
      Tree a
Tip -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr -> Int -> a -> b -> b
gr Int
sz a
x (Int -> a -> Tree a -> Tree a -> b
go Int
rsz a
rx Tree a
rl Tree a
rr)
        Tree a
Tip -> a -> b
g a
x
{-# INLINE fold #-}

foldSimple :: b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple :: forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple b
tip Int -> a -> b -> b -> b
f = b
-> (Int -> a -> b -> b -> b)
-> (Int -> a -> b -> b)
-> (Int -> a -> b -> b)
-> (a -> b)
-> Tree a
-> b
forall b a.
b
-> (Int -> a -> b -> b -> b)
-> (Int -> a -> b -> b)
-> (Int -> a -> b -> b)
-> (a -> b)
-> Tree a
-> b
fold b
tip Int -> a -> b -> b -> b
f Int -> a -> b -> b
gl Int -> a -> b -> b
gr a -> b
g
  where
    gl :: Int -> a -> b -> b
gl !Int
sz a
x b
ml = Int -> a -> b -> b -> b
f Int
sz a
x b
ml b
tip
    {-# INLINE gl #-}
    gr :: Int -> a -> b -> b
gr !Int
sz a
x b
mr = Int -> a -> b -> b -> b
f Int
sz a
x b
tip b
mr
    {-# INLINE gr #-}
    g :: a -> b
g a
x = Int -> a -> b -> b -> b
f Int
1 a
x b
tip b
tip
    {-# INLINE g #-}
{-# INLINE foldSimple #-}

traverse :: Applicative f => (a -> f b) -> Tree a -> f (Tree b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Tree a -> f (Tree b)
traverse a -> f b
f = f (Tree b)
-> (Int -> a -> f (Tree b) -> f (Tree b) -> f (Tree b))
-> (Int -> a -> f (Tree b) -> f (Tree b))
-> (Int -> a -> f (Tree b) -> f (Tree b))
-> (a -> f (Tree b))
-> Tree a
-> f (Tree b)
forall b a.
b
-> (Int -> a -> b -> b -> b)
-> (Int -> a -> b -> b)
-> (Int -> a -> b -> b)
-> (a -> b)
-> Tree a
-> b
fold (Tree b -> f (Tree b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree b
forall a. Tree a
Tip) Int -> a -> f (Tree b) -> f (Tree b) -> f (Tree b)
glr Int -> a -> f (Tree b) -> f (Tree b)
gl Int -> a -> f (Tree b) -> f (Tree b)
gr a -> f (Tree b)
g
  where
    glr :: Int -> a -> f (Tree b) -> f (Tree b) -> f (Tree b)
glr !Int
sz a
x f (Tree b)
ml f (Tree b)
mr = (Tree b -> b -> Tree b -> Tree b)
-> f (Tree b) -> f b -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3R' ((b -> Tree b -> Tree b -> Tree b)
-> Tree b -> b -> Tree b -> Tree b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz)) f (Tree b)
ml (a -> f b
f a
x) f (Tree b)
mr
    -- See Note [Traverse liftA3R']
    {-# INLINE glr #-}
    gl :: Int -> a -> f (Tree b) -> f (Tree b)
gl !Int
sz a
x f (Tree b)
ml = (Tree b -> b -> Tree b) -> f (Tree b) -> f b -> f (Tree b)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2 (\Tree b
l' b
x' -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x' Tree b
l' Tree b
forall a. Tree a
Tip) f (Tree b)
ml (a -> f b
f a
x)
    {-# INLINE gl #-}
    gr :: Int -> a -> f (Tree b) -> f (Tree b)
gr !Int
sz a
x f (Tree b)
mr = (b -> Tree b -> Tree b) -> f b -> f (Tree b) -> f (Tree b)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2 (\b
x' Tree b
r' -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x' Tree b
forall a. Tree a
Tip Tree b
r') (a -> f b
f a
x) f (Tree b)
mr
    {-# INLINE gr #-}
    g :: a -> f (Tree b)
g a
x = (b -> Tree b) -> f b -> f (Tree b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> Tree b
forall a. a -> Tree a
singleton (a -> f b
f a
x)
    {-# INLINE g #-}
{-# INLINE traverse #-}

itraverse :: Applicative f => (Int -> a -> f b) -> Int -> Tree a -> f (Tree b)
itraverse :: forall (f :: * -> *) a b.
Applicative f =>
(Int -> a -> f b) -> Int -> Tree a -> f (Tree b)
itraverse Int -> a -> f b
f !Int
i0 = \case
  Bin Int
sz a
x Tree a
l Tree a
r -> Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go Int
i0 Int
sz a
x Tree a
l Tree a
r
  Tree a
Tip -> Tree b -> f (Tree b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree b
forall a. Tree a
Tip
  where
    go :: Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go !Int
i !Int
sz a
x Tree a
l Tree a
r = case Tree a
l of
      Bin Int
lsz a
lx Tree a
ll Tree a
lr -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr ->
          (Tree b -> b -> Tree b -> Tree b)
-> f (Tree b) -> f b -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3R'
            ((b -> Tree b -> Tree b -> Tree b)
-> Tree b -> b -> Tree b -> Tree b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz))
            (Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go Int
i Int
lsz a
lx Tree a
ll Tree a
lr)
            (Int -> a -> f b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsz) a
x)
            (Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lszInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
rsz a
rx Tree a
rl Tree a
rr)
          -- See Note [Traverse liftA3R']
        Tree a
Tip ->
          (Tree b -> b -> Tree b) -> f (Tree b) -> f b -> f (Tree b)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2
            (\Tree b
l' b
x' -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x' Tree b
l' Tree b
forall a. Tree a
Tip)
            (Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go Int
i Int
lsz a
lx Tree a
ll Tree a
lr)
            (Int -> a -> f b
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsz) a
x)
      Tree a
Tip -> case Tree a
r of
        Bin Int
rsz a
rx Tree a
rl Tree a
rr ->
          (b -> Tree b -> Tree b) -> f b -> f (Tree b) -> f (Tree b)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2 (\b
x' Tree b
r' -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x' Tree b
forall a. Tree a
Tip Tree b
r') (Int -> a -> f b
f Int
i a
x) (Int -> Int -> a -> Tree a -> Tree a -> f (Tree b)
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
rsz a
rx Tree a
rl Tree a
rr)
        Tree a
Tip ->
          (b -> Tree b) -> f b -> f (Tree b)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\b
x' -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x' Tree b
forall a. Tree a
Tip Tree b
forall a. Tree a
Tip) (Int -> a -> f b
f Int
i a
x)
    -- See Note [Traverse]
{-# INLINE itraverse #-}

-- Note [Traverse liftA3R']
-- ~~~~~~~~~~~~~~~~~~~~~~~~
-- We want to associate to the right because we define foldMap using traverse
-- and ifoldMap using itraverse. It is more appropriate to be right-associative
-- for <>.

-- Right associative and strict
liftA3R' :: Applicative f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3R' :: forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3R' a -> b -> c -> d
f f a
mx f b
my f c
mz =
  (a -> S2 b c -> d) -> f a -> f (S2 b c) -> f d
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2
    (\a
x (U.S2 b
y c
z) -> a -> b -> c -> d
f a
x b
y c
z)
    f a
mx
    ((b -> c -> S2 b c) -> f b -> f c -> f (S2 b c)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
Ap.liftA2 b -> c -> S2 b c
forall a b. a -> b -> S2 a b
U.S2 f b
my f c
mz)
{-# INLINE liftA3R' #-}

--------------
-- Construct
--------------

generateA :: Applicative f => (Int -> f a) -> Int -> Int -> f (Tree a)
generateA :: forall (f :: * -> *) a.
Applicative f =>
(Int -> f a) -> Int -> Int -> f (Tree a)
generateA Int -> f a
f = Int -> Int -> f (Tree a)
go
  where
    go :: Int -> Int -> f (Tree a)
go !Int
i Int
n
      | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Tree a -> f (Tree a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
forall a. Tree a
Tip
      | Bool
otherwise =
          (Tree a -> a -> Tree a -> Tree a)
-> f (Tree a) -> f a -> f (Tree a) -> f (Tree a)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
Ap.liftA3
            ((a -> Tree a -> Tree a -> Tree a)
-> Tree a -> a -> Tree a -> Tree a
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
n))
            (Int -> Int -> f (Tree a)
go Int
i Int
lsz)
            (Int -> f a
f (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsz))
            (Int -> Int -> f (Tree a)
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lszInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lszInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
      where
        lsz :: Int
lsz = (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
{-# INLINE generateA #-}

----------
-- Index
----------

-- Precondition: 0 <= i < size xs
adjustF :: Functor f => (a -> f a) -> Int -> Tree a -> f (Tree a)
adjustF :: forall (f :: * -> *) a.
Functor f =>
(a -> f a) -> Int -> Tree a -> f (Tree a)
adjustF a -> f a
f = Int -> Tree a -> f (Tree a)
go
  where
    go :: Int -> Tree a -> f (Tree a)
go !Int
i = \case
      Bin Int
sz a
x Tree a
l Tree a
r -> case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
szl of
        Ordering
LT -> (Tree a -> Tree a) -> f (Tree a) -> f (Tree a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Tree a
l' -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz a
x Tree a
l' Tree a
r) (Int -> Tree a -> f (Tree a)
go Int
i Tree a
l)
        Ordering
EQ -> (a -> Tree a) -> f a -> f (Tree a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
x' -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz a
x' Tree a
l Tree a
r) (a -> f a
f a
x)
        Ordering
GT -> (Tree a -> Tree a) -> f (Tree a) -> f (Tree a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz a
x Tree a
l) (Int -> Tree a -> f (Tree a)
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szlInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Tree a
r)
        where
          szl :: Int
szl = Tree a -> Int
forall a. Tree a -> Int
size Tree a
l
      Tree a
Tip -> String -> f (Tree a)
forall a. String -> a
errorOutOfBounds String
"Tree.adjustF"
{-# INLINE adjustF #-}

-- Inserts at ends if not in bounds
insertAt :: Int -> a -> Tree a -> Tree a
insertAt :: forall a. Int -> a -> Tree a -> Tree a
insertAt !Int
i a
x (Bin Int
_ a
y Tree a
l Tree a
r)
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
szl = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
y (Int -> a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a
insertAt Int
i a
x Tree a
l) Tree a
r
  | Bool
otherwise = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
y Tree a
l (Int -> a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a
insertAt (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szlInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a
x Tree a
r)
  where
    szl :: Int
szl = Tree a -> Int
forall a. Tree a -> Int
size Tree a
l
insertAt Int
_ a
x Tree a
Tip = a -> Tree a
forall a. a -> Tree a
singleton a
x

-- Precondition: 0 <= i < size xs
deleteAt :: Int -> Tree a -> Tree a
deleteAt :: forall a. Int -> Tree a -> Tree a
deleteAt !Int
i (Bin Int
_ a
x Tree a
l Tree a
r) = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
szl of
  Ordering
LT -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
x (Int -> Tree a -> Tree a
forall a. Int -> Tree a -> Tree a
deleteAt Int
i Tree a
l) Tree a
r
  Ordering
EQ -> Tree a -> Tree a -> Tree a
forall a. Tree a -> Tree a -> Tree a
glue Tree a
l Tree a
r
  Ordering
GT -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
x Tree a
l (Int -> Tree a -> Tree a
forall a. Int -> Tree a -> Tree a
deleteAt (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szlInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Tree a
r)
  where
    szl :: Int
szl = Tree a -> Int
forall a. Tree a -> Int
size Tree a
l
deleteAt Int
_ Tree a
Tip = String -> Tree a
forall a. String -> a
errorOutOfBounds String
"Tree.deleteAt"

----------
-- Slice
----------

cons :: a -> Tree a -> Tree a
cons :: forall a. a -> Tree a -> Tree a
cons a
x Tree a
Tip = a -> Tree a
forall a. a -> Tree a
singleton a
x
cons a
x (Bin Int
_ a
y Tree a
l Tree a
r) = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
y (a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a
cons a
x Tree a
l) Tree a
r

snoc :: Tree a -> a -> Tree a
snoc :: forall a. Tree a -> a -> Tree a
snoc Tree a
Tip a
x = a -> Tree a
forall a. a -> Tree a
singleton a
x
snoc (Bin Int
_ a
y Tree a
l Tree a
r) a
x = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
y Tree a
l (Tree a -> a -> Tree a
forall a. Tree a -> a -> Tree a
snoc Tree a
r a
x)

uncons :: Tree a -> U.SMaybe (U.S2 a (Tree a))
uncons :: forall a. Tree a -> SMaybe (S2 a (Tree a))
uncons (Bin Int
_ a
x Tree a
l Tree a
r) = S2 a (Tree a) -> SMaybe (S2 a (Tree a))
forall a. a -> SMaybe a
U.SJust (a -> Tree a -> Tree a -> S2 a (Tree a)
forall a. a -> Tree a -> Tree a -> S2 a (Tree a)
unconsSure a
x Tree a
l Tree a
r)
uncons Tree a
Tip = SMaybe (S2 a (Tree a))
forall a. SMaybe a
U.SNothing
{-# INLINE uncons #-}

unconsSure :: a -> Tree a -> Tree a -> U.S2 a (Tree a)
unconsSure :: forall a. a -> Tree a -> Tree a -> S2 a (Tree a)
unconsSure a
x (Bin Int
_ a
lx Tree a
ll Tree a
lr) Tree a
r = case a -> Tree a -> Tree a -> S2 a (Tree a)
forall a. a -> Tree a -> Tree a -> S2 a (Tree a)
unconsSure a
lx Tree a
ll Tree a
lr of
  U.S2 a
y Tree a
l' -> a -> Tree a -> S2 a (Tree a)
forall a b. a -> b -> S2 a b
U.S2 a
y (a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
x Tree a
l' Tree a
r)
unconsSure a
x Tree a
Tip Tree a
r = a -> Tree a -> S2 a (Tree a)
forall a b. a -> b -> S2 a b
U.S2 a
x Tree a
r

unsnoc :: Tree a -> U.SMaybe (U.S2 (Tree a) a)
unsnoc :: forall a. Tree a -> SMaybe (S2 (Tree a) a)
unsnoc (Bin Int
_ a
x Tree a
l Tree a
r) = S2 (Tree a) a -> SMaybe (S2 (Tree a) a)
forall a. a -> SMaybe a
U.SJust (S2 (Tree a) a -> SMaybe (S2 (Tree a) a))
-> S2 (Tree a) a -> SMaybe (S2 (Tree a) a)
forall a b. (a -> b) -> a -> b
$ a -> Tree a -> Tree a -> S2 (Tree a) a
forall a. a -> Tree a -> Tree a -> S2 (Tree a) a
unsnocSure a
x Tree a
l Tree a
r
unsnoc Tree a
Tip = SMaybe (S2 (Tree a) a)
forall a. SMaybe a
U.SNothing
{-# INLINE unsnoc #-}

unsnocSure :: a -> Tree a -> Tree a -> U.S2 (Tree a) a
unsnocSure :: forall a. a -> Tree a -> Tree a -> S2 (Tree a) a
unsnocSure a
x Tree a
l (Bin Int
_ a
rx Tree a
rl Tree a
rr) = case a -> Tree a -> Tree a -> S2 (Tree a) a
forall a. a -> Tree a -> Tree a -> S2 (Tree a) a
unsnocSure a
rx Tree a
rl Tree a
rr of
  U.S2 Tree a
r' a
y -> Tree a -> a -> S2 (Tree a) a
forall a b. a -> b -> S2 a b
U.S2 (a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
x Tree a
l Tree a
r') a
y
unsnocSure a
x Tree a
l Tree a
Tip = Tree a -> a -> S2 (Tree a) a
forall a b. a -> b -> S2 a b
U.S2 Tree a
l a
x

-- Precondition: 0 <= i < size xs
splitAtF
  :: U.Biapplicative f
  => Int -> Tree a -> f (Tree a) (U.S2 a (Tree a))
splitAtF :: forall (f :: * -> * -> *) a.
Biapplicative f =>
Int -> Tree a -> f (Tree a) (S2 a (Tree a))
splitAtF = Int -> Tree a -> f (Tree a) (S2 a (Tree a))
forall (f :: * -> * -> *) a.
Biapplicative f =>
Int -> Tree a -> f (Tree a) (S2 a (Tree a))
go
  where
    go :: Int -> Tree a -> p (Tree a) (S2 a (Tree a))
go !Int
i (Bin Int
_ a
x Tree a
l Tree a
r) = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
i Int
szl of
      Ordering
LT -> (S2 a (Tree a) -> S2 a (Tree a))
-> p (Tree a) (S2 a (Tree a)) -> p (Tree a) (S2 a (Tree a))
forall b c a. (b -> c) -> p a b -> p a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ((Tree a -> Tree a) -> S2 a (Tree a) -> S2 a (Tree a)
forall b c a. (b -> c) -> S2 a b -> S2 a c
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (\Tree a
lr -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
link a
x Tree a
lr Tree a
r)) (Int -> Tree a -> p (Tree a) (S2 a (Tree a))
go Int
i Tree a
l)
      Ordering
EQ -> Tree a -> S2 a (Tree a) -> p (Tree a) (S2 a (Tree a))
forall a b. a -> b -> p a b
forall (p :: * -> * -> *) a b. Biapplicative p => a -> b -> p a b
U.bipure Tree a
l (a -> Tree a -> S2 a (Tree a)
forall a b. a -> b -> S2 a b
U.S2 a
x Tree a
r)
      Ordering
GT -> (Tree a -> Tree a)
-> p (Tree a) (S2 a (Tree a)) -> p (Tree a) (S2 a (Tree a))
forall a b c. (a -> b) -> p a c -> p b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
link a
x Tree a
l) (Int -> Tree a -> p (Tree a) (S2 a (Tree a))
go (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
szlInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Tree a
r)
      where
        szl :: Int
szl = Tree a -> Int
forall a. Tree a -> Int
size Tree a
l
    go Int
_ Tree a
Tip = String -> p (Tree a) (S2 a (Tree a))
forall a. String -> a
errorOutOfBounds String
"Tree.splitAtF"
{-# INLINE splitAtF #-}

--------------
-- Transform
--------------

mapMaybeA :: Applicative f => (a -> f (Maybe b)) -> Tree a -> f (Tree b)
mapMaybeA :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f (Maybe b)) -> Tree a -> f (Tree b)
mapMaybeA a -> f (Maybe b)
f = f (Tree b)
-> (Int -> a -> f (Tree b) -> f (Tree b) -> f (Tree b))
-> Tree a
-> f (Tree b)
forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple f (Tree b)
forall {a}. f (Tree a)
tip Int -> a -> f (Tree b) -> f (Tree b) -> f (Tree b)
forall {p}. p -> a -> f (Tree b) -> f (Tree b) -> f (Tree b)
g
  where
    tip :: f (Tree a)
tip = Tree a -> f (Tree a)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
forall a. Tree a
Tip
    {-# INLINE tip #-}
    g :: p -> a -> f (Tree b) -> f (Tree b) -> f (Tree b)
g p
_ a
x f (Tree b)
ml f (Tree b)
mr = (\Tree b -> Maybe b -> Tree b -> Tree b
h -> (Tree b -> Maybe b -> Tree b -> Tree b)
-> f (Tree b) -> f (Maybe b) -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
Ap.liftA3 Tree b -> Maybe b -> Tree b -> Tree b
h f (Tree b)
ml (a -> f (Maybe b)
f a
x) f (Tree b)
mr) ((Tree b -> Maybe b -> Tree b -> Tree b) -> f (Tree b))
-> (Tree b -> Maybe b -> Tree b -> Tree b) -> f (Tree b)
forall a b. (a -> b) -> a -> b
$ \Tree b
l Maybe b
my Tree b
r ->
      case Maybe b
my of
        Maybe b
Nothing -> Tree b -> Tree b -> Tree b
forall a. Tree a -> Tree a -> Tree a
merge Tree b
l Tree b
r
        Just b
y -> b -> Tree b -> Tree b -> Tree b
forall a. a -> Tree a -> Tree a -> Tree a
link b
y Tree b
l Tree b
r
    {-# INLINE g #-}
{-# INLINE mapMaybeA #-}

mapEitherA
  :: Applicative f
  => (a -> f (Either b c)) -> Tree a -> f (U.S2 (Tree b) (Tree c))
mapEitherA :: forall (f :: * -> *) a b c.
Applicative f =>
(a -> f (Either b c)) -> Tree a -> f (S2 (Tree b) (Tree c))
mapEitherA a -> f (Either b c)
f = f (S2 (Tree b) (Tree c))
-> (Int
    -> a
    -> f (S2 (Tree b) (Tree c))
    -> f (S2 (Tree b) (Tree c))
    -> f (S2 (Tree b) (Tree c)))
-> Tree a
-> f (S2 (Tree b) (Tree c))
forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple f (S2 (Tree b) (Tree c))
forall {a} {a}. f (S2 (Tree a) (Tree a))
tip Int
-> a
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
forall {p :: * -> * -> *} {p}.
Biapplicative p =>
p
-> a
-> f (p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
g
  where
    tip :: f (S2 (Tree a) (Tree a))
tip = S2 (Tree a) (Tree a) -> f (S2 (Tree a) (Tree a))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> Tree a -> S2 (Tree a) (Tree a)
forall a b. a -> b -> S2 a b
forall (p :: * -> * -> *) a b. Biapplicative p => a -> b -> p a b
U.bipure Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
    {-# INLINE tip #-}
    g :: p
-> a
-> f (p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
g p
_ a
x f (p (Tree b) (Tree c))
ml f (p (Tree b) (Tree c))
mr = (\p (Tree b) (Tree c)
-> Either b c -> p (Tree b) (Tree c) -> p (Tree b) (Tree c)
h -> (p (Tree b) (Tree c)
 -> Either b c -> p (Tree b) (Tree c) -> p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
-> f (Either b c)
-> f (p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
Ap.liftA3 p (Tree b) (Tree c)
-> Either b c -> p (Tree b) (Tree c) -> p (Tree b) (Tree c)
h f (p (Tree b) (Tree c))
ml (a -> f (Either b c)
f a
x) f (p (Tree b) (Tree c))
mr) ((p (Tree b) (Tree c)
  -> Either b c -> p (Tree b) (Tree c) -> p (Tree b) (Tree c))
 -> f (p (Tree b) (Tree c)))
-> (p (Tree b) (Tree c)
    -> Either b c -> p (Tree b) (Tree c) -> p (Tree b) (Tree c))
-> f (p (Tree b) (Tree c))
forall a b. (a -> b) -> a -> b
$ \p (Tree b) (Tree c)
l Either b c
my p (Tree b) (Tree c)
r ->
      case Either b c
my of
        Left b
y -> (Tree b -> Tree b -> Tree b)
-> (Tree c -> Tree c -> Tree c)
-> p (Tree b) (Tree c)
-> p (Tree b) (Tree c)
-> p (Tree b) (Tree c)
forall a b c d e f.
(a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f
forall (p :: * -> * -> *) a b c d e f.
Biapplicative p =>
(a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f
U.biliftA2 (b -> Tree b -> Tree b -> Tree b
forall a. a -> Tree a -> Tree a -> Tree a
link b
y) Tree c -> Tree c -> Tree c
forall a. Tree a -> Tree a -> Tree a
merge p (Tree b) (Tree c)
l p (Tree b) (Tree c)
r
        Right c
y -> (Tree b -> Tree b -> Tree b)
-> (Tree c -> Tree c -> Tree c)
-> p (Tree b) (Tree c)
-> p (Tree b) (Tree c)
-> p (Tree b) (Tree c)
forall a b c d e f.
(a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f
forall (p :: * -> * -> *) a b c d e f.
Biapplicative p =>
(a -> b -> c) -> (d -> e -> f) -> p a d -> p b e -> p c f
U.biliftA2 Tree b -> Tree b -> Tree b
forall a. Tree a -> Tree a -> Tree a
merge (c -> Tree c -> Tree c -> Tree c
forall a. a -> Tree a -> Tree a -> Tree a
link c
y) p (Tree b) (Tree c)
l p (Tree b) (Tree c)
r
    {-# INLINE g #-}
{-# INLINE mapEitherA #-}

------------------
-- Zip and unzip
------------------

zipWithStreamM :: Monad m => (a -> b -> m c) -> Tree a -> Stream b -> m (Tree c)
zipWithStreamM :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Tree a -> Stream b -> m (Tree c)
zipWithStreamM a -> b -> m c
f Tree a
t (Stream s -> Step s b
step s
s) = SStateT s m (Tree c) -> s -> m (Tree c)
forall (m :: * -> *) s a. Functor m => SStateT s m a -> s -> m a
U.evalSStateT (SStateT s m (Tree c)
-> (Int
    -> a
    -> SStateT s m (Tree c)
    -> SStateT s m (Tree c)
    -> SStateT s m (Tree c))
-> Tree a
-> SStateT s m (Tree c)
forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple SStateT s m (Tree c)
forall {a}. SStateT s m (Tree a)
tip Int
-> a
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
forall {p}.
p
-> a
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
g Tree a
t) s
s
  where
    tip :: SStateT s m (Tree a)
tip = Tree a -> SStateT s m (Tree a)
forall a. a -> SStateT s m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Tree a
forall a. Tree a
Tip
    {-# INLINE tip #-}
    g :: p
-> a
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
-> SStateT s m (Tree c)
g p
_ a
x SStateT s m (Tree c)
ml SStateT s m (Tree c)
mr = (s -> m (S2 s (Tree c))) -> SStateT s m (Tree c)
forall s (m :: * -> *) a. (s -> m (S2 s a)) -> SStateT s m a
U.SStateT ((s -> m (S2 s (Tree c))) -> SStateT s m (Tree c))
-> (s -> m (S2 s (Tree c))) -> SStateT s m (Tree c)
forall a b. (a -> b) -> a -> b
$ \s
s2 -> do
      U.S2 s
s3 Tree c
l <- SStateT s m (Tree c) -> s -> m (S2 s (Tree c))
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
U.runSStateT SStateT s m (Tree c)
ml s
s2
      case s -> Step s b
step s
s3 of
        Step s b
Done -> S2 s (Tree c) -> m (S2 s (Tree c))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (S2 s (Tree c) -> m (S2 s (Tree c)))
-> S2 s (Tree c) -> m (S2 s (Tree c))
forall a b. (a -> b) -> a -> b
$ s -> Tree c -> S2 s (Tree c)
forall a b. a -> b -> S2 a b
U.S2 s
s3 Tree c
l
        Yield b
y s
s4 -> do
          c
z <- a -> b -> m c
f a
x b
y
          U.S2 s
s5 Tree c
r <- SStateT s m (Tree c) -> s -> m (S2 s (Tree c))
forall s (m :: * -> *) a. SStateT s m a -> s -> m (S2 s a)
U.runSStateT SStateT s m (Tree c)
mr s
s4
          S2 s (Tree c) -> m (S2 s (Tree c))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (S2 s (Tree c) -> m (S2 s (Tree c)))
-> S2 s (Tree c) -> m (S2 s (Tree c))
forall a b. (a -> b) -> a -> b
$! s -> Tree c -> S2 s (Tree c)
forall a b. a -> b -> S2 a b
U.S2 s
s5 (c -> Tree c -> Tree c -> Tree c
forall a. a -> Tree a -> Tree a -> Tree a
link c
z Tree c
l Tree c
r)
    {-# INLINE g #-}
{-# INLINE zipWithStreamM #-}

unzipWithA
  :: Applicative f => (a -> f (b, c)) -> Tree a -> f (U.S2 (Tree b) (Tree c))
unzipWithA :: forall (f :: * -> *) a b c.
Applicative f =>
(a -> f (b, c)) -> Tree a -> f (S2 (Tree b) (Tree c))
unzipWithA a -> f (b, c)
f = f (S2 (Tree b) (Tree c))
-> (Int
    -> a
    -> f (S2 (Tree b) (Tree c))
    -> f (S2 (Tree b) (Tree c))
    -> f (S2 (Tree b) (Tree c)))
-> Tree a
-> f (S2 (Tree b) (Tree c))
forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple f (S2 (Tree b) (Tree c))
forall {a} {a}. f (S2 (Tree a) (Tree a))
tip Int
-> a
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
g
  where
    tip :: f (S2 (Tree a) (Tree a))
tip = S2 (Tree a) (Tree a) -> f (S2 (Tree a) (Tree a))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> Tree a -> S2 (Tree a) (Tree a)
forall a b. a -> b -> S2 a b
U.S2 Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
    {-# INLINE tip #-}
    g :: Int
-> a
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
g !Int
sz a
x f (S2 (Tree b) (Tree c))
ml f (S2 (Tree b) (Tree c))
mr = (\S2 (Tree b) (Tree c)
-> (b, c) -> S2 (Tree b) (Tree c) -> S2 (Tree b) (Tree c)
h -> (S2 (Tree b) (Tree c)
 -> (b, c) -> S2 (Tree b) (Tree c) -> S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
-> f (b, c)
-> f (S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
Ap.liftA3 S2 (Tree b) (Tree c)
-> (b, c) -> S2 (Tree b) (Tree c) -> S2 (Tree b) (Tree c)
h f (S2 (Tree b) (Tree c))
ml (a -> f (b, c)
f a
x) f (S2 (Tree b) (Tree c))
mr) ((S2 (Tree b) (Tree c)
  -> (b, c) -> S2 (Tree b) (Tree c) -> S2 (Tree b) (Tree c))
 -> f (S2 (Tree b) (Tree c)))
-> (S2 (Tree b) (Tree c)
    -> (b, c) -> S2 (Tree b) (Tree c) -> S2 (Tree b) (Tree c))
-> f (S2 (Tree b) (Tree c))
forall a b. (a -> b) -> a -> b
$
      \(U.S2 Tree b
l1 Tree c
l2) (b
x1,c
x2) (U.S2 Tree b
r1 Tree c
r2) ->
        Tree b -> Tree c -> S2 (Tree b) (Tree c)
forall a b. a -> b -> S2 a b
U.S2 (Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x1 Tree b
l1 Tree b
r1) (Int -> c -> Tree c -> Tree c -> Tree c
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz c
x2 Tree c
l2 Tree c
r2)
    {-# INLINE g #-}
{-# INLINE unzipWithA #-}

unzipWith3A
  :: Applicative f
  => (a -> f (b, c, d))
  -> Tree a
  -> f (U.S3 (Tree b) (Tree c) (Tree d))
unzipWith3A :: forall (f :: * -> *) a b c d.
Applicative f =>
(a -> f (b, c, d)) -> Tree a -> f (S3 (Tree b) (Tree c) (Tree d))
unzipWith3A a -> f (b, c, d)
f = f (S3 (Tree b) (Tree c) (Tree d))
-> (Int
    -> a
    -> f (S3 (Tree b) (Tree c) (Tree d))
    -> f (S3 (Tree b) (Tree c) (Tree d))
    -> f (S3 (Tree b) (Tree c) (Tree d)))
-> Tree a
-> f (S3 (Tree b) (Tree c) (Tree d))
forall b a. b -> (Int -> a -> b -> b -> b) -> Tree a -> b
foldSimple f (S3 (Tree b) (Tree c) (Tree d))
forall {a} {a} {a}. f (S3 (Tree a) (Tree a) (Tree a))
tip Int
-> a
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
g
  where
    tip :: f (S3 (Tree a) (Tree a) (Tree a))
tip = S3 (Tree a) (Tree a) (Tree a) -> f (S3 (Tree a) (Tree a) (Tree a))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> Tree a -> Tree a -> S3 (Tree a) (Tree a) (Tree a)
forall a b c. a -> b -> c -> S3 a b c
U.S3 Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
    {-# INLINE tip #-}
    g :: Int
-> a
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
g !Int
sz a
x f (S3 (Tree b) (Tree c) (Tree d))
ml f (S3 (Tree b) (Tree c) (Tree d))
mr = (\S3 (Tree b) (Tree c) (Tree d)
-> (b, c, d)
-> S3 (Tree b) (Tree c) (Tree d)
-> S3 (Tree b) (Tree c) (Tree d)
h -> (S3 (Tree b) (Tree c) (Tree d)
 -> (b, c, d)
 -> S3 (Tree b) (Tree c) (Tree d)
 -> S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (b, c, d)
-> f (S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
Ap.liftA3 S3 (Tree b) (Tree c) (Tree d)
-> (b, c, d)
-> S3 (Tree b) (Tree c) (Tree d)
-> S3 (Tree b) (Tree c) (Tree d)
h f (S3 (Tree b) (Tree c) (Tree d))
ml (a -> f (b, c, d)
f a
x) f (S3 (Tree b) (Tree c) (Tree d))
mr) ((S3 (Tree b) (Tree c) (Tree d)
  -> (b, c, d)
  -> S3 (Tree b) (Tree c) (Tree d)
  -> S3 (Tree b) (Tree c) (Tree d))
 -> f (S3 (Tree b) (Tree c) (Tree d)))
-> (S3 (Tree b) (Tree c) (Tree d)
    -> (b, c, d)
    -> S3 (Tree b) (Tree c) (Tree d)
    -> S3 (Tree b) (Tree c) (Tree d))
-> f (S3 (Tree b) (Tree c) (Tree d))
forall a b. (a -> b) -> a -> b
$
      \(U.S3 Tree b
l1 Tree c
l2 Tree d
l3) (b
x1,c
x2,d
x3) (U.S3 Tree b
r1 Tree c
r2 Tree d
r3) ->
        Tree b -> Tree c -> Tree d -> S3 (Tree b) (Tree c) (Tree d)
forall a b c. a -> b -> c -> S3 a b c
U.S3 (Int -> b -> Tree b -> Tree b -> Tree b
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz b
x1 Tree b
l1 Tree b
r1) (Int -> c -> Tree c -> Tree c -> Tree c
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz c
x2 Tree c
l2 Tree c
r2) (Int -> d -> Tree d -> Tree d -> Tree d
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
sz d
x3 Tree d
l3 Tree d
r3)
    {-# INLINE g #-}
{-# INLINE unzipWith3A #-}

-----------
-- Errors
-----------

errorOutOfBounds :: String -> a
errorOutOfBounds :: forall a. String -> a
errorOutOfBounds String
name = String -> a
forall a. HasCallStack => String -> a
error (String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": out of bounds")

------------
-- Balance
------------

-- O(|log n1 - log n2|). Link two trees with a value in between.
link :: a -> Tree a -> Tree a -> Tree a
link :: forall a. a -> Tree a -> Tree a -> Tree a
link !a
x Tree a
Tip Tree a
r = a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a
cons a
x Tree a
r
link a
x Tree a
l Tree a
Tip = Tree a -> a -> Tree a
forall a. Tree a -> a -> Tree a
snoc Tree a
l a
x
link a
x l :: Tree a
l@(Bin Int
ls a
lx Tree a
ll Tree a
lr) r :: Tree a
r@(Bin Int
rs a
rx Tree a
rl Tree a
rr)
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
ls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rs = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
rx (a -> Int -> Tree a -> Tree a -> Tree a
forall a. a -> Int -> Tree a -> Tree a -> Tree a
linkL a
x Int
ls Tree a
l Tree a
rl) Tree a
rr
  | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
rs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ls = a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
lx Tree a
ll (a -> Tree a -> Int -> Tree a -> Tree a
forall a. a -> Tree a -> Int -> Tree a -> Tree a
linkR a
x Tree a
lr Int
rs Tree a
r)
  | Bool
otherwise     = Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
l Tree a
r
{-# INLINE link #-}

linkL :: a -> Int -> Tree a -> Tree a -> Tree a
linkL :: forall a. a -> Int -> Tree a -> Tree a -> Tree a
linkL !a
x !Int
ls !Tree a
l Tree a
r = case Tree a
r of
  Bin Int
rs a
rx Tree a
rl Tree a
rr
    | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
ls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rs -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
rx (a -> Int -> Tree a -> Tree a -> Tree a
forall a. a -> Int -> Tree a -> Tree a -> Tree a
linkL a
x Int
ls Tree a
l Tree a
rl) Tree a
rr
    | Bool
otherwise     -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
l Tree a
r
  Tree a
Tip -> String -> Tree a
forall a. HasCallStack => String -> a
error String
"Tree.linkL: impossible"

linkR :: a -> Tree a -> Int -> Tree a -> Tree a
linkR :: forall a. a -> Tree a -> Int -> Tree a -> Tree a
linkR !a
x Tree a
l !Int
rs !Tree a
r = case Tree a
l of
  Bin Int
ls a
lx Tree a
ll Tree a
lr
    | Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
rs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ls -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
lx Tree a
ll (a -> Tree a -> Int -> Tree a -> Tree a
forall a. a -> Tree a -> Int -> Tree a -> Tree a
linkR a
x Tree a
lr Int
rs Tree a
r)
    | Bool
otherwise     -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
l Tree a
r
  Tree a
Tip -> String -> Tree a
forall a. HasCallStack => String -> a
error String
"Tree.linkR: impossible"

-- O(log (n1 + n2)). Link two trees.
merge :: Tree a -> Tree a -> Tree a
merge :: forall a. Tree a -> Tree a -> Tree a
merge Tree a
Tip Tree a
r = Tree a
r
merge Tree a
l Tree a
Tip = Tree a
l
merge l :: Tree a
l@(Bin Int
ls a
lx Tree a
ll Tree a
lr) r :: Tree a
r@(Bin Int
rs a
rx Tree a
rl Tree a
rr)
  | Int
ls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
rs = case a -> Tree a -> Tree a -> S2 (Tree a) a
forall a. a -> Tree a -> Tree a -> S2 (Tree a) a
unsnocSure a
lx Tree a
ll Tree a
lr of U.S2 Tree a
l' a
mx -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
link a
mx Tree a
l' Tree a
r
  | Bool
otherwise = case a -> Tree a -> Tree a -> S2 a (Tree a)
forall a. a -> Tree a -> Tree a -> S2 a (Tree a)
unconsSure a
rx Tree a
rl Tree a
rr of U.S2 a
mx Tree a
r' -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
link a
mx Tree a
l Tree a
r'
{-# INLINE merge #-}

-- O(log (n1 + n2)). Link two trees. Precondition: The trees must be balanced
-- wrt each other.
glue :: Tree a -> Tree a -> Tree a
glue :: forall a. Tree a -> Tree a -> Tree a
glue Tree a
Tip Tree a
r = Tree a
r
glue Tree a
l Tree a
Tip = Tree a
l
glue l :: Tree a
l@(Bin Int
ls a
lx Tree a
ll Tree a
lr) r :: Tree a
r@(Bin Int
rs a
rx Tree a
rl Tree a
rr)
  | Int
ls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
rs = case a -> Tree a -> Tree a -> S2 (Tree a) a
forall a. a -> Tree a -> Tree a -> S2 (Tree a) a
unsnocSure a
lx Tree a
ll Tree a
lr of U.S2 Tree a
l' a
m -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceR a
m Tree a
l' Tree a
r
  | Bool
otherwise = case a -> Tree a -> Tree a -> S2 a (Tree a)
forall a. a -> Tree a -> Tree a -> S2 a (Tree a)
unconsSure a
rx Tree a
rl Tree a
rr of U.S2 a
m Tree a
r' -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Tree a -> Tree a -> Tree a
balanceL a
m Tree a
l Tree a
r'
{-# INLINE glue #-}

-- Note [Balance]
-- ~~~~~~~~~~~~~~
-- The balancing code here is largely influenced by the implementation of
-- for the Set type in containers: https://hackage.haskell.org/package/containers
-- The linked papers in Data.Seqn.Seq describe the structure in greater detail.
--
-- To summarize:
-- * A tree is balanced if size(left child) < delta*size(right child) and vice
--   versa, which a special case for Tips. See `balanceOk` in `valid`, which is
--   used to check that balance holds in tests.
-- * Rebalancing involves rotations. The rotation can be single or double. The
--   constant `ratio` determines whether a double rotation is performed.

delta, ratio :: Int
delta :: Int
delta = Int
3
ratio :: Int
ratio = Int
2

-- O(1). Restores balance with at most one right rotation. Precondition: One
-- right rotation must be enough to restore balance. This is the case when the
-- left tree might have been inserted to or the right tree deleted from.
balanceL :: a -> Tree a -> Tree a -> Tree a
balanceL :: forall a. a -> Tree a -> Tree a -> Tree a
balanceL !a
x Tree a
l Tree a
r = case Tree a
r of
  Tree a
Tip -> case Tree a
l of
    Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip
    Bin Int
_ a
lx Tree a
ll Tree a
lr -> case Tree a
lr of
      Tree a
Tip -> case Tree a
ll of
        Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
2 a
x Tree a
l Tree a
forall a. Tree a
Tip
        Bin Int
_ a
_ Tree a
_ Tree a
_ -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
3 a
lx Tree a
ll (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
      Bin Int
_ a
lrx Tree a
_ Tree a
_ -> case Tree a
ll of
        Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
3 a
lrx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
lx Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
        Bin Int
_ a
_ Tree a
_ Tree a
_ -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
4 a
lx Tree a
ll (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
2 a
x Tree a
lr Tree a
forall a. Tree a
Tip)
  Bin Int
rs a
_ Tree a
_ Tree a
_ -> case Tree a
l of
    Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
forall a. Tree a
Tip Tree a
r
    Bin Int
ls a
lx Tree a
ll Tree a
lr
      | Int
ls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
rs -> case (Tree a
ll, Tree a
lr) of
        (Bin Int
lls a
_ Tree a
_ Tree a
_, Bin Int
lrs a
lrx Tree a
lrl Tree a
lrr)
          | Int
lrs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
lls -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
lx Tree a
ll (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lrs) a
x Tree a
lr Tree a
r)
          | Bool
otherwise -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
lrx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
llsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Tree a -> Int
forall a. Tree a -> Int
size Tree a
lrl) a
lx Tree a
ll Tree a
lrl) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Tree a -> Int
forall a. Tree a -> Int
size Tree a
lrr) a
x Tree a
lrr Tree a
r)
        (Tree a, Tree a)
_ -> String -> Tree a
forall a. HasCallStack => String -> a
error String
"Tree.balanceL: impossible"
      | Bool
otherwise -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
l Tree a
r
{-# NOINLINE balanceL #-}

-- O(1). Restores balance with at most one left rotation. Precondition: One left
-- rotation must be enough to restore balance. This is the case when the right
-- tree might have been inserted to or the left tree deleted from.
balanceR :: a -> Tree a -> Tree a -> Tree a
balanceR :: forall a. a -> Tree a -> Tree a -> Tree a
balanceR !a
x Tree a
l Tree a
r = case Tree a
l of
  Tree a
Tip -> case Tree a
r of
    Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip
    Bin Int
_ a
rx Tree a
rl Tree a
rr -> case Tree a
rl of
      Tree a
Tip -> case Tree a
rr of
        Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
2 a
x Tree a
forall a. Tree a
Tip Tree a
r
        Bin Int
_ a
_ Tree a
_ Tree a
_ -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
3 a
rx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip) Tree a
rr
      Bin Int
_ a
rlx Tree a
_ Tree a
_ -> case Tree a
rr of
        Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
3 a
rlx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
x Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
1 a
rx Tree a
forall a. Tree a
Tip Tree a
forall a. Tree a
Tip)
        Bin Int
_ a
_ Tree a
_ Tree a
_ -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
4 a
rx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin Int
2 a
x Tree a
forall a. Tree a
Tip Tree a
rl) Tree a
rr
  Bin Int
ls a
_ Tree a
_ Tree a
_ -> case Tree a
r of
    Tree a
Tip -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
ls) a
x Tree a
l Tree a
forall a. Tree a
Tip
    Bin Int
rs a
rx Tree a
rl Tree a
rr
      | Int
rs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
deltaInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
ls -> case (Tree a
rl, Tree a
rr) of
        (Bin Int
rls a
rlx Tree a
rll Tree a
rlr, Bin Int
rrs a
_ Tree a
_ Tree a
_)
          | Int
rls Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ratioInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
rrs -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
rx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rls) a
x Tree a
l Tree a
rl) Tree a
rr
          | Bool
otherwise -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
rlx (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Tree a -> Int
forall a. Tree a -> Int
size Tree a
rll) a
x Tree a
l Tree a
rll) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rrsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Tree a -> Int
forall a. Tree a -> Int
size Tree a
rlr) a
rx Tree a
rlr Tree a
rr)
        (Tree a, Tree a)
_ -> String -> Tree a
forall a. HasCallStack => String -> a
error String
"Tree.balanceR: impossible"
      | Bool
otherwise -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. Int -> a -> Tree a -> Tree a -> Tree a
Bin (Int
1Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
rs) a
x Tree a
l Tree a
r
{-# NOINLINE balanceR #-}

------------
-- Testing
------------

valid :: Tree a -> Bool
valid :: forall a. Tree a -> Bool
valid Tree a
s = Tree a -> Bool
forall a. Tree a -> Bool
balanceOk Tree a
s Bool -> Bool -> Bool
&& Tree a -> Bool
forall a. Tree a -> Bool
sizeOk Tree a
s
  where
    balanceOk :: Tree a -> Bool
balanceOk = \case
      Bin Int
_ a
_ Tree a
l Tree a
r -> Bool
ok Bool -> Bool -> Bool
&& Tree a -> Bool
balanceOk Tree a
l Bool -> Bool -> Bool
&& Tree a -> Bool
balanceOk Tree a
r
        where
          ok :: Bool
ok = Tree a -> Int
forall a. Tree a -> Int
size Tree a
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Tree a -> Int
forall a. Tree a -> Int
size Tree a
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 Bool -> Bool -> Bool
||
               (Tree a -> Int
forall a. Tree a -> Int
size Tree a
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
delta Int -> Int -> Int
forall a. Num a => a -> a -> a
* Tree a -> Int
forall a. Tree a -> Int
size Tree a
r Bool -> Bool -> Bool
&& Tree a -> Int
forall a. Tree a -> Int
size Tree a
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
delta Int -> Int -> Int
forall a. Num a => a -> a -> a
* Tree a -> Int
forall a. Tree a -> Int
size Tree a
l)
      Tree a
Tip -> Bool
True

    sizeOk :: Tree a -> Bool
sizeOk = \case
      Bin Int
sz a
_ Tree a
l Tree a
r -> Tree a -> Bool
sizeOk Tree a
l Bool -> Bool -> Bool
&& Tree a -> Bool
sizeOk Tree a
r Bool -> Bool -> Bool
&& Tree a -> Int
forall a. Tree a -> Int
size Tree a
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Tree a -> Int
forall a. Tree a -> Int
size Tree a
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
sz
      Tree a
Tip -> Bool
True

debugShowsPrec :: Show a => Int -> Tree a -> ShowS
debugShowsPrec :: forall a. Show a => Int -> Tree a -> String -> String
debugShowsPrec Int
p = \case
  Bin Int
sz a
x Tree a
l Tree a
r ->
    Bool -> (String -> String) -> String -> String
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) ((String -> String) -> String -> String)
-> (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$
      String -> String -> String
showString String
"Bin " (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      Int -> String -> String
forall a. Show a => a -> String -> String
shows Int
sz (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      String -> String -> String
showString String
" " (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      Int -> a -> String -> String
forall a. Show a => Int -> a -> String -> String
showsPrec Int
11 a
x (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      String -> String -> String
showString String
" " (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      Int -> Tree a -> String -> String
forall a. Show a => Int -> Tree a -> String -> String
debugShowsPrec Int
11 Tree a
l (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      String -> String -> String
showString String
" " (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
      Int -> Tree a -> String -> String
forall a. Show a => Int -> Tree a -> String -> String
debugShowsPrec Int
11 Tree a
r
  Tree a
Tip -> String -> String -> String
showString String
"Tip"