{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, Safe #-}

{-
  This module is part of Chatty.
  Copyleft (c) 2014 Marvin Cohrs

  All wrongs reversed. Sharing is an act of love, not crime.
  Please share Chatty with everyone you like.

  Chatty is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  Chatty is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with Chatty. If not, see <http://www.gnu.org/licenses/>.
-}

-- | Provides an AVL tree.
module Data.Chatty.AVL (avlMax,avlMin,avlLookup,avlHeight,avlSize,avlInsert,avlRemove,AVL (EmptyAVL,AVL),avlRoot,avlPreorder,avlPostorder,avlInorder) where

import Data.Maybe
import Data.Chatty.BST
import Data.Chatty.None

instance Indexable i o v => AnyBST AVL i o v where
  anyBstMax :: AVL i -> Maybe i
anyBstMax = AVL i -> Maybe i
forall i. AVL i -> Maybe i
avlMax
  anyBstMin :: AVL i -> Maybe i
anyBstMin = AVL i -> Maybe i
forall i. AVL i -> Maybe i
avlMin
  anyBstLookup :: o -> AVL i -> Maybe v
anyBstLookup = o -> AVL i -> Maybe v
forall i o v. Indexable i o v => o -> AVL i -> Maybe v
avlLookup
  anyBstEmpty :: AVL i
anyBstEmpty = AVL i
forall a. AVL a
EmptyAVL
  anyBstInsert :: i -> AVL i -> AVL i
anyBstInsert = i -> AVL i -> AVL i
forall i o v. Indexable i o v => i -> AVL i -> AVL i
avlInsert
  anyBstRemove :: o -> AVL i -> AVL i
anyBstRemove = o -> AVL i -> AVL i
forall i o v. Indexable i o v => o -> AVL i -> AVL i
avlRemove
  anyBstHead :: AVL i -> Maybe i
anyBstHead = AVL i -> Maybe i
forall i. AVL i -> Maybe i
avlHead
  anyBstInorder :: AVL i -> [i]
anyBstInorder = AVL i -> [i]
forall i. AVL i -> [i]
avlInorder

instance None (AVL a) where
  none :: AVL a
none = AVL a
forall a. AVL a
EmptyAVL

-- | An AVL tree.
data AVL a = EmptyAVL | AVL a Int Int !(AVL a) !(AVL a)

-- | Get the greatest element.
avlMax :: AVL i -> Maybe i
avlMax :: AVL i -> Maybe i
avlMax AVL i
EmptyAVL = Maybe i
forall a. Maybe a
Nothing
avlMax (AVL i
a Int
_ Int
_ AVL i
_ AVL i
EmptyAVL) = i -> Maybe i
forall a. a -> Maybe a
Just i
a
avlMax (AVL i
_ Int
_ Int
_ AVL i
_ AVL i
r) = AVL i -> Maybe i
forall i. AVL i -> Maybe i
avlMax AVL i
r

-- | Get the least element.
avlMin :: AVL i -> Maybe i
avlMin :: AVL i -> Maybe i
avlMin AVL i
EmptyAVL = Maybe i
forall a. Maybe a
Nothing
avlMin (AVL i
a Int
_ Int
_ AVL i
EmptyAVL AVL i
_) = i -> Maybe i
forall a. a -> Maybe a
Just i
a
avlMin (AVL i
_ Int
_ Int
_ AVL i
l AVL i
_) = AVL i -> Maybe i
forall i. AVL i -> Maybe i
avlMin AVL i
l

-- | Lookup a given key.
avlLookup :: Indexable i o v => o -> AVL i -> Maybe v
avlLookup :: o -> AVL i -> Maybe v
avlLookup o
_ AVL i
EmptyAVL = Maybe v
forall a. Maybe a
Nothing
avlLookup o
o (AVL i
a Int
_ Int
_ AVL i
l AVL i
r)
  | o
o o -> o -> Bool
forall a. Eq a => a -> a -> Bool
== i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a = v -> Maybe v
forall a. a -> Maybe a
Just (v -> Maybe v) -> v -> Maybe v
forall a b. (a -> b) -> a -> b
$ i -> v
forall i o v. Indexable i o v => i -> v
valueOf i
a
  | o
o o -> o -> Bool
forall a. Ord a => a -> a -> Bool
< i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a = o -> AVL i -> Maybe v
forall i o v. Indexable i o v => o -> AVL i -> Maybe v
avlLookup o
o AVL i
l
  | o
o o -> o -> Bool
forall a. Ord a => a -> a -> Bool
> i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a = o -> AVL i -> Maybe v
forall i o v. Indexable i o v => o -> AVL i -> Maybe v
avlLookup o
o AVL i
r

-- | Lookup if a given key is contained
avlContains :: Indexable i o v => o -> AVL i -> Bool
avlContains :: o -> AVL i -> Bool
avlContains o
o = Maybe v -> Bool
forall a. Maybe a -> Bool
isJust (Maybe v -> Bool) -> (AVL i -> Maybe v) -> AVL i -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. o -> AVL i -> Maybe v
forall i o v. Indexable i o v => o -> AVL i -> Maybe v
avlLookup o
o

-- | Get the height of the tree.
avlHeight :: AVL i -> Int
avlHeight :: AVL i -> Int
avlHeight AVL i
EmptyAVL = Int
0
avlHeight (AVL i
_ Int
_ Int
h AVL i
_ AVL i
_) = Int
h

-- | Get the size of the tree.
avlSize :: AVL i -> Int
avlSize :: AVL i -> Int
avlSize AVL i
EmptyAVL = Int
0
avlSize (AVL i
_ Int
s Int
_ AVL i
_ AVL i
_) = Int
s

avlBalance :: AVL i -> AVL i
avlBalance :: AVL i -> AVL i
avlBalance AVL i
EmptyAVL = AVL i
forall a. AVL a
EmptyAVL
avlBalance t :: AVL i
t@(AVL i
a Int
_ Int
_ AVL i
l AVL i
r)
  | Int -> Int
forall a. Num a => a -> a
abs (AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
r) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = AVL i
t
  | AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
r = case AVL i
r of
    AVL i
a1 Int
_ Int
_ AVL i
l1 AVL i
r1 ->
      let child :: AVL i
child = i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l AVL i
l1) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l AVL i
l1) AVL i
l AVL i
l1
      in i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a1 (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
child AVL i
r1) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
child AVL i
r1) AVL i
child AVL i
r1
  | Bool
