Pairing queues with linear interface

gistfile1.txt

{-# language BangPatterns #-}
{-# language LinearTypes #-}
{-# language GADTs #-}
{-# language StandaloneKindSignatures #-}
{-# language KindSignatures #-}
{-# language TypeApplications #-}
{-# language DataKinds #-}
{-# language ScopedTypeVariables #-}

-- | Pairing heap loosely based on one Donnacha Oisín Kidney wrote for
-- 'Data.Sequence.unstableSortBy'. This one is more general (it can be empty),
-- but it's probably also slower.
module Data.Linear.Pairing where

import Data.Unrestricted.Linear
import qualified Data.List.Linear as LL
import Prelude hiding (Foldable (..))
import qualified Data.Foldable as F
import qualified Data.Functor.Linear as LF
import qualified Data.Monoid.Linear as LM
import qualified Prelude.Linear as PL
import Data.List.NonEmpty (NonEmpty (..))

data PQueue a where
  Empty :: PQueue a
  Q :: {-# UNPACK #-} !(NEQueue a) -> PQueue a
data NEQueue a where
  NEQ :: !a -> !(NEQList a) -> NEQueue a
data NEQList a where
  Nil :: NEQList a
  NEQCons :: {-# UNPACK #-} !(NEQueue a)
          -> !(NEQList a)
          -> NEQList a

withPQueueFromList :: Ord a => [Ur a] %1-> (PQueue a %1-> Ur r) %1-> Ur r
withPQueueFromList as f = f (fromList as)

withNEQueueFromNonEmpty :: Ord a => NonEmpty (Ur a) %1-> (NEQueue a %1-> Ur r) %1-> Ur r
withNEQueueFromNonEmpty as f = f (fromNonEmptyNE as)

withEmptyPQueue :: Ord a => (PQueue a %1-> Ur r) %1-> Ur r
withEmptyPQueue f = f Empty

withNEQueueFromSingleton :: Ord a => a -> (NEQueue a %1-> Ur r) %1-> Ur r
withNEQueueFromSingleton a f = f (singletonNEQ a)

fromAscList :: [Ur a] %1-> Ur (PQueue a)
fromAscList [] = Ur Empty
fromAscList (x : xs) = case fromAscNonEmptyNEh x xs of
  Ur r -> Ur (Q r)

fromAscNonEmptyNE :: NonEmpty (Ur a) %1-> Ur (NEQueue a)
fromAscNonEmptyNE (x :| xs) = fromAscNonEmptyNEh x xs

fromAscNonEmptyNEh :: Ur a %1-> [Ur a] %1-> Ur (NEQueue a)
fromAscNonEmptyNEh (Ur a) [] = Ur (NEQ a Nil)
fromAscNonEmptyNEh (Ur a) (b : bs) = case fromAscNonEmptyNEh b bs of
  Ur r -> Ur (NEQ a (NEQCons r Nil))

instance Consumable (PQueue a) where
  consume Empty = ()
  consume (Q _) = ()

instance Ord a => Dupable (PQueue a) where
  dupR q = case move q of
    Ur q' -> LF.pure q'

instance Ord a => Semigroup (PQueue a) where
  x <> y = merge x y
instance Ord a => Semigroup (NEQueue a) where
  x <> y = mergeNEQ x y
instance Ord a => Monoid (PQueue a) where
  mempty = empty

instance Ord a => LM.Semigroup (PQueue a) where
  x <> y = merge x y
instance Ord a => LM.Semigroup (NEQueue a) where
  x <> y = mergeNEQ x y
instance Ord a => LM.Monoid (PQueue a) where
  mempty = empty

instance Ord a => Eq (PQueue a) where
  p == q = toPlainList p == toPlainList q
instance Ord a => PL.Eq (PQueue a) where
  p == q = toList p PL.== toList q
instance Ord a => Ord (PQueue a) where
  p `compare` q = toPlainList p `compare` toPlainList q
instance Ord a => PL.Ord (PQueue a) where
  p `compare` q = toList p `PL.compare` toList q

-- | Rebalances the queue in \(O(n \log n)\) time so subsequent operations are
-- cheap.
instance Ord a => Movable (PQueue a) where
  move q = fromAscList (toList q)

instance Consumable (NEQueue a) where
  consume (NEQ _ _) = ()

instance Ord a => Dupable (NEQueue a) where
  dupR q = case move q of
    Ur q' -> LF.pure q'

-- | Rebalances the queue in \(O(n \log n)\) time so subsequent operations are
-- cheap.
--
-- Question: Is it possible to be more conservative? This always generates a
-- perfectly balanced queue, but it's not obvious to me that we need to go
-- that far.
instance Ord a => Movable (NEQueue a) where
  move q = fromAscNonEmptyNE (toNonEmptyNE q)


mergeNEQ :: Ord a => NEQueue a %1-> NEQueue a %1-> NEQueue a
mergeNEQ (NEQ x1 ts1) (NEQ x2 ts2)
  | x1 <= x2
  = NEQ x1 (NEQ x2 ts2 `NEQCons` ts1)
  | otherwise
  = NEQ x2 (NEQ x1 ts1 `NEQCons` ts2)

merge :: Ord a => PQueue a %1-> PQueue a %1-> PQueue a
merge Empty q = q
merge q Empty = q
merge (Q ne1) (Q ne2) = Q (mergeNEQ ne1 ne2)

singletonNEQ :: a -> NEQueue a
singletonNEQ a = NEQ a Nil

singleton :: a -> PQueue a
singleton a = Q (singletonNEQ a)

popMinNEQ :: Ord a => NEQueue a %1-> (Ur a, PQueue a)
popMinNEQ (NEQ x xs) = (Ur x, mergeNEQs xs)
  where
    mergeNEQs Nil = Empty
    mergeNEQs (t `NEQCons` ts) = Q $ go t ts

    go t Nil = t
    go t1 (t2 `NEQCons` Nil) = t1 <+> t2
    go t1 (t2 `NEQCons` (t3 `NEQCons` ts)) = (t1 <+> t2) <+> go t3 ts
    (<+>) = mergeNEQ

minViewNE :: Ord a => NEQueue a %1-> (Ur a, Maybe (NEQueue a))
minViewNE ne = case popMinNEQ ne of
  (ua, Empty) -> (ua, Nothing)
  (ua, Q ne') -> (ua, Just ne')

minView :: Ord a => PQueue a %1-> Maybe (Ur a, PQueue a)
minView Empty = Nothing
minView (Q ne) = case popMinNEQ ne of
  (the_min, the_rest) -> Just (the_min, the_rest)
{-# INLINE minView #-}

insert :: Ord a => a -> PQueue a %1 -> PQueue a
insert a q = merge (singleton a) q

empty :: PQueue a
empty = Empty

fromList :: Ord a => [Ur a] %1-> PQueue a
fromList = LL.foldl' (\acc (Ur a) -> insert a acc) empty

fromNonEmptyNE :: Ord a => NonEmpty (Ur a) %1-> NEQueue a
fromNonEmptyNE (Ur a :| as) = case fromList as of
     Empty -> NEQ a Nil
     Q ne -> NEQ a (NEQCons ne Nil)

fromListNonLinear :: Ord a => [a] -> PQueue a
fromListNonLinear = F.foldl' (\acc a -> insert a acc) empty

fromNonEmptyNonLinearNE :: Ord a => NonEmpty a -> NEQueue a
fromNonEmptyNonLinearNE (a :| as) = NEQ a $
  case fromListNonLinear as of
    Empty -> Nil
    Q ne -> NEQCons ne Nil

foldrNonLinear :: forall a b. Ord a => (a -> b -> b) -> b -> PQueue a -> b
foldrNonLinear c n = \q -> go q
  where
    go :: PQueue a -> b
    go q = case minView q of
      Nothing -> n
      Just (Ur a, q') -> a `c` go q'

foldr :: forall a b. Ord a => (a -> b %1 -> b) -> b %1 -> PQueue a %1 -> b
foldr c = \n q -> go n q
  where
    go :: b %1-> PQueue a %1-> b
    go n q = case minView q of
      Nothing -> n
      Just (Ur a, q') -> a `c` go n q'

foldrNE1 :: forall a b. Ord a => (a -> b %1 -> b) -> (a -> b) %1 -> NEQueue a %1 -> b
foldrNE1 c = \n q -> go n q
  where
    go :: (a -> b) %1-> NEQueue a %1-> b
    go n q = case popMinNEQ q of
      (Ur a, Empty) -> n a
      (Ur a, Q more) -> a `c` go n more

foldl' :: forall a b. Ord a => (b %1-> a -> b) -> b %1-> PQueue a %1-> b
foldl' f b q = foldr (\a r !acc -> r (f acc a)) (\x -> x) q b

-- | Lazily convert a 'PQueue` to a list.
toList :: Ord a => PQueue a %1-> [Ur a]
toList = foldr (\a r -> Ur a : r) []

toNonEmptyNE :: Ord a => NEQueue a %1-> NonEmpty (Ur a)
toNonEmptyNE ne = case popMinNEQ ne of
  (ura, q) -> ura :| toList q

toPlainNonEmptyNE :: Ord a => NEQueue a %1-> NonEmpty a
toPlainNonEmptyNE ne = case popMinNEQ ne of
  (Ur a, q) -> a :| toPlainList q

-- | Less powerful than 'toList', but easier to handle in non-linear code.
toPlainList :: Ord a => PQueue a %1-> [a]
toPlainList = foldr (:) []

instance (Show a, Ord a) => Show (PQueue a) where
  showsPrec p q = showParen (p > 10) (showString "fromAscList " . showList (toPlainList q))
instance (Show a, Ord a) => Show (NEQueue a) where
  showsPrec p q = showParen (p > 10) (showString "fromAscNonEmptyNE " . shows (toPlainNonEmptyNE q))