> {-# OPTIONS -F -pgmF arrowp-ext #-}

Homogeneous (or depth-preserving) functions over perfectly balanced trees.

> module Hom where

> import Control.Arrow
> import Control.Category
> import Data.Complex
> import Prelude hiding (id, (.))

> infixr 4 :&:

Consider the following non-regular type of perfectly balanced trees,
or `powertrees' (cf Jayadev Misra's powerlists):

> data Pow a = Zero a | Succ (Pow (Pair a))
> 	deriving Show

> type Pair a = (a, a)

Here are some example elements:

> tree0 = Zero 1
> tree1 = Succ (Zero (1, 2))
> tree2 = Succ (Succ (Zero ((1, 2), (3, 4))))
> tree3 = Succ (Succ (Succ (Zero (((1, 2), (3, 4)), ((5, 6), (7, 8))))))

The elements of this type have a string of constructors expressing
a depth n as a Peano numeral, enclosing a nested pair tree of 2^n
elements.  The type definition ensures that all elements of this type
are perfectly balanced binary trees of this form.  (Such things arise
in circuit design, eg Ruby, and descriptions of parallel algorithms.)
And the type system will ensure that all legal programs preserve
this structural invariant.
 
The only problem is that the type constraint is too restrictive, rejecting
many of the standard operations on these trees.  Typically you want to
split a tree into two subtrees, do some processing on the subtrees and
combine the results.  But the type system cannot discover that the two
results are of the same depth (and thus combinable).  We need a type
that says a function preserves depth.  Here it is:

> data Hom a b = (a -> b) :&: Hom (Pair a) (Pair b)

A homogeneous (or depth-preserving) function is an infinite sequence of
functions of type Pair^n a -> Pair^n b, one for each depth n.  We can
apply a homogeneous function to a powertree by selecting the function
for the required depth:

> apply :: Hom a b -> Pow a -> Pow b
> apply (f :&: fs) (Zero x) = Zero (f x)
> apply (f :&: fs) (Succ t) = Succ (apply fs t)

Having defined apply, we can forget about powertrees and do all our
programming with Hom's.  Firstly, Hom is an arrow:


> instance Category Hom where
>  id = id :&: id
>  (g :&: gs) . (f :&: fs) = (g . f) :&: (fs >>> gs)
>
> instance Arrow Hom where
> 	arr f = f :&: arr (f *** f)
> 	first (f :&: fs) =
> 		first f :&: (arr transpose >>> first fs >>> arr transpose)

> transpose :: ((a,b), (c,d)) -> ((a,c), (b,d))
> transpose ((a,b), (c,d)) = ((a,c), (b,d))

arr maps f over the leaves of a powertree.

The composition >>> composes sequences of functions pairwise.
 
The *** operator unriffles a powertree of pairs into a pair of powertrees,
applies the appropriate function to each and riffles the results.
It defines a categorical product for this arrow category.

When describing algorithms, one often provides a pure function for the
base case (trees of one element) and a (usually recursive) expression
for trees of pairs.

For example, a common divide-and-conquer pattern is the butterfly, where
one recursive call processes the odd-numbered elements and the other
processes the even ones (cf Geraint Jones and Mary Sheeran's Ruby papers):

> butterfly :: (Pair a -> Pair a) -> Hom a a
> butterfly f = id :&: proc (x, y) -> do
> 		x' <- butterfly f -< x
> 		y' <- butterfly f -< y
> 		returnA -< f (x', y')

The recursive calls operate on halves of the original tree, so the
recursion is well-defined.

Some examples of butterflies:

> rev :: Hom a a
> rev = butterfly swap
>	where	swap (x, y) = (y, x)

> unriffle :: Hom (Pair a) (Pair a)
> unriffle = butterfly transpose

Batcher's sorter for bitonic sequences:

> bisort :: Ord a => Hom a a
> bisort = butterfly cmp
> 	where	cmp (x, y) = (min x y, max x y)

This can be used (with rev) as the merge phase of a merge sort.
 
> sort :: Ord a => Hom a a
> sort = id :&: proc (x, y) -> do
>		x' <- sort -< x
>		y' <- sort -< y
>		yr <- rev -< y'
>		p <- unriffle -< (x', yr)
>		bisort2 -< p
>	where _ :&: bisort2 = bisort

Here is the scan operation, using the algorithm of Ladner and Fischer:

> scan :: (a -> a -> a) -> a -> Hom a a
> scan op b = id :&: proc (x, y) -> do
> 		y' <- scan op b -< op x y
> 		l <- rsh b -< y'
> 		returnA -< (op l x, y')

The auxiliary function rsh b shifts each element in the tree one place to
the right, placing b in the now-vacant leftmost position, and discarding
the old rightmost element:

> rsh :: a -> Hom a a
> rsh b = const b :&: proc (x, y) -> do
> 		w <- rsh b -< y
> 		returnA -< (w, x)

Finally, here is the Fast Fourier Transform:

> type C = Complex Double

> fft :: Hom C C
> fft = id :&: proc (x, y) -> do
> 		x' <- fft -< x
> 		y' <- fft -< y
> 		r <- roots (-1) -< ()
> 		let z = r*y'
> 		unriffle -< (x' + z, x' - z)

The auxiliary function roots r (where r is typically a root of unity)
populates a tree of size n (necessarily a power of 2) with the values
1, w, w^2, ..., w^(n-1), where w^n = r.

> roots :: C -> Hom () C
> roots r = const 1 :&: proc _ -> do
> 		x <- roots r' -< ()
> 		unriffle -< (x, x*r')
>	where	r' = if imagPart s >= 0 then -s else s
> 		s = sqrt r

Miscellaneous functions:

> rrot :: Hom a a
> rrot = id :&: proc (x, y) -> do
> 		w <- rrot -< y
> 		returnA -< (w, x)

> ilv :: Hom a a -> Hom (Pair a) (Pair a)
> ilv f = proc (x, y) -> do
> 		x' <- f -< x
> 		y' <- f -< y
> 		returnA -< (x', y')

> scan' :: (a -> a -> a) -> a -> Hom a a
> scan' op b = proc x -> do
>		l <- rsh b -< x
>		(id :&: ilv (scan' op b)) -< op l x

> riffle :: Hom (Pair a) (Pair a)
> riffle = id :&: proc ((x1, y1), (x2, y2)) -> do
>		x <- riffle -< (x1, x2)
>		y <- riffle -< (y1, y2)
>		returnA -< (x, y)

> invert :: Hom a a
> invert = id :&: proc (x, y) -> do
> 		x' <- invert -< x
> 		y' <- invert -< y
>		unriffle -< (x', y')

> carryLookaheadAdder :: Hom (Bool, Bool) Bool
> carryLookaheadAdder = proc (x, y) -> do
>		carryOut <- rsh (Just False) -<
>			if x == y then Just x else Nothing
>		Just carryIn <- scan plusMaybe Nothing -< carryOut
>		returnA -< x `xor` y `xor` carryIn
>	where	plusMaybe x Nothing = x
>		plusMaybe x (Just y) = Just y
>		False `xor` b = b
>		True `xor` b = not b

Global conditional for SIMD

> ifAll :: Hom a b -> Hom a b -> Hom (a, Bool) b
> ifAll fs gs = ifAllAux snd (arr fst >>> fs) (arr fst >>> gs)
>	where	ifAllAux :: (a -> Bool) -> Hom a b -> Hom a b -> Hom a b
>		ifAllAux p (f :&: fs) (g :&: gs) =
>			liftIf p f g :&: ifAllAux (liftAnd p) fs gs
>		liftIf p f g x = if p x then f x else g x
>		liftAnd p (x, y) = p x && p y

> maybeAll :: Hom a c -> Hom (a, b) c -> Hom (a, Maybe b) c
> maybeAll (n :&: ns) (j :&: js) =
>	choose :&: (arr dist >>> maybeAll ns (arr transpose >>> js))
>	where	choose (a, Nothing) = n a
>		choose (a, Just b) = j (a, b)
>		dist ((a1, b1), (a2, b2)) = ((a1, a2), zipMaybe b1 b2)
>		zipMaybe (Just x) (Just y) = Just (x, y)
>		zipMaybe _ _ = Nothing