{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Implementation of pairing heaps stored off-heap
module Foreign.Heap where

import qualified Data.List as List
import Foreign.List (List)
import qualified Foreign.List as List
import Foreign.Marshal.Pure (Box, Pool)
import qualified Foreign.Marshal.Pure as Manual
import Prelude.Linear hiding (foldl)

data Heap k a
  = Empty
  | NonEmpty (Box (NEHeap k a))

data NEHeap k a
  = Heap k a (Box (List (NEHeap k a)))

instance
  (Manual.Representable k, Manual.Representable a) =>
  Manual.MkRepresentable (NEHeap k a) (k, a, Box (List (NEHeap k a)))
  where
  toRepr (Heap k a l) = (k, a, l)
  ofRepr (k, a, l) = Heap k a l

instance (Manual.Representable k, Manual.Representable a) => Manual.Representable (NEHeap k a) where
  type AsKnown (NEHeap k a) = Manual.AsKnown (k, a, (Box (List (NEHeap k a))))

-- * Non-empty heap primitives

singletonN :: (Manual.Representable k, Manual.Representable a) => k %1 -> a %1 -> Pool %1 -> NEHeap k a
singletonN k a pool = Heap k a (Manual.alloc List.Nil pool)

-- XXX: (Movable k, Ord k) is a bit stronger than strictly required. We could
-- give a linear version of `Ord` instead.
mergeN :: forall k a. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => NEHeap k a %1 -> NEHeap k a %1 -> Pool %1 -> NEHeap k a
mergeN (Heap k1 a1 h1) (Heap k2 a2 h2) pool =
  testAndRebuild (move k1) a1 h1 (move k2) a2 h2 pool
  where
    --- XXX: this is a good example of why we need a working `case` and/or
    --- `let`.
    testAndRebuild :: Ur k %1 -> a %1 -> Box (List (NEHeap k a)) %1 -> Ur k %1 -> a %1 -> Box (List (NEHeap k a)) %1 -> Pool %1 -> NEHeap k a
    testAndRebuild (Ur k1') a1' h1' (Ur k2') a2' h2' =
      if k1' <= k2'
        then helper k1' a1' k2' a2' h1' h2'
        else helper k2' a2' k1' a1' h2' h1'

    helper :: k -> a %1 -> k -> a %1 -> Box (List (NEHeap k a)) %1 -> Box (List (NEHeap k a)) %1 -> Pool %1 -> NEHeap k a
    helper k1'' a1'' k2'' a2'' h1'' h2'' pool'' = Heap k1'' a1'' (Manual.alloc ((List.Cons :: b %1 -> Box (List b) %1 -> List b) ((Heap :: c %1 -> b %1 -> Box (List (NEHeap c b)) %1 -> NEHeap c b) k2'' a2'' h2'') h1'') pool'')

-- XXX: the type signatures for List.Cons and Heap are necessary for certain
-- older versions of the compiler, and as such are temporary. See PR #38
-- and PR #380 in tweag/ghc/linear-types.

mergeN' :: forall k a. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => NEHeap k a %1 -> Heap k a %1 -> Pool %1 -> NEHeap k a
mergeN' h Empty pool = pool `lseq` h
mergeN' h (NonEmpty h') pool = mergeN h (Manual.deconstruct h') pool

extractMinN :: (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => NEHeap k a %1 -> Pool %1 -> (k, a, Heap k a)
extractMinN (Heap k a h) pool = (k, a, pairUp (Manual.deconstruct h) pool)

pairUp :: forall k a. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => List (NEHeap k a) %1 -> Pool %1 -> Heap k a
pairUp List.Nil pool = pool `lseq` Empty
pairUp (List.Cons h r) pool = pairOne h (Manual.deconstruct r) (dup pool)
  where
    pairOne :: NEHeap k a %1 -> List (NEHeap k a) %1 -> (Pool, Pool) %1 -> Heap k a
    pairOne h' r' (pool1, pool2) =
      NonEmpty $ Manual.alloc (pairOne' h' r' (dup3 pool1)) pool2

    pairOne' :: NEHeap k a %1 -> List (NEHeap k a) %1 -> (Pool, Pool, Pool) %1 -> NEHeap k a
    pairOne' h1 List.Nil pools =
      pools `lseq` h1
    pairOne' h1 (List.Cons h2 r') (pool1, pool2, pool3) =
      mergeN' (mergeN h1 h2 pool1) (pairUp (Manual.deconstruct r') pool2) pool3

