{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_HADDOCK ignore-exports #-}

-- |
-- Module      : Data.Tree.AVL.Extern.Balance
-- Description : Balancing algorithm over ITree trees
-- Copyright   : (c) Nicolás Rodríguez, 2021
-- License     : GPL-3
-- Maintainer  : Nicolás Rodríguez
-- Stability   : experimental
-- Portability : POSIX
--
-- Implementation of the balancing algorithm over ITree trees for externalist AlmostAVL trees.
module Data.Tree.AVL.Extern.Balance
  ( Balanceable (Balance, balance),
    Balanceable' (Balance', balance'),
    Rotateable (Rotate, rotate),
  )
where

import Data.Proxy (Proxy (Proxy))
import Data.Tree.AVL.Invariants
  ( BS (Balanced, LeftHeavy, RightHeavy),
    BalancedState,
    Height,
    US (LeftUnbalanced, NotUnbalanced, RightUnbalanced),
    UnbalancedState,
  )
import Data.Tree.ITree
  ( ITree (EmptyITree, ForkITree),
    Tree (EmptyTree, ForkTree),
  )
import Data.Tree.Node (Node)
import Prelude ()

-- | This type class provides the functionality to balance
-- a tree @t@ without checking any structural invariant (key ordering or height balance).
-- The insertion is defined at the value level and the type level;
-- the verification of the @BST@/@AVL@ restrictions is performed after the insertion.
class Balanceable (t :: Tree) where
  type Balance (t :: Tree) :: Tree
  balance :: ITree t -> ITree (Balance t)

instance Balanceable 'EmptyTree where
  type Balance 'EmptyTree = 'EmptyTree
  balance :: ITree 'EmptyTree -> ITree (Balance 'EmptyTree)
