StablePhases.hs

{-# LANGUAGE TupleSections, GADTs, DeriveFunctor, DeriveFoldable #-}

import Control.Applicative

import Data.Foldable (traverse_,toList)
import Test.QuickCheck
import qualified Data.List
import Control.Monad.State

--------------------------------------------------------------------------------
-- Heap type (Like phases, doesn't track insertion order or size)
--------------------------------------------------------------------------------

data Heap k f a where
  Pure :: a -> Heap k f a
  Root :: !k
       -> {-# UNPACK #-} !Int
       -> (x -> y -> a) 
       -> f x
       -> Heaps k f y 
       -> Heap k f a

data Heaps k f a where
  Nil :: Heaps k f ()
  App :: !k
      -> {-# UNPACK #-} !Int
      -> f x 
      -> Heaps k f y
      -> Heaps k f z
      -> Heaps k f (x,y,z)

instance Functor (Heap k f) where
  fmap f (Pure x) = Pure (f x)
  fmap f (Root k n c x xs) = Root k n (\a b -> f (c a b)) x xs

instance Ord k => Applicative (Heap k f) where
  pure = Pure
  Pure f <*> xs = fmap f xs
  xs <*> Pure f = fmap ($ f) xs

  Root xk xn xc xs xss <*> Root yk yn yc ys yss
    | (xk,xn) <= (yk,yn) = Root xk xn (\a (b,c,d) -> xc a d (yc b c)) xs (App yk (yn - xn) ys yss xss)
    | otherwise          = Root yk yn (\a (b,c,d) -> xc b c (yc a d)) ys (App xk (xn - yn) xs xss yss)

--------------------------------------------------------------------------------
-- Phases type (Tracks insertion order, making reordering stable)
--------------------------------------------------------------------------------

data Phases k f a = Phases {-# UNPACK #-} !Word !(Heap k f a)
  deriving Functor
  
instance Ord k => Applicative (Phases k f) where
  pure = Phases 0 . pure
  Phases n xs <*> Phases m ys = Phases (n+m) (xs <*> delay n ys)
    where
      delay _ xs@(Pure _) = xs
      delay n (Root k m c x xs) = Root k (fromEnum n + m) c x xs

phase :: k -> f a -> Phases k f a
phase k fa = Phases 1 (Root k 0 const fa Nil)

-- |   
-- >>> :{
-- let out c = c <$ putStrLn (c : " out")
-- in runPhases $ sequenceA
--  [ phase 3 (out 'a')
--  , phase 2 (out 'b')
--  , phase 1 (out 'c')
--  , phase 2 (out 'd')
--  , phase 3 (out 'e') ]
-- :}
-- c out
-- b out
-- d out
-- a out
-- e out
-- "abcde"
runPhases :: (Ord k, Applicative f) => Phases k f a -> f a
runPhases (Phases _ h) = runHeap h
  where
    runHeap :: (Ord k, Applicative f) => Heap k f a -> f a
    runHeap (Pure x) = pure x
    runHeap (Root _ _ c x xs) = liftA2 c x (runHeap (merges xs))
    
    merges :: (Ord k, Applicative f) => Heaps k f a -> Heap k f a
    merges Nil = Pure ()
    merges (App k1 n1 e1 t1 Nil) = Root k1 n1 (,,()) e1 t1
    merges (App k1 n1 e1 t1 (App k2 n2 e2 t2 xs)) = 
       (Root k1 n1 (\a b cd es -> (a,b, cd es)) e1 t1 <*> Root k2 n2 (,,) e2 t2) 
         <*> merges xs
       
--------------------------------------------------------------------------------
-- Tests
--------------------------------------------------------------------------------

sortOn :: Ord k => (a -> k) -> [a] -> [a]
sortOn k = fst . runPhases . traverse_ (\x -> phase (k x) ([x],()))

sortOnIsStable :: [(Word,Word)] -> Property
sortOnIsStable xs = Data.List.sortOn snd xs === sortOn snd xs

data Tree a 
  = Leaf a 
  | Tree a :*: Tree a
  deriving (Show, Eq, Ord, Functor, Foldable)
  
instance Traversable Tree where
  traverse f (Leaf x) = Leaf <$> f x
  traverse f (xs :*: ys) = liftA2 (:*:) (traverse f xs) (traverse f ys)
  
instance Arbitrary a => Arbitrary (Tree a) where
  arbitrary = sized go
    where
      go n | n <= 1 = Leaf <$> arbitrary
      go n = frequency [(1,Leaf <$> arbitrary), (n, branch)]
        where 
          branch = do
            m <- choose (1,n-1)
            liftA2 (:*:) (go m) (go (n-m))
            
  shrink (Leaf _) = []
  shrink (xs :*: ys) = xs : ys : map (uncurry (:*:)) (shrink (xs,ys))

fill :: Traversable t => t a -> [b] -> t b
fill = evalState . traverse (\_ -> state (\(x:xs) -> (x,xs)))
  
sortOnT :: (Ord k, Traversable t) => (a -> k) -> t a -> t a
sortOnT k xs = fill xs lst
  where
    lst = fst (runPhases (traverse (\x -> phase (k x) ([x],())) xs))

sortOnTree :: Ord k => (a -> k) -> Tree a -> Tree a
sortOnTree k t = fill t (Data.List.sortOn k (toList t))
    
sortOnTIsStable :: Tree (Word,Word) -> Property
sortOnTIsStable xs = sortOnTree snd xs === sortOnT snd xs