otherwise = case AVL i
l of
    AVL i
a1 Int
_ Int
_ AVL i
l1 AVL i
r1 ->
      let child :: AVL i
child = i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
r1 AVL i
r) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
r1 AVL i
r) AVL i
r1 AVL i
r
      in i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a1 (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l1 AVL i
child) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l1 AVL i
child) AVL i
l1 AVL i
child

findSize :: AVL i -> AVL i -> Int
findSize :: AVL i -> AVL i -> Int
findSize AVL i
a AVL i
b = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ AVL i -> Int
forall i. AVL i -> Int
avlSize AVL i
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ AVL i -> Int
forall i. AVL i -> Int
avlSize AVL i
b

findHeight :: AVL i -> AVL i -> Int
findHeight :: AVL i -> AVL i -> Int
findHeight AVL i
a AVL i
b = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
a) (AVL i -> Int
forall i. AVL i -> Int
avlHeight AVL i
b)

-- | Insert into the tree.
avlInsert :: Indexable i o v => i -> AVL i -> AVL i
avlInsert :: i -> AVL i -> AVL i
avlInsert i
a AVL i
EmptyAVL = i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a Int
1 Int
1 AVL i
forall a. AVL a
EmptyAVL AVL i
forall a. AVL a
EmptyAVL
avlInsert i
a (AVL i
a1 Int
s Int
h AVL i
l AVL i
r)
  | i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a o -> o -> Bool
forall a. Eq a => a -> a -> Bool
== i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a1 = i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a Int
s Int
h AVL i
l AVL i
r
  | i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a o -> o -> Bool
forall a. Ord a => a -> a -> Bool
< i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a1 =
    let l' :: AVL i
l' = i -> AVL i -> AVL i
forall i o v. Indexable i o v => i -> AVL i -> AVL i
avlInsert i
a AVL i
l
    in AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a1 (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l' AVL i
r) AVL i
l' AVL i
r
  | Bool
otherwise =
    let r' :: AVL i