balance ITree 'EmptyTree
_ = ITree 'EmptyTree
ITree (Balance 'EmptyTree)
EmptyITree

instance
  ( us ~ UnbalancedState (Height l) (Height r),
    Balanceable' ('ForkTree l (Node n a) r) us
  ) =>
  Balanceable ('ForkTree l (Node n a) r)
  where
  type Balance ('ForkTree l (Node n a) r) = Balance' ('ForkTree l (Node n a) r) (UnbalancedState (Height l) (Height r))
  balance :: ITree ('ForkTree l (Node n a) r)
-> ITree (Balance ('ForkTree l (Node n a) r))
balance ITree ('ForkTree l (Node n a) r)
t = ITree ('ForkTree l (Node n a) r)
-> Proxy us -> ITree (Balance' ('ForkTree l (Node n a) r) us)
forall (t :: Tree) (us :: US).
Balanceable' t us =>
ITree t -> Proxy us -> ITree (Balance' t us)
balance' ITree ('ForkTree l (Node n a) r)
t (Proxy us
forall {k} (t :: k). Proxy t
Proxy :: Proxy us)

-- | This type class provides the functionality to balance
-- a tree @t@ without checking any structural invariant (key ordering or height balance).
-- It's only used by the 'Balanceable' class and it has one extra parameter @us@,
-- which is the `UnbalancedState` of the two sub trees of @t@.
class Balanceable' (t :: Tree) (us :: US) where
  type Balance' (t :: Tree) (us :: US) :: Tree
  balance' :: ITree t -> Proxy us -> ITree (Balance' t us)

instance Balanceable' ('ForkTree l (Node n a) r) 'NotUnbalanced where
  type Balance' ('ForkTree l (Node n a) r) 'NotUnbalanced = ('ForkTree l (Node n a) r)
  balance' :: ITree ('ForkTree l (Node n a) r)
-> Proxy 'NotUnbalanced
-> ITree (Balance' ('ForkTree l (Node n a) r) 'NotUnbalanced)
balance' ITree ('ForkTree l (Node n a) r)
t Proxy 'NotUnbalanced
_ = ITree ('ForkTree l (Node n a) r)
ITree (Balance' ('ForkTree l (Node n a) r) 'NotUnbalanced)
t

instance
  ( bs ~ BalancedState (Height ll) (Height lr),
    Rotateable ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced bs
  ) =>
  Balanceable' ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced
  where
  type
    Balance' ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced =
      Rotate ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced (BalancedState (Height ll) (Height lr))
  balance' :: ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
-> Proxy 'LeftUnbalanced
-> ITree
     (Balance'
        ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
        'LeftUnbalanced)
balance' ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
t Proxy 'LeftUnbalanced
pus = ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
-> Proxy 'LeftUnbalanced
-> Proxy bs
-> ITree
     (Rotate
        ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
        'LeftUnbalanced
        bs)
forall (t :: Tree) (us :: US) (bs :: BS).
Rotateable t us bs =>
ITree t -> Proxy us -> Proxy bs -> ITree (Rotate t us bs)
rotate ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
t Proxy 'LeftUnbalanced
pus (Proxy bs
forall {k} (t :: k). Proxy t
Proxy :: Proxy bs)

instance
  ( bs ~ BalancedState (Height rl) (Height rr),
    Rotateable ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced bs
  ) =>
  Balanceable' ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced
  where
  type
    Balance' ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced =
      Rotate ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced (BalancedState (Height rl) (Height rr))
  balance' :: ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
-> Proxy 'RightUnbalanced
-> ITree
     (Balance'
        ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
        'RightUnbalanced)
balance' ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
t Proxy 'RightUnbalanced
pus = ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
-> Proxy 'RightUnbalanced
-> Proxy bs
-> ITree
     (Rotate
        ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
        'RightUnbalanced
        bs)
forall (t :: Tree) (us :: US) (bs :: BS).
Rotateable t us bs =>
ITree t -> Proxy us -> Proxy bs -> ITree (Rotate t us bs)
rotate ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
t Proxy 'RightUnbalanced
pus (Proxy bs
forall {k} (t :: k). Proxy t
Proxy :: Proxy bs)

-- | This type class provides the functionality to apply a rotation to
-- a tree @t@ without checking any structural invariant (key ordering or height balance).
-- The rotation is defined at the value level and the type level;
-- the verification of the @BST@/@AVL@ restrictions is performed after the insertion.
class Rotateable (t :: Tree) (us :: US) (bs :: BS) where
  type Rotate (t :: Tree) (us :: US) (bs :: BS) :: Tree
  rotate :: ITree t -> Proxy us -> Proxy bs -> ITree (Rotate t us bs)

-- Left-Left case (Right rotation)
instance Rotateable ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced 'LeftHeavy where
  type
    Rotate ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced 'LeftHeavy =
      ('ForkTree ll (Node ln la) ('ForkTree lr (Node n a) r))
  rotate :: ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
-> Proxy 'LeftUnbalanced
-> Proxy 'LeftHeavy
-> ITree
     (Rotate
        ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
        'LeftUnbalanced
        'LeftHeavy)
rotate (ForkITree (ForkITree ITree l
ll Node n a
lnode ITree r
lr) Node n a
node ITree r
r) Proxy 'LeftUnbalanced
_ Proxy 'LeftHeavy
_ = ITree l
-> Node n a
-> ITree ('ForkTree r (Node n a) r)
-> ITree ('ForkTree l (Node n a) ('ForkTree r (Node n a) r))
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
ll Node n a
lnode (ITree r -> Node n a -> ITree r -> ITree ('ForkTree r (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree r
lr Node n a
node ITree r
r)

instance Rotateable ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced 'Balanced where
  type
    Rotate ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r) 'LeftUnbalanced 'Balanced =
      ('ForkTree ll (Node ln la) ('ForkTree lr (Node n a) r))
  rotate :: ITree ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
-> Proxy 'LeftUnbalanced
-> Proxy 'Balanced
-> ITree
     (Rotate
        ('ForkTree ('ForkTree ll (Node ln la) lr) (Node n a) r)
        'LeftUnbalanced
        'Balanced)
rotate (ForkITree (ForkITree ITree l
ll Node n a
lnode ITree r
lr) Node n a
node ITree r
r) Proxy 'LeftUnbalanced
_ Proxy 'Balanced
_ = ITree l
-> Node n a
-> ITree ('ForkTree r (Node n a) r)
-> ITree ('ForkTree l (Node n a) ('ForkTree r (Node n a) r))
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
ll Node n a
lnode (ITree r -> Node n a -> ITree r -> ITree ('ForkTree r (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree r
lr Node n a
node ITree r
r)

-- Right-Right case (Left rotation)
instance Rotateable ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced 'RightHeavy where
  type
    Rotate ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced 'RightHeavy =
      ('ForkTree ('ForkTree l (Node n a) rl) (Node rn ra) rr)
  rotate :: ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
-> Proxy 'RightUnbalanced
-> Proxy 'RightHeavy
-> ITree
     (Rotate
        ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
        'RightUnbalanced
        'RightHeavy)
rotate (ForkITree ITree l
l Node n a
node (ForkITree ITree l
rl Node n a
rnode ITree r
rr)) Proxy 'RightUnbalanced
_ Proxy 'RightHeavy
_ = ITree ('ForkTree l (Node n a) l)
-> Node n a
-> ITree r
-> ITree ('ForkTree ('ForkTree l (Node n a) l) (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree (ITree l -> Node n a -> ITree l -> ITree ('ForkTree l (Node n a) l)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
l Node n a
node ITree l
rl) Node n a
rnode ITree r
rr

instance Rotateable ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced 'Balanced where
  type
    Rotate ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr)) 'RightUnbalanced 'Balanced =
      ('ForkTree ('ForkTree l (Node n a) rl) (Node rn ra) rr)
  rotate :: ITree ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
-> Proxy 'RightUnbalanced
-> Proxy 'Balanced
-> ITree
     (Rotate
        ('ForkTree l (Node n a) ('ForkTree rl (Node rn ra) rr))
        'RightUnbalanced
        'Balanced)
rotate (ForkITree ITree l
l Node n a
node (ForkITree ITree l
rl Node n a
rnode ITree r
rr)) Proxy 'RightUnbalanced
_ Proxy 'Balanced
_ = ITree ('ForkTree l (Node n a) l)
-> Node n a
-> ITree r
-> ITree ('ForkTree ('ForkTree l (Node n a) l) (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree (ITree l -> Node n a -> ITree l -> ITree ('ForkTree l (Node n a) l)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
l Node n a
node ITree l
rl) Node n a
rnode ITree r
rr

-- Left-Right case (First left rotation, then right rotation)
instance Rotateable ('ForkTree ('ForkTree ll (Node ln la) ('ForkTree lrl (Node lrn lra) lrr)) (Node n a) r) 'LeftUnbalanced 'RightHeavy where
  type
    Rotate ('ForkTree ('ForkTree ll (Node ln la) ('ForkTree lrl (Node lrn lra) lrr)) (Node n a) r) 'LeftUnbalanced 'RightHeavy =
      ('ForkTree ('ForkTree ll (Node ln la) lrl) (Node lrn lra) ('ForkTree lrr (Node n a) r))
  rotate :: ITree
  ('ForkTree
     ('ForkTree ll (Node ln la) ('ForkTree lrl (Node lrn lra) lrr))
     (Node n a)
     r)
-> Proxy 'LeftUnbalanced
-> Proxy 'RightHeavy
-> ITree
     (Rotate
        ('ForkTree
           ('ForkTree ll (Node ln la) ('ForkTree lrl (Node lrn lra) lrr))
           (Node n a)
           r)
        'LeftUnbalanced
        'RightHeavy)
rotate (ForkITree (ForkITree ITree l
ll Node n a
lnode (ForkITree ITree l
lrl Node n a
lrnode ITree r
lrr)) Node n a
node ITree r
r) Proxy 'LeftUnbalanced
_ Proxy 'RightHeavy
_ =
    ITree ('ForkTree l (Node n a) l)
-> Node n a
-> ITree ('ForkTree r (Node n a) r)
-> ITree
     ('ForkTree
        ('ForkTree l (Node n a) l) (Node n a) ('ForkTree r (Node n a) r))
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree (ITree l -> Node n a -> ITree l -> ITree ('ForkTree l (Node n a) l)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
ll Node n a
lnode ITree l
lrl) Node n a
lrnode (ITree r -> Node n a -> ITree r -> ITree ('ForkTree r (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree r
lrr Node n a
node ITree r
r)

-- Right-Left case (First right rotation, then left rotation)
instance Rotateable ('ForkTree l (Node n a) ('ForkTree ('ForkTree rll (Node rln rla) rlr) (Node rn ra) rr)) 'RightUnbalanced 'LeftHeavy where
  type
    Rotate ('ForkTree l (Node n a) ('ForkTree ('ForkTree rll (Node rln rla) rlr) (Node rn ra) rr)) 'RightUnbalanced 'LeftHeavy =
      ('ForkTree ('ForkTree l (Node n a) rll) (Node rln rla) ('ForkTree rlr (Node rn ra) rr))
  rotate :: ITree
  ('ForkTree
     l
     (Node n a)
     ('ForkTree ('ForkTree rll (Node rln rla) rlr) (Node rn ra) rr))
-> Proxy 'RightUnbalanced
-> Proxy 'LeftHeavy
-> ITree
     (Rotate
        ('ForkTree
           l
           (Node n a)
           ('ForkTree ('ForkTree rll (Node rln rla) rlr) (Node rn ra) rr))
        'RightUnbalanced
        'LeftHeavy)
rotate (ForkITree ITree l
l Node n a
node (ForkITree (ForkITree ITree l
rll Node n a
rlnode ITree r
rlr) Node n a
rnode ITree r
rr)) Proxy 'RightUnbalanced
_ Proxy 'LeftHeavy
_ =
    ITree ('ForkTree l (Node n a) l)
-> Node n a
-> ITree ('ForkTree r (Node n a) r)
-> ITree
     ('ForkTree
        ('ForkTree l (Node n a) l) (Node n a) ('ForkTree r (Node n a) r))
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree (ITree l -> Node n a -> ITree l -> ITree ('ForkTree l (Node n a) l)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree l
l Node n a
node ITree l
rll) Node n a
rlnode (ITree r -> Node n a -> ITree r -> ITree ('ForkTree r (Node n a) r)
forall a (l :: Tree) (n :: Nat) (r :: Tree).
Show a =>
ITree l -> Node n a -> ITree r -> ITree ('ForkTree l (Node n a) r)
ForkITree ITree r
rlr Node n a
rnode ITree r
rr)