Finger Trees

FingerTrees.hs

{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
module Main where
import Data.Foldable
import Data.Traversable
import Data.Monoid

class Monoid v => Measured a v | a -> v where
  measure :: a -> v

data Digit v a = One !a
               | Two !a !a
               | Three !a !a !a
               | Four !a !a !a !a
             deriving(Functor, Foldable, Traversable, Show)

data Node v a = Node2 a a
              | Node3 a a a
            deriving(Functor, Foldable, Traversable, Show)

data FingerTree v a = Empty
                    | Single !a
                    | Deep !v (Digit v a) (FingerTree v (Node v a)) (Digit v a)
                    deriving(Functor, Foldable, Traversable, Show)

instance Measured a v => Measured (Digit v a) v where
  measure = foldMap measure
instance Measured a v => Measured (Node v a) v where
  measure = foldMap measure
instance Measured a v => Measured (FingerTree v a) v where
  measure Empty = mempty
  measure (Single a) = measure a
  measure (Deep v _ _ _) = v

deep :: Measured a v =>
        Digit v a -> FingerTree v (Node v a) -> Digit v a -> FingerTree v a
deep pr m fs = Deep (measure pr <> measure m <> measure fs) pr m fs

(|>) :: Measured a v => a -> FingerTree v a -> FingerTree v a
a |> Empty = Single a
a |> Single b = Deep (measure a <> measure b) (One a) Empty (One b)
a |> Deep _ (Four b c d e) m sf = deep (Two a b) (Node3 c d e |> m) sf
a |> Deep _ pr m sf = deep pr' m sf
  where pr' = case pr of
               One b -> Two a b
               Two b c -> Three a b c
               Three b c d -> Four a b c d
               Four {} -> error "Impossible"
infixr 5 |>
(<|) :: Measured a v => FingerTree v a -> a -> FingerTree v a
Empty <| a = Single a
Single a <| b = deep (One a) Empty (One b)
Deep _ pr m (Four a b c d) <| e = deep pr (m <| Node3 a b c) (Two d e)
Deep _ pr m sf <| x = deep pr m sf'
  where sf' = case sf of
               One a -> Two a x
               Two a b -> Three a b x
               Three a b c -> Four a b c x
               Four {} -> error "Impossible"
infixl 5 <|

nodeToFix :: Node v a -> Digit v a
nodeToFix (Node2 a b) = Two a b
nodeToFix (Node3 a b c) = Three a b c

uncons :: Measured a v => FingerTree v a -> Maybe (a, FingerTree v a)
uncons Empty = Nothing
uncons (Single a) = Just (a, Empty)
uncons (Deep _ (Two a b) m sf) = Just (a, deep (One b) m sf)
uncons (Deep _ (Three a b c) m sf) = Just (a, deep (Two b c) m sf)
uncons (Deep _ (Four a b c d) m sf) = Just (a, deep (Three b c d) m sf)
uncons (Deep _ (One x) m sf) = Just (x, viewRest)
  where viewRest = case uncons m of
          Just (node, r) -> deep (nodeToFix node) r sf
          Nothing -> case sf of
            One a -> Single a
            Two a b -> deep (One a) Empty (One b)
            Three a b c -> deep (Two a b) Empty (One c)
            Four a b c d -> deep (Two a b) Empty (Two c d)

unsnoc :: Measured a v => FingerTree v a -> Maybe (a, FingerTree v a)
unsnoc Empty = Nothing
unsnoc (Single a) = Just (a, Empty)
unsnoc (Deep _ pr m (Two a b)) = Just (b, deep pr m (One a))
unsnoc (Deep _ pr m (Three a b c)) = Just (c, deep pr m (Two a b))
unsnoc (Deep _ pr m (Four a b c d)) = Just (d, deep pr m (Three a b c))
unsnoc (Deep _ pr m (One x)) = Just (x, viewRest)
  where viewRest = case unsnoc m of
          Just (node, r) -> deep (nodeToFix node) r pr
          Nothing -> case pr of
            One a -> Single a
            Two a b -> deep (One a) Empty (One b)
            Three a b c -> deep (Two a b) Empty (One c)
            Four a b c d -> deep (Two a b) Empty (Two c d)

head :: Measured a v => FingerTree v a -> Maybe a
head = fmap fst . uncons

last :: Measured a v => FingerTree v a -> Maybe a
last = fmap fst . unsnoc

--- Demo Time ---

instance Measured Int () where
  measure = const ()

buildN :: Int -> FingerTree () Int -> FingerTree () Int
buildN 0 f = f
buildN n f = let f' = n |> f
                 n' = n - 1
             in n' `seq` buildN n' f'

main :: IO ()
main = print . foldMap Sum . buildN 10000000 $ Empty