r' = i -> AVL i -> AVL i
forall i o v. Indexable i o v => i -> AVL i -> AVL i
avlInsert i
a AVL i
r
    in AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a1 (Int
sInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l AVL i
r') AVL i
l AVL i
r' 

-- | Remove from the tree.
avlRemove :: Indexable i o v => o -> AVL i -> AVL i
avlRemove :: o -> AVL i -> AVL i
avlRemove o
_ AVL i
EmptyAVL = AVL i
forall a. AVL a
EmptyAVL
avlRemove o
o t :: AVL i
t@(AVL i
a Int
_ Int
_ AVL i
EmptyAVL AVL i
EmptyAVL)
  | i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a o -> o -> Bool
forall a. Eq a => a -> a -> Bool
== o
o = AVL i
forall a. AVL a
EmptyAVL
  | Bool
otherwise = AVL i
t
avlRemove o
o t :: AVL i
t@(AVL i
a Int
_ Int
_ AVL i
l AVL i
r)
  | i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a o -> o -> Bool
forall a. Eq a => a -> a -> Bool
== o
o =
    case AVL i
t of
      AVL i
_ Int
_ Int
_ AVL i
EmptyAVL AVL i
_ -> case AVL i -> (Maybe i, AVL i)
forall i. AVL i -> (Maybe i, AVL i)
getLeft AVL i
r of
        (Just i
a',AVL i
r') -> AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a' (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
forall a. AVL a
EmptyAVL AVL i
r') (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
forall a. AVL a
EmptyAVL AVL i
r') AVL i
forall a. AVL a
EmptyAVL AVL i
r'
      AVL i
_ -> case AVL i -> (Maybe i, AVL i)
forall i. AVL i -> (Maybe i, AVL i)
getRight AVL i
l of
        (Just i
a',AVL i
l') -> AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a' (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l' AVL i
r) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l' AVL i
r) AVL i
l' AVL i
r
  | o
o o -> o -> Bool
forall a. Ord a => a -> a -> Bool
< i -> o
forall i o v. Indexable i o v => i -> o
indexOf i
a =
    let l' :: AVL i
l' = o -> AVL i -> AVL i
forall i o v. Indexable i o v => o -> AVL i -> AVL i
avlRemove o
o AVL i
l
    in AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l' AVL i
r) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l' AVL i
r) AVL i
l' AVL i
r
  | Bool
otherwise =
    let r' :: AVL i
r' = o -> AVL i -> AVL i
forall i o v. Indexable i o v => o -> AVL i -> AVL i
avlRemove o
o AVL i
r
    in AVL i -> AVL i
forall i. AVL i -> AVL i
avlBalance (AVL i -> AVL i) -> AVL i -> AVL i
forall a b. (a -> b) -> a -> b
$ i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l AVL i
r') (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l AVL i
r') AVL i
l AVL i
r'

getLeft :: AVL i -> (Maybe i,AVL i)
getLeft :: AVL i -> (Maybe i, AVL i)
getLeft AVL i
EmptyAVL = (Maybe i
forall a. Maybe a
Nothing,AVL i
forall a. AVL a
EmptyAVL)
getLeft (AVL i
a Int
_ Int
_ AVL i
EmptyAVL AVL i
EmptyAVL) = (i -> Maybe i
forall a. a -> Maybe a
Just i
a,AVL i
forall a. AVL a
EmptyAVL)
getLeft (AVL i
a Int
_ Int
_ AVL i
EmptyAVL AVL i
r) = (i -> Maybe i
forall a. a -> Maybe a
Just i
a,AVL i
r)
getLeft (AVL i
a Int
_ Int
_ AVL i
l AVL i
r) =
  case AVL i -> (Maybe i, AVL i)
forall i. AVL i -> (Maybe i, AVL i)
getLeft AVL i
l of
    (Maybe i
p, AVL i
t2) -> (Maybe i
p, i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
r AVL i
t2) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
r AVL i
t2) AVL i
t2 AVL i
r)

getRight :: AVL i -> (Maybe i,AVL i)
getRight :: AVL i -> (Maybe i, AVL i)
getRight AVL i
EmptyAVL = (Maybe i
forall a. Maybe a
Nothing,AVL i
forall a. AVL a
EmptyAVL)
getRight (AVL i
a Int
_ Int
_ AVL i
EmptyAVL AVL i
EmptyAVL) = (i -> Maybe i
forall a. a -> Maybe a
Just i
a,AVL i
forall a. AVL a
EmptyAVL)
getRight (AVL i
a Int
_ Int
_ AVL i
l AVL i
EmptyAVL) = (i -> Maybe i
forall a. a -> Maybe a
Just i
a,AVL i
l)
getRight (AVL i
a Int
_ Int
_ AVL i
l AVL i
r) =
  case AVL i -> (Maybe i, AVL i)
