{- Example of AVL trees by michaelbeaumont -} {-@ LIQUID "--totality" @-} module AVL (Tree, singleton, insert) where -- Basic functions {-@ data Tree [ht] a = Nil | Tree (x::a) (l::Tree a) (r::Tree a) @-} data Tree a = Nil | Tree a (Tree a) (Tree a) deriving Show {-@ measure ht @-} ht :: Tree a -> Int ht Nil = 0 ht (Tree x l r) = if (ht l) > (ht r) then (1 + ht l) else (1 + ht r) {-@ invariant {v:Tree a | 0 <= ht v} @-} {-@ measure bFac @-} bFac Nil = 0 bFac (Tree v l r) = ht l - ht r {-@ htDiff :: s:Tree a -> t: Tree a -> {v: Int | HtDiff s t v} @-} htDiff :: Tree a -> Tree a -> Int htDiff l r = ht l - ht r {-@ emp :: {v: AVLTree | ht v == 0} @-} emp = Nil {-@ singleton :: a -> {v: AVLTree | ht v == 1 }@-} singleton a = Tree a Nil Nil -- | Insert functions {-@ decrease insert 3 @-} {-@ insert :: a -> s: AVLTree -> {t: AVLTree | EqHt t s || HtDiff t s 1 } @-} insert :: (Ord a) => a -> Tree a -> Tree a insert a Nil = singleton a insert a t@(Tree v l r) = case compare a v of LT -> insL GT -> insR EQ -> t where r' = insert a r l' = insert a l insL | siblDiff > 1 && bFac l' == (0-1) = rebalanceLR v l' r | siblDiff > 1 && bFac l' == 1 = rebalanceLL v l' r | siblDiff <= 1 = Tree v l' r | otherwise = t where siblDiff = htDiff l' r insR | siblDiff > 1 && bFac r' == 1 = rebalanceRL v l r' | siblDiff > 1 && bFac r' == (0-1) = rebalanceRR v l r' | siblDiff <= 1 = Tree v l r' | otherwise = t where siblDiff = htDiff r' l {-@ rebalanceLL :: a -> l:{AVLTree | LeftHeavy l } -> r:{AVLTree | HtDiff l r 2} -> {t:AVLTree | EqHt t l } @-} rebalanceLL v (Tree lv ll lr) r = Tree lv ll (Tree v lr r) {-@ rebalanceLR :: a -> l:{AVLTree | RightHeavy l } -> r:{AVLTree | HtDiff l r 2 } -> {t: AVLTree | EqHt t l } @-} rebalanceLR v (Tree lv ll (Tree lrv lrl lrr)) r = Tree lrv (Tree lv ll lrl) (Tree v lrr r) {-@ rebalanceRR :: a -> l: AVLTree -> r: {AVLTree | RightHeavy r && HtDiff r l 2 } -> {t: AVLTree | EqHt t r } @-} rebalanceRR v l (Tree rv rl rr) = Tree rv (Tree v l rl) rr {-@ rebalanceRL :: a -> l: AVLTree -> r:{AVLTree | LeftHeavy r && HtDiff r l 2} -> {t: AVLTree | EqHt t r } @-} rebalanceRL v l (Tree rv (Tree rlv rll rlr) rr) = Tree rlv (Tree v l rll) (Tree rv rlr rr) -- Test main = do mapM_ print [a,b,c,d] where a = singleton 5 b = insert 2 a c = insert 3 b d = insert 7 c -- Liquid Haskell {-@ predicate HtDiff S T D = (ht S) - (ht T) == D @-} {-@ predicate EqHt S T = (ht S) == (ht T) @-} {-@ predicate LeftHeavy T = bFac T == 1 @-} {-@ predicate RightHeavy T = bFac T == -1 @-} {-@ measure balanced :: Tree a -> Bool balanced (Nil) = true balanced (Tree v l r) = ((ht l) <= (ht r) + 1) && (ht r <= ht l + 1) && (balanced l) && (balanced r) @-} {-@ type AVLTree = {v: Tree a | balanced v} @-}