-- * Heap primitives

empty :: Heap k a
empty = Empty

singleton :: forall k a. (Manual.Representable k, Manual.Representable a) => k %1 -> a %1 -> Pool %1 -> Heap k a
singleton k a pool = NonEmpty $ singletonAlloc k a (dup pool)
  where
    singletonAlloc :: k %1 -> a %1 -> (Pool, Pool) %1 -> Box (NEHeap k a)
    singletonAlloc k' a' (pool1, pool2) =
      Manual.alloc (singletonN k' a' pool1) pool2

extractMin :: (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => Heap k a %1 -> Pool %1 -> Maybe (k, a, Heap k a)
extractMin Empty pool = pool `lseq` Nothing
extractMin (NonEmpty h) pool = Just $ extractMinN (Manual.deconstruct h) pool

merge :: forall k a. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => Heap k a %1 -> Heap k a %1 -> Pool %1 -> Heap k a
merge Empty h' pool = pool `lseq` h'
merge (NonEmpty h) h' pool = NonEmpty $ neMerge (Manual.deconstruct h) h' (dup pool)
  where
    neMerge :: NEHeap k a %1 -> Heap k a %1 -> (Pool, Pool) %1 -> Box (NEHeap k a)
    neMerge h1 h2 (pool1, pool2) =
      Manual.alloc (mergeN' h1 h2 pool1) pool2

-- * Heap sort

-- | Guaranteed to yield pairs in ascending key order
foldl :: forall k a b. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => (b %1 -> k %1 -> a %1 -> b) -> b %1 -> Heap k a %1 -> Pool %1 -> b
foldl f acc h pool = go acc h (dup pool)
  where
    go :: b %1 -> Heap k a %1 -> (Pool, Pool) %1 -> b
    go acc' h' (pool1, pool2) = dispatch acc' (extractMin h' pool1) pool2

    dispatch :: b %1 -> Maybe (k, a, Heap k a) %1 -> Pool %1 -> b
    dispatch acc' Nothing pool' = pool' `lseq` acc'
    dispatch acc' (Just (k, a, h')) pool' =
      foldl f (f acc' k a) h' pool'

-- | Strict: stream must terminate.
unfold :: forall k a s. (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => (s -> Maybe ((k, a), s)) -> s -> Pool %1 -> Heap k a
unfold step seed pool = dispatch (step seed) pool
  where
    dispatch :: (Maybe ((k, a), s)) -> Pool %1 -> Heap k a
    dispatch Nothing pool' = pool' `lseq` Empty
    dispatch (Just ((k, a), next)) pool' = mkStep k a next (dup3 pool')

    mkStep :: k -> a -> s -> (Pool, Pool, Pool) %1 -> Heap k a
    mkStep k a next (pool1, pool2, pool3) =
      merge (singleton k a pool1) (unfold step next pool2) pool3

-- TODO: linear unfold: could apply to off-heap lists!

ofList :: (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => [(k, a)] -> Pool %1 -> Heap k a
ofList l pool = unfold List.uncons l pool

-- XXX: sorts in reverse
toList :: (Manual.Representable k, Manual.Representable a, Movable k, Ord k) => Heap k a %1 -> Pool %1 -> [(k, a)]
toList h pool = foldl (\l k a -> (k, a) : l) [] h pool

sort :: forall k a. (Manual.Representable k, Manual.Representable a, Movable k, Ord k, Movable a) => [(k, a)] -> [(k, a)]
sort l = unur $ Manual.withPool (\pool -> move $ sort' l (dup pool))
  where
    -- XXX: can we avoid this call to `move`?

    sort' :: [(k, a)] -> (Pool, Pool) %1 -> [(k, a)]
    sort' l' (pool1, pool2) = toList (ofList l' pool1) pool2