forall i. AVL i -> (Maybe i, AVL i)
getRight AVL i
r of
    (Maybe i
p, AVL i
t2) -> (Maybe i
p, i -> Int -> Int -> AVL i -> AVL i -> AVL i
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL i
a (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findSize AVL i
l AVL i
t2) (AVL i -> AVL i -> Int
forall i. AVL i -> AVL i -> Int
findHeight AVL i
l AVL i
t2) AVL i
l AVL i
t2)

instance Functor AVL where
  fmap :: (a -> b) -> AVL a -> AVL b
fmap a -> b
_ AVL a
EmptyAVL = AVL b
forall a. AVL a
EmptyAVL
  fmap a -> b
f (AVL a
a Int
s Int
h AVL a
l AVL a
r) = b -> Int -> Int -> AVL b -> AVL b -> AVL b
forall a. a -> Int -> Int -> AVL a -> AVL a -> AVL a
AVL (a -> b
f a
a) Int
s Int
h ((a -> b) -> AVL a -> AVL b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f AVL a
l) ((a -> b) -> AVL a -> AVL b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f AVL a
r)

-- | Get the root of the tree.
avlRoot :: AVL i -> i
avlRoot :: AVL i -> i
avlRoot AVL i
EmptyAVL = [Char] -> i
forall a. HasCallStack => [Char] -> a
error [Char]
"Trying to get the root of an empty AVL tree."
avlRoot (AVL i
a Int
_ Int
_ AVL i
_ AVL i
_) = i
a

-- | Get the root of the tree (safely)
avlHead :: AVL i -> Maybe i
avlHead :: AVL i -> Maybe i
avlHead AVL i
EmptyAVL = Maybe i
forall a. Maybe a
Nothing
avlHead AVL i
t = i -> Maybe i
forall a. a -> Maybe a
Just (i -> Maybe i) -> i -> Maybe i
forall a b. (a -> b) -> a -> b
$ AVL i -> i
forall i. AVL i -> i
avlRoot AVL i
t

-- | Traverse the tree, order (head, left, right)
avlPreorder :: AVL i -> [i]
avlPreorder :: AVL i -> [i]
avlPreorder AVL i
EmptyAVL = []
avlPreorder (AVL i
a Int
_ Int
_ AVL i
l AVL i
r) = i
a i -> [i] -> [i]
forall a. a -> [a] -> [a]
: AVL i -> [i]
forall i. AVL i -> [i]
avlPreorder AVL i
l [i] -> [i] -> [i]
forall a. [a] -> [a] -> [a]
++ AVL i -> [i]
forall i. AVL i -> [i]
avlPreorder AVL i
r

-- | Traverse the tree, order (left, right, head)
avlPostorder :: AVL i -> [i]
avlPostorder :: AVL i -> [i]
avlPostorder AVL i
EmptyAVL = []
avlPostorder (AVL i
a Int
_ Int
_ AVL i
l AVL i
r) = AVL i -> [i]
forall i. AVL i -> [i]
avlPostorder AVL i
l [i] -> [i] -> [i]
forall a. [a] -> [a] -> [a]
++ AVL i -> [i]
forall i. AVL i -> [i]
avlPostorder AVL i
r [i] -> [i] -> [i]
forall a. [a] -> [a] -> [a]
++ [i
a]

-- | Traverse the tree, order (left, head, right)
avlInorder :: AVL i -> [i]
avlInorder :: AVL i -> [i]
avlInorder AVL i
EmptyAVL = []
avlInorder (AVL i
a Int
_ Int
_ AVL i
l AVL i
r) = AVL i -> [i]
forall i. AVL i -> [i]
avlInorder AVL i
l [i] -> [i] -> [i]
forall a. [a] -> [a] -> [a]
++ [i
a] [i] -> [i] -> [i]
forall a. [a] -> [a] -> [a]
++ AVL i -> [i]
forall i. AVL i -> [i]
avlInorder AVL i
r