Foldables vs Recursion Schemes

bt.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

import Data.Functor.Foldable
import Data.Functor.Foldable.TH

-- it's important that a comes after the left sub-tree so that default foldable
-- implementation picks up the correct ordering

data Tree a = TNil | Tree a (Tree a) (Tree a)
    deriving (Show)

makeBaseFunctor 'Tree

-- generic implementation, equivalent with default foldMap if we derived Foldable for Tree
-- the problem would be the ordering of the fields, "a" would have to be the second field for this to work
-- this is also a bit more relaxed by not requiring the default element to be an identity of the semigroup operation

foldMapTreeWithDefault :: (Semigroup m ) => m -> (a -> m) -> Tree a -> m
foldMapTreeWithDefault d f = cata alg
    where alg TNilF = d
          alg (TreeF x l r) = l <> f x <> r

----------------------------------------------
-- instances for Range
----------------------------------------------

data Range a = Range a a 
             | REmpty
             | RInvalid
    deriving (Show)

relaxLowerBound :: (Ord a) => a -> Range a -> Range a
relaxLowerBound x RInvalid = RInvalid
relaxLowerBound x REmpty = Range x x
relaxLowerBound x (Range y z) = if x <= y then Range x z else RInvalid

instance Ord a => Semigroup (Range a) where
    REmpty <> x = x
    RInvalid <> _ = RInvalid
    (Range a x) <> y = (relaxLowerBound a . relaxLowerBound x) y

instance Ord a => Monoid (Range a) where
    mempty = REmpty

----------------------------------------------
-- instances for BalanceInfo
----------------------------------------------

data BalanceInfo a = BalanceInfo a
                   | BNode
                   | BInvalid

instance (Num a, Ord a) => Semigroup (BalanceInfo a) where
    BNode <> x = x
    x <> BNode = x
    (BalanceInfo x) <> (BalanceInfo y) | abs (x - y) <= 1 = BalanceInfo $ max x y + 1
    _ <> _ = BInvalid

----------------------------------------------
-- catamorphisms
----------------------------------------------

foldMapTree :: (Monoid m) => (a -> m) -> Tree a -> m
foldMapTree = foldMapTreeWithDefault mempty

getRange :: (Ord a) => Tree a -> Range a
-- here the default is mempty so we could also just use foldMapTree (eg. FoldMap)
getRange = foldMapTreeWithDefault REmpty  (\x -> Range x x)
--getRange = foldMapTree (\x -> Range x x)

getBalanceInfo :: (Num a, Ord a) => Tree a -> BalanceInfo a
-- here the default BalanceInfo 0 is not an identity so foldMap is not powerful enough
getBalanceInfo = foldMapTreeWithDefault (BalanceInfo 0) (\x -> BNode)

----------------------------------------------
-- wrappers
----------------------------------------------

isBinarySearchTree t = 
    case getRange t of
        Range _ _ -> True
        _ -> False

isBalancedTree t = 
    case getBalanceInfo t of
        BalanceInfo x -> True
        _ -> False


example :: Tree Int
example = 
    Tree 
        12
        (Tree 
            5
            (Tree 3 TNil TNil) 
            (Tree 10 TNil TNil))
        (Tree 
            14
            (Tree 13 TNil TNil) 
            (Tree 20 TNil TNil))

main :: IO ()
main = do
    print $ isBinarySearchTree example
    print $ isBalancedTree example