A Treap implementation in Haskell with QuickCheck tests

Treap.hs

module Treap
  ( Treap (..)
  , empty
  , null
  , insert
  , delete
  , member
  , size
  , rInsert
  , rFromList
  , toList
  ) where

import Prelude hiding (null)
import System.Random
import System.Environment

-- | A Treap (https://en.wikipedia.org/wiki/Treap) is a data structure with is both a binary search
-- tree and a heap.
-- Every node in a treap has a value and a priority. The nodes form a BST over the values
-- and a heap over the priorities.
-- If the priorities are random numbers generated from a uniform distribution then the resulting
-- treap is balanced.
-- This treap implementation uses a max heap and keeps only unique elements.
data Treap k = Empty | Node !Int !k !(Treap k) !(Treap k) deriving (Eq)

instance (Show k) => Show (Treap k) where
  show Empty                   = "X"
  show (Node _ key left right) = show key ++ " [" ++ show left ++ " " ++ show right ++ "]"

-- | Creates an empty treap
empty :: (Ord k) => Treap k
empty = Empty

-- | Checks if a treap is empty
null :: (Ord k) => Treap k -> Bool
null Empty = True
null _     = False

priority :: Treap k -> Int
priority Empty          = minBound
priority (Node p _ _ _) = p

rotateLeft :: Treap k -> Treap k
rotateLeft (Node p k l (Node rp rk rl rr)) = Node rp rk (Node p k l rl) rr
rotateLeft _ = error "Wrong rotation (rotateLeft)"

rotateRight :: Treap k -> Treap k
rotateRight (Node p k (Node lp lk ll lr) r) = Node lp lk ll (Node p k lr r)
rotateRight _ = error "Wrong rotation (rotateRight)"

-- | Inserts a key in a treap using a given priority
insert :: (Ord k)
       => k       -- ^ key to insert
       -> Int     -- ^ priority of the key
       -> Treap k -- ^ treap to insert in
       -> Treap k -- ^ resulting treap
insert key prio Empty = Node prio key Empty Empty
insert key prio t@(Node p k l r)
  | key < k   = rotate $ Node p k (insert key prio l) r
  | key > k   = rotate $ Node p k l (insert key prio r)
  | otherwise = t
  where
    rotate Empty = Empty
    rotate t@(Node p _ l r)
      | p < priority l = rotateRight t
      | p < priority r = rotateLeft t
      | otherwise      = t

-- | Deletes a key from a treap
delete :: (Ord k)
       => k       -- ^ key to delete
       -> Treap k -- ^ treap to delete from
       -> Treap k -- ^ resulting treap
delete _ Empty = Empty
delete key t@(Node p k l r)
  | key < k    = Node p k (delete key l) r
  | key > k    = Node p k l (delete key r)
  | otherwise  = delRoot t
  where
    delRoot Empty                  = Empty
    delRoot (Node _ _ Empty Empty) = Empty
    delRoot t@(Node _ _ l r)
      | priority l < priority r = let (Node p k l' r') = rotateLeft t in Node p k (delRoot l') r'
      | otherwise               = let (Node p k l' r') = rotateRight t in Node p k l' (delRoot r')

-- | Checks if a given key exists in a treap
member :: (Ord k)
       => k       -- ^ key to check
       -> Treap k -- ^ treap to check in
       -> Bool
member _ Empty = False
member key (Node _ k l r)
  | key == k  = True
  | key < k   = member key l
  | otherwise = member key r

-- | Returns the size of a treap
size :: Treap k -> Int
size Empty          = 0
size (Node _ _ l r) = 1 + size l + size r

-- | Inserts a key in a treap with a random priority
rInsert :: (Ord k, RandomGen g)
        => k            -- ^ key to insert
        -> g            -- ^ a random generator
        -> Treap k      -- ^ treap to insert in
        -> (Treap k, g) -- ^ pair of the resulting treap and the next random generator
rInsert key gen tree = let (prio, gen') = random gen in (insert key prio tree, gen')

-- | Inserts a list of keys in a treap with random priorities
rFromList :: (Ord k, RandomGen g)
          => [k]          -- ^ keys to insert
          -> g            -- ^ a random generator
          -> (Treap k, g) -- ^ pair of the resulting treap and the next random generator
rFromList keys gen = foldl (\(t, g) k -> rInsert k g t) (empty, gen) keys

-- | Returns a list of all keys in the treap in ascending order
toList :: (Ord k) => Treap k -> [k]
toList Empty          = []
toList (Node _ k l r) = toList l ++ [k] ++ toList r

-- Count unique words in a file using a treap
countUniqueWords :: FilePath -> IO Int
countUniqueWords fileName = do
  wrds <- fmap words $ readFile fileName
  gen  <- newStdGen
  let (tree, _) = rFromList wrds gen
  return $ size tree

-- main
main :: IO ()
main = do
  args  <- getArgs
  count <- countUniqueWords (head args)
  putStrLn $ "Unique word count = " ++ show count

TreapTest.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- Tests for Treap. Run with: runhaskell TreapTest.hs -a 1000 --fail-fast
module TreapTest where

import Treap

import Control.Monad (liftM4)
import System.Random
import Test.Hspec
import Test.QuickCheck

-- Arbitarty instance for Treap
instance (Ord a, Arbitrary a, Num a, Random a) => Arbitrary (Treap a)  where
  arbitrary = gen 0 1000 maxBound where
    gen min max _ | (max - min) <= 3 = return Empty
    gen min max prio =
      if prio < 0
        then return Empty
        else do
          elt <- choose (min, max)
          p   <- choose (0, prio)
          frequency [ (1, return Empty)
                    , (6, liftM4 Node (return p) (return elt)
                                      (gen min (elt - 1) (p - 1)) (gen (elt + 1) max (p - 1)))]

newtype BoundedInt = BoundedInt Int deriving (Eq, Ord, Num)

instance Show BoundedInt where
  show (BoundedInt x) = show x

instance Random BoundedInt where
  randomR (BoundedInt lo, BoundedInt hi) g =
    let (a, g') = randomR (lo, hi) g in (BoundedInt a, g')

  random g = let (a, g') = random g in (BoundedInt a, g')

instance Arbitrary BoundedInt where
  arbitrary = fmap BoundedInt $ choose (0, 100)

isBST :: (Ord k) => Treap k -> Bool
isBST Empty                                          = True
isBST (Node _ _ Empty Empty)                         = True
isBST (Node _ k l@(Node _ lk _ _) Empty)             = lk < k && isBST l
isBST (Node _ k Empty r@(Node _ rk _ _))             = k < rk && isBST r
isBST (Node _ k l@(Node _ lk _ _) r@(Node _ rk _ _)) = lk < k && k < rk && isBST l && isBST r

isHeap :: Treap k -> Bool
isHeap Empty                                          = True
isHeap (Node _ _ Empty Empty)                         = True
isHeap (Node p _ l@(Node lp _ _ _) Empty)             = p >= lp && isHeap l
isHeap (Node p _ Empty r@(Node rp _ _ _))             = p >= rp && isHeap r
isHeap (Node p _ l@(Node lp _ _ _) r@(Node rp _ _ _)) = p >= lp && p >= rp && isHeap l && isHeap r

isTreap :: (Ord k) => Treap k -> Bool
isTreap t = isBST t && isHeap t

isAscUniq :: (Ord a) => [a] -> Bool
isAscUniq []       = True
isAscUniq [_]      = True
isAscUniq (x:y:ys) = x < y && isAscUniq (y:ys)

main :: IO ()
main = hspec $ do
  describe "empty" $
    it "is a valid Treap" $
      property $ isTreap (empty :: Treap Int)

  describe "insert" $ do
    it "preserves a valid Treap" $
      property $ \t (x :: Int) p -> isTreap t && isTreap (insert x p t)

    it "adds the element if it is not present in the tree" $
      property $ \t (x :: Int) p -> not (x `member` t) ==> member x (insert x p t)

    it "does not change the elements in the tree if the element is present in the tree" $
      property $ \t (BoundedInt x) p -> x `member` t ==> toList (insert x p t) == toList t

    it "increases size by one if the element is not present in the tree" $
      property $ \t (x :: Int) p -> not (x `member` t) ==> size (insert x p t) == size t + 1

    it "does not remove a present element" $
      property $ \t (BoundedInt x) (y :: Int) p -> x `member` t ==> x `member` insert y p t

  describe "delete" $ do
    it "preserves a valid Treap" $
      property $ \t (x :: Int) -> isTreap t && isTreap (delete x t)

    it "removes an element if it is present in the tree" $
      property $ \t (BoundedInt x) -> x `member` t ==> not (member x (delete x t))

    it "does not change the tree if the element is not present in the tree" $
      property $ \t (x :: Int) -> not (x `member` t) ==> delete x t == t

    it "decreases size by one if the element is present in the tree" $
      property $ \t (BoundedInt x) -> x `member` t ==> size (delete x t) == size t - 1

    it "does not insert an absent element" $
      property $ \t (x :: Int) (y :: Int) -> not (x `member` t) ==> not (x `member` delete y t)

  describe "member" $ do
    it "returns true for an element inserted into the tree" $
      property $ \t (x :: Int) p -> member x $ insert x p t

    it "returns false for an element deleted from the tree" $
      property $ \t (x :: Int) -> not $ member x $ delete x t

  describe "insert after delete for a present element" $
    it "does not change the elements in the tree" $
      property $ \t (BoundedInt x) p -> x `member` t ==> toList (insert x p (delete x t)) == toList t

  describe "delete after insert for an absent element" $
    it "does not change the tree" $
      property $ \t (x :: Int) p -> not (x `member` t) ==> delete x (insert x p t) == t

  describe "size" $
    it "is same as the size of list of elements in the tree" $
      property $ \(t :: Treap Int) -> size t == length (toList t)

  describe "toList" $
    it "produces an ascending list of unique elements" $
      property $ \(t :: Treap Int) -> isAscUniq $ toList t