Monad-like trees with fun applications like integer factoring

annotree.hs

#!/usr/bin/env runhaskell
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, TypeSynonymInstances, FunctionalDependencies #-}

{- Monad-like trees can do cool things like factoring integers and 
 - pruning leaves with bind (>>=).
 -
 - Internal nodes are annotated with a measure of some sort that can be used
 - for efficient indexing and sizing. 
 -
 - Measured typeclass and tagged nodes from 
 - http://apfelmus.nfshost.com/articles/monoid-fingertree.html
 -}

module AnnoTree where

import Prelude 
import Data.Monoid

class Sequence s a where
    -- A Sequence is anything isomorphic to a list
    fromList :: [a] -> s a
    toList :: s a -> [a]
    -- the above let us lift list ops to the sequence
    -- including filter
    sfilter :: (a -> Bool) -> s a -> [a]
    sfilter p v = filter p $ toList v
    -- and foldr
    sfoldr :: (a -> b -> b) -> b  -> s a -> b
    sfoldr f v0 v = foldr f v0 $ toList v
    -- and monoid ops
    sempty :: s a
    sempty = fromList []
    sappend :: s a -> s a -> s a
    sappend v w = fromList $ (toList v) ++ (toList w)
    sconcat :: [s a] -> s a
    sconcat [] = sempty
    sconcat (v:vs) = sappend v (sconcat vs)

-- LeafyTree instances
data LeafyTree a = LTEmpty | LTLeaf a | LTBranch (LeafyTree a) (LeafyTree a) deriving (Show)

-- builds a tree from the bottom up: divide, construct, and join
ltbuild :: [LeafyTree a] -> [LeafyTree a]
ltbuild [t] = [t]
ltbuild [t1,t2] = [LTBranch t1 t2]
ltbuild xs = ltbuild $ (ltbuild xs1) ++ (ltbuild xs2)
    where l = length xs
          l2 = (l `div` 2) + (l `mod` 2)
          (xs1, xs2) = splitAt l2 xs

-- removes empty leaves
prune :: LeafyTree a -> LeafyTree a
prune (LTBranch LTEmpty t) = prune t
prune (LTBranch t LTEmpty) = prune t
prune t = t

-- collapsing our tree into a list gets us filtering and 
-- monoid ops 'for free', expensively
instance Sequence LeafyTree a where
    fromList [] = LTEmpty
    fromList xs = (ltbuild $ map LTLeaf xs) !! 0
    toList LTEmpty = []
    toList (LTLeaf x) = [x]
    toList (LTBranch l r) = (toList l) ++ (toList r)

instance Functor LeafyTree where
    fmap _ LTEmpty = LTEmpty
    fmap f (LTLeaf x) = LTLeaf (f x)
    fmap f (LTBranch l r) = prune $ LTBranch (fmap f l) (fmap f r)

instance Monad LeafyTree where
    return x = LTLeaf x
    LTEmpty >>= f = LTEmpty
    (LTLeaf x) >>= f = (f x) 
    (LTBranch l r) >>= f = prune $ LTBranch (l >>= f) (r >>= f)

instance Monoid (LeafyTree a) where
    mempty = sempty
    mappend = sappend

-- some examples of construction
test_leafy_tree = do
    putStrLn "test_leafy_tree"
    putStrLn . show $ (fromList [] :: LeafyTree Int)
    putStrLn . show $ (fromList [1] :: LeafyTree Int)
    putStrLn . show $ (fromList [1..2] :: LeafyTree Int)
    putStrLn . show $ (fromList [1..3] :: LeafyTree Int)
    putStrLn . show $ (fromList [1..4] :: LeafyTree Int)
    putStrLn . show $ (fromList [1..5] :: LeafyTree Int)
    putStrLn . show $ (fromList [1..6] :: LeafyTree Int)

-- error is awful... should be working in the Maybe monad
-- but there are only so many hours in the day
-- and days in a life
first :: (a -> Bool) -> [a] -> a
first p [] = error "Nothing matches criterion"
first p (x:xs) | p x = x
               | otherwise = first p xs

-- the poor man's factoring routine.
firstFactor :: Int -> Int
firstFactor x | x <= 0 = error "naturals only" 
              | x <= 3 = x
              | (mod x 2 == 0) = 2
              | otherwise = first (\f -> (mod x f) == 0) ([3,5..(div x 2)]++[x])

-- we can bind a function into our tree monad to factor integers
treeFactor :: Int -> LeafyTree Int
treeFactor x | x <= 0 = error "naturals only"
             | x <= 3 = LTLeaf x
             | otherwise = 
                if (f == x) then (LTLeaf x)
                            else LTBranch (treeFactor f) (treeFactor (div x f))
                where f = firstFactor x

-- some examples of integer factoring
test_factor_ints = do
    putStrLn "test_factor_ints"
    putStrLn . show $ firstFactor 5
    putStrLn . show $ (return 5 :: LeafyTree Int) >>= treeFactor
    putStrLn . show $ (return 10 :: LeafyTree Int) >>= treeFactor
    putStrLn . show $ (return 20 :: LeafyTree Int) >>= treeFactor
    putStrLn . show $ (return 12345678 :: LeafyTree Int) >>= treeFactor
    let t = ((return 12345678 :: LeafyTree Int) >>= treeFactor)
    let t2 = t >>= treeFactor
    let t3 = (LTBranch (LTLeaf 20) (LTLeaf 341)) >>= treeFactor
    let t4 = (sappend (LTLeaf 2) t3) >>= treeFactor
    putStrLn . show $ toList t
    putStrLn . show $ toList t2
    putStrLn . show $ toList t3 
    putStrLn . show $ toList t4
    putStrLn . show $ sfoldr (*) 1 t
    putStrLn . show $ sfoldr (*) 1 t2
    putStrLn . show $ sfoldr (*) 1 t3
    putStrLn . show $ sfoldr (*) 1 t4

-- we can also bind a function to prune our tree of any even factors
oddClobber :: Int -> LeafyTree Int
oddClobber x | mod x 2 == 0 = LTEmpty
             | otherwise = treeFactor x

test_pruning = do
    putStrLn "test_pruning"
    let t = treeFactor 12345678
    let t2 = t >>= oddClobber
    putStrLn . show $ t
    putStrLn . show $ t2

main = do
    test_leafy_tree
    test_factor_ints
    test_pruning