Fast Applicative sorting

gistfile1.txt

{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# OPTIONS_GHC -Wall -fwarn-incomplete-uni-patterns #-}

module AS2 where

import Data.Type.Equality
import Unsafe.Coerce (unsafeCoerce)
import Data.Proxy

-- | Natural numbers
data Nat = Z | S Nat

-- | A sorted vector
data SVec :: Nat -> * -> * where
  Nil :: SVec 'Z a
  Cons :: a -> SVec n a -> SVec ('S n) a

-- Convenience function for turning vectors into proxies
prox :: SVec n a -> Proxy n
prox _ = Proxy

type family (:+) m n where
  'Z :+ n = n
  ('S m) :+ n = 'S (m :+ n)

-- --------------------------------------------------------------------
--
-- Several basic proofs about natural number arithmetic. We use rewrite
-- rules to assert that these will actually terminate when run, which
-- saves a *lot* of time.
--
-- --------------------------------------------------------------------

plusSucc :: SVec m a -> proxy n q -> (m :+ 'S n) :~: ('S (m :+ n))
plusSucc Nil _ = Refl
plusSucc (Cons _ xs) p = case plusSucc xs p of Refl -> Refl
{-# NOINLINE plusSucc #-}

-- This rule declares that plusSucc m n will terminate, which makes
-- merging *much* more efficient.
{-# RULES
"plusSucc" forall v p . plusSucc v p = unsafeCoerce (Refl :: 'Z :~: 'Z)
 #-}

plusZero :: SVec m a -> (m :+ 'Z) :~: m
plusZero Nil = Refl
plusZero (Cons _ xs) = case plusZero xs of Refl -> Refl
{-# NOINLINE plusZero #-}

-- This rule declares that plusZero m will terminate.
{-# RULES
"plusZero" forall v . plusZero v = unsafeCoerce (Refl :: 'Z :~: 'Z)
 #-}

plusAssoc :: SVec m a -> proxy n -> proxy o -> (m :+ (n :+ o)) :~: ((m :+ n) :+ o)
plusAssoc Nil _ _ = Refl
plusAssoc (Cons _ m) n p = case plusAssoc m n p of Refl -> Refl
{-# NOINLINE plusAssoc #-}

{-# RULES
"plusAssoc" forall v p q. plusAssoc v p q = unsafeCoerce (Refl :: 'Z :~: 'Z)
 #-}

-- -------------------------------------------------------------------

-- | Merge two sorted vectors into one.
merge :: Ord x => SVec m x -> SVec n x -> SVec (m :+ n) x
merge Nil ys = ys
merge xs Nil = case plusZero xs of Refl -> xs
merge xss@(Cons x xs) yss@(Cons y ys)
  | x <= y = Cons x (merge xs yss)
  | otherwise = case plusSucc xss ys of Refl -> Cons y (merge xss ys)

-- This is where the fun happens. We make an Applicative out of a sorted vector
-- of length m along with a function that consumes any sorted vector of length
-- at least m and returns a result and the remainder of the vector.
data Sort x a where
  Sort :: (forall n. Proxy n -> SVec (m :+ n) x -> (a, SVec n x)) -> SVec m x -> Sort x a
  -- The proxy argument isn't strictly necessary, but it makes the type easier
  -- to work with.

instance Functor (Sort x) where
  fmap f (Sort g xs) = Sort (\p v -> case g p v of (a, r) -> (f a, r)) xs

instance Ord x => Applicative (Sort x) where
  pure a = Sort (\_ xs -> (a, xs)) Nil

  -- (<*>) merges the vectors and produces a compound function that
  -- uses the first part of the vector to produce one result and
  -- the next to produce the other, then applies the first to the second.
  (<*>) :: forall a b . Sort x (a -> b) -> Sort x a -> Sort x b
  Sort f (xs :: SVec m x) <*> Sort g (ys :: SVec n x) =
    Sort h (merge xs ys)
    where
      h :: forall o . Proxy o -> SVec ((m :+ n) :+ o) x -> (b, SVec o x)
      h p v = case plusAssoc xs (prox ys) p of
        Refl -> case f (Proxy :: Proxy (n :+ o)) v of { (a, v') ->
                case g (Proxy :: Proxy o) v' of { (b, v'') ->
                  (a b, v'')}}

liftSort :: x -> Sort x x
liftSort a = Sort (\_ (Cons x xs) -> (x, xs)) (Cons a Nil)

runSort :: forall x a . Sort x a -> a
runSort (Sort f xs) = case plusZero xs of
    Refl -> fst $ f (Proxy :: Proxy 'Z) xs

-- | Sort an arbitrary Traversable container. For a list-like instance,
-- this will perform an insertion sort. For a tree-like instance, it will
-- perform a merge sort.
sortTraversable :: (Ord a, Traversable t) => t a -> t a
sortTraversable = runSort . traverse liftSort

sortTraversal :: ((a -> Sort a a) -> t -> Sort a t) -> t -> t
sortTraversal trav = runSort . trav liftSort