Bird tree solution

BirdTree.hs

module BirdTree where

import Control.Monad.Maybe
import Control.Monad.State

import Data.ByteString.Char8 ( ByteString )
import qualified Data.ByteString.Char8 as B

import Data.Char ( isSpace )
import Data.Functor ( (<$>) )
import Data.List ( intercalate )
import Data.Maybe ( fromJust )
import Data.Ratio ( (%) )

{-

A bird tree is an infinite binary tree, whose definition is as follows:

bird =         1/1

              /   \
             /     \
            /       \

     1/(bird+1)   (1/bird)+1

-}

------------------------------------------------------------------------
-- Entry point
------------------------------------------------------------------------

main :: IO ()
main = (run <$> B.readFile "bird.in") >>= B.writeFile "bird.out"
    where
        run = convey . contemplate . consume

consume :: ByteString -> [Rational]
consume = fromJust . consume'
    where
        consume' = parse $ do
            n <- extractor B.readInt
            replicateM n (skipSpaces >> readFraction)

contemplate :: [Rational] -> [[Direction]]
contemplate = map findInTree

convey :: [[Direction]] -> ByteString
convey = B.unlines . map (B.pack . map directionToChar)

directionToChar :: Direction -> Char
directionToChar L = 'L'
directionToChar R = 'R'


------------------------------------------------------------------------
-- Main code
------------------------------------------------------------------------

birdTree :: Tree Rational
birdTree = Branch (1 % 1) --> fmap (\n -> 1 / (n + 1)) birdTree
                          --> fmap (\n -> (1 / n) + 1) birdTree
    where
        f --> x = f x

data Direction = L | R
    deriving (Enum, Eq, Ord, Show)

findInTree :: Rational -> [Direction]
findInTree = search birdTree False
    where
        search (Branch value left right) isOdd target
            | value == target = []
            | isOdd `xor` (target < value) = L : search left (not isOdd) target
            | otherwise                    = R : search right (not isOdd) target


------------------------------------------------------------------------
-- Parsing
------------------------------------------------------------------------

-- | A rudimentary non-backtracking parser, implemented using monad
-- transformers.
type Extractor a = MaybeT (State ByteString) a

-- | Wrap a plain reading function (such as 'B.readInt') in an
-- extractor monad.
extractor :: (ByteString -> Maybe (a, ByteString)) -> Extractor a
extractor f = do
    input <- get
    case f input of
        Just (result, remainder) -> do
            put remainder
            return result
        Nothing -> mzero

-- | Parse a ByteString, returning a Just if it was successful or
-- Nothing if not.
parse :: Extractor a -> ByteString -> Maybe a
parse = evalState . runMaybeT

readFraction :: Extractor Rational
readFraction = do
    num <- extractor B.readInteger -- Numerator
    slash <- extractor B.uncons
    guard (slash == '/')           -- Slash
    den <- extractor B.readInteger -- Denominator
    return (num % den)

skipSpaces :: Extractor ()
skipSpaces = modify (B.dropWhile isSpace)


------------------------------------------------------------------------
-- Data structures and utility functions
------------------------------------------------------------------------

-- | An infinite binary tree. Each node has an attached value and two
-- children.
data Tree v = Branch v (Tree v) (Tree v)

instance Functor Tree where
    fmap f (Branch value left right) = Branch (f value) (fmap f left) (fmap f right)

-- | Get a node's attached value.
rootValue :: Tree v -> v
rootValue (Branch value _left _right) = value

-- | Get the direct children of a node, as a two element list.
children :: Tree v -> [Tree v]
children (Branch _value left right) = [left, right]

-- | Get a list of nodes at each level of the tree.
levels :: Tree v -> [[Tree v]]
levels tree = iterate (concatMap children) [tree]

-- | Dump a human-readable representation of the tree, down to a
-- specified depth. Useful for debugging.
--
-- The depth should be greater than or equal to 0.
showTree :: Show v => Int -> Tree v -> String
showTree depth | depth < 0 = error "BirdTree.showTree: depth must be positive"
               | otherwise = showLevels . take depth . levels

showLevels :: Show v => [[Tree v]] -> String
showLevels = unlines . map (intercalate " | " . map (show . rootValue))

-- | Boolean XOR.
xor :: Bool -> Bool -> Bool
xor True False = True
xor False True = True
xor _ _ = False