okasaki

Okasaki's Purely Functional Data Structures
git clone git://git.jtobin.io/okasaki.git
Log | Files | Refs | LICENSE

Tree.hs (4153B)

      1 {-# OPTIONS_GHC -Wall -fno-warn-unused-top-binds #-}
      2 {-# LANGUAGE TemplateHaskell #-}
      3 
      4 module Okasaki.Tree (
      5     TreeF(..)
      6   , Tree
      7   , lef
      8   , nod
      9 
     10   , spy
     11 
     12   , put
     13   , pat
     14   , rub
     15 
     16   , has
     17   , haz
     18 
     19   , gas
     20   ) where
     21 
     22 import Data.Eq.Deriving (deriveEq1)
     23 import Data.Fix hiding (cata, ana, hylo)
     24 import Data.Functor.Foldable
     25 import Data.Monoid
     26 import Okasaki.Orphans ()
     27 import qualified Okasaki.Tree.CPS as CPS
     28 import Text.Show.Deriving
     29 
     30 data TreeF a r =
     31     LeafF
     32   | NodeF r !a r
     33   deriving (Functor, Show)
     34 
     35 $(deriveShow1 ''TreeF)
     36 $(deriveEq1 ''TreeF)
     37 
     38 type Tree a = Fix (TreeF a)
     39 
     40 lef :: Tree a
     41 lef = Fix LeafF
     42 
     43 nod :: a -> Tree a -> Tree a -> Tree a
     44 nod x l r = Fix (NodeF l x r)
     45 
     46 put :: Ord a => a -> Tree a -> Tree a
     47 put x = apo lag where
     48   lag pin = case project pin of
     49     LeafF -> NodeF (Left lef) x (Left lef)
     50     NodeF l e r -> case compare x e of
     51       EQ -> NodeF (Left l) e (Left r)
     52       LT -> NodeF (Right l) e (Left r)
     53       GT -> NodeF (Left l) e (Right r)
     54 
     55 --  NB would be interesting to benchmark these without optimizations to
     56 --     figure out the speed/memory profiles compared to the standard
     57 --     versions
     58 
     59 spy :: Ord a => a -> Tree a -> Maybe a
     60 spy x = getLast . cata alg where
     61   alg = \case
     62     LeafF -> mempty
     63     NodeF l e r
     64       | x < e     -> l
     65       | otherwise -> Last (Just e) <> r
     66 
     67 -- exercise 2.2 (max d + 1 comparisons)
     68 haz :: Ord a => a -> Tree a -> Bool
     69 haz x t = case spy x t of
     70   Nothing -> False
     71   Just s  -> s == x
     72 
     73 -- exercise 2.3 (no unnecessary copying)
     74 pat :: Ord a => a -> Tree a -> Tree a
     75 pat x t = tug id t where
     76   tug k s = case project s of
     77     LeafF -> k (nod x lef lef)
     78     NodeF l e r -> case compare x e of
     79       EQ -> t
     80       LT -> tug (\a -> k (nod e a r)) l
     81       GT -> tug (\a -> k (nod e l a)) r
     82 
     83 -- exercise 2.4 (no unnecessary copying, max d + 1 comparisons)
     84 rub :: Ord a => a -> Tree a -> Tree a
     85 rub x t = tug Nothing id t where
     86   tug c k s = case project s of
     87     LeafF -> case c of
     88       Nothing -> k (nod x lef lef)
     89       Just a
     90         | a == x    -> t
     91         | otherwise -> k (nod x lef lef)
     92 
     93     NodeF l e r
     94       | x < e     -> tug c (\a -> k (nod e a r)) l
     95       | otherwise -> tug (pure e) (\a -> k (nod e l a)) r
     96 
     97 -- exercise 2.5a (construct balanced binary trees of depth n)
     98 dap :: Ord a => a -> Int -> Tree a
     99 dap x n = ana lag (n, lef) where
    100   lag (j, t)
    101     | j <= 0 = LeafF
    102     | otherwise =
    103         let s = (pred j, t)
    104         in  NodeF s x s
    105 
    106 -- exercise 2.5b (construct mostly-balanced binary trees of size n)
    107 sap :: Ord a => a -> Int -> Tree a
    108 sap x n = ana lag (n, lef) where
    109   lag (j, t)
    110     | j <= 0 = LeafF
    111     | odd j  =
    112         let s = (j `quot` 2, t)
    113         in  NodeF s x s
    114     | otherwise =
    115         let l = j `quot` 2
    116             r = pred (j `quot` 2)
    117         in  NodeF (l, t) x (r, t)
    118 
    119 gas :: Ord a => [a] -> Tree a
    120 gas = cata $ \case
    121   Nil      -> lef
    122   Cons h t -> put h t
    123 
    124 dep :: Integral b => Tree a -> b
    125 dep = getSum . cata alg where
    126   alg = \case
    127     LeafF       -> mempty
    128     NodeF l _ r -> 1 <> max l r
    129 
    130 wyt :: Integral b => Tree a -> b
    131 wyt = getSum . cata alg where
    132   alg = \case
    133     LeafF       -> mempty
    134     NodeF l _ r -> 1 <> l <> r
    135 
    136 has :: Ord a => a -> Tree a -> Bool
    137 has x = cata $ \case
    138   LeafF       -> False
    139   NodeF l e r -> case compare x e of
    140     EQ -> True
    141     LT -> l
    142     GT -> r
    143 
    144 ver :: CPS.Tree a -> Tree a
    145 ver = ana lag where
    146   lag (project -> CPS.TreeF c) = c LeafF NodeF
    147 
    148 rev :: Tree a -> CPS.Tree a
    149 rev = ana lag where
    150   lag pin = case project pin of
    151     LeafF       -> CPS.lefF
    152     NodeF l e r -> CPS.nodF e l r
    153 
    154 test0 :: Num a => Tree a
    155 test0 = nod 1 lef lef
    156 
    157 -- 1
    158 -- | \
    159 -- L  L
    160 
    161 test1 :: Num a => Tree a
    162 test1 = nod 2 (nod 1 lef lef) (nod 3 lef lef)
    163 
    164 -- 2
    165 -- |   \
    166 -- 1    3
    167 -- | \  | \
    168 -- L  L L  L
    169 
    170 test2 :: (Ord a, Num a) => Tree a
    171 test2 = put 0 test1
    172 
    173 -- 2
    174 -- |   \
    175 -- 1    3
    176 -- | \  | \
    177 -- 0  L L  L
    178 -- | \
    179 -- L  L
    180 
    181 test3 :: (Ord a, Num a) => Tree a
    182 test3 = put 5 test1
    183 
    184 -- 2
    185 -- |   \
    186 -- 1    3
    187 -- | \  | \
    188 -- L  L L  5
    189 --         | \
    190 --         L  L
    191 
    192 test4 :: (Ord a, Num a) => Tree a
    193 test4 = put 4 test3
    194 
    195 -- 2
    196 -- |   \
    197 -- 1    3
    198 -- | \  | \
    199 -- L  L L  5
    200 --         | \
    201 --         4  L
    202 --         | \
    203 --         L  L
    204