{-# LANGUAGE GADTs #-}

module Test.Credit.Sortable.MergeSort where

import Prettyprinter (Pretty)
import Control.Monad.Credit
import Test.Credit
import Test.Credit.Sortable.Base

data MergeSort a m = MergeSort Size (Thunk m (MLazyCon m) [[a]])

mrg :: MonadCredit m => Ord a => [a] -> [a] -> m [a]
mrg :: forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [a] -> m [a]
mrg [] [a]
ys = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
ys
mrg [a]
xs [] = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
xs
mrg (a
x:[a]
xs) (a
y:[a]
ys) = m ()
forall (m :: * -> *). MonadCount m => m ()
tick m () -> m [a] -> m [a]
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
  if a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
y
    then (a
x:) ([a] -> [a]) -> m [a] -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [a] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [a] -> m [a]
mrg [a]
xs (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys)
    else (a
y:) ([a] -> [a]) -> m [a] -> m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> [a] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [a] -> m [a]
mrg (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs) [a]
ys

addSeg :: MonadCredit m => Ord a => [a] -> [[a]] -> Size -> m [[a]]
addSeg :: forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> Size -> m [[a]]
addSeg [a]
seg [[a]]
segs Size
size =
  if Size
size Size -> Size -> Size
forall a. Integral a => a -> a -> a
`mod` Size
2 Size -> Size -> Bool
forall a. Eq a => a -> a -> Bool
== Size
0 then [[a]] -> m [[a]]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([[a]] -> m [[a]]) -> [[a]] -> m [[a]]
forall a b. (a -> b) -> a -> b
$ [a]
seg [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [[a]]
segs
  else do -- technically we should have a tick here, but Okasaki doesn't and we follow his presentation
    let ([a]
seg1:[[a]]
segs') = [[a]]
segs
    [a]
seg' <- [a] -> [a] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [a] -> m [a]
mrg [a]
seg [a]
seg1 
    [a] -> [[a]] -> Size -> m [[a]]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> Size -> m [[a]]
addSeg [a]
seg' [[a]]
segs' (Size
size Size -> Size -> Size
forall a. Integral a => a -> a -> a
`div` Size
2)

psi :: Size -> Credit
psi :: Size -> Credit
psi Size
n = Credit
2 Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
* Size -> Credit
linear Size
n Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
- Credit
2 Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
* [Credit] -> Credit
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Size -> Credit
linear (Size -> Credit) -> Size -> Credit
forall a b. (a -> b) -> a -> b
$ Credit -> Size
b Credit
i Size -> Size -> Size
forall a. Num a => a -> a -> a
* (Size
n Size -> Size -> Size
forall a. Integral a => a -> a -> a
`mod` Size
2Size -> Credit -> Size
forall a b. (Num a, Integral b) => a -> b -> a
^Credit
i Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
1) | Credit
i <- [Credit
0..(Size -> Credit
log2 Size
n Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
+ Credit
1)] ]
  where
    b :: Credit -> Size
b Credit
i = (Size
n Size -> Size -> Size
forall a. Integral a => a -> a -> a
`div` Size
2Size -> Credit -> Size
forall a b. (Num a, Integral b) => a -> b -> a
^Credit
i) Size -> Size -> Size
forall a. Integral a => a -> a -> a
`mod` Size
2

data MLazyCon m a where
  Empty :: MLazyCon m [[a]]
  AddSeg :: Ord a => [a] -> Thunk m (MLazyCon m) [[a]] -> Size -> MLazyCon m [[a]]

instance MonadCredit m => HasStep (MLazyCon m) m where
  step :: forall a. MLazyCon m a -> m a
step MLazyCon m a
Empty = a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  step (AddSeg [a]
seg Thunk m (MLazyCon m) [[a]]
segs Size
size) = do
    Thunk m (MLazyCon m) [[a]] -> Credit -> m ()
forall (m :: * -> *) (t :: * -> *) a.
MonadCredit m =>
Thunk m t a -> Credit -> m ()
forall (t :: * -> *) a. Thunk m t a -> Credit -> m ()
creditWith Thunk m (MLazyCon m) [[a]]
segs (Size -> Credit
psi Size
size)
    [[a]]
segs' <- Thunk m (MLazyCon m) [[a]] -> m [[a]]
forall (m :: * -> *) (t :: * -> *) a.
(MonadLazy m, HasStep t m) =>
Thunk m t a -> m a
forall (t :: * -> *) a. HasStep t m => Thunk m t a -> m a
force Thunk m (MLazyCon m) [[a]]
segs
    [a] -> [[a]] -> Size -> m [[a]]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> Size -> m [[a]]
addSeg [a]
seg [[a]]
segs' Size
size

mrgAll :: MonadCredit m => Ord a => [a] -> [[a]] -> m [a]
mrgAll :: forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> m [a]
mrgAll [a]
xs [] = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a]
xs
mrgAll [a]
xs ([a]
seg:[[a]]
segs) = m ()
forall (m :: * -> *). MonadCount m => m ()
tick m () -> m [a] -> m [a]
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> do
  [a]
seg' <- [a] -> [a] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [a] -> m [a]
mrg [a]
xs [a]
seg
  [a] -> [[a]] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> m [a]
mrgAll [a]
seg' [[a]]
segs

instance Sortable MergeSort where
  empty :: forall (m :: * -> *) a. MonadCredit m => m (MergeSort a m)
empty = do
    Thunk m (MLazyCon m) [[a]]
segs <- MLazyCon m [[a]] -> m (Thunk m (MLazyCon m) [[a]])
forall (m :: * -> *) (t :: * -> *) a.
MonadLazy m =>
t a -> m (Thunk m t a)
forall (t :: * -> *) a. t a -> m (Thunk m t a)
delay MLazyCon m [[a]]
forall (m :: * -> *) a. MLazyCon m [[a]]
Empty
    MergeSort a m -> m (MergeSort a m)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MergeSort a m -> m (MergeSort a m))
-> MergeSort a m -> m (MergeSort a m)
forall a b. (a -> b) -> a -> b
$ Size -> Thunk m (MLazyCon m) [[a]] -> MergeSort a m
forall a (m :: * -> *).
Size -> Thunk m (MLazyCon m) [[a]] -> MergeSort a m
MergeSort Size
0 Thunk m (MLazyCon m) [[a]]
segs
  add :: forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
a -> MergeSort a m -> m (MergeSort a m)
add a
x (MergeSort Size
size Thunk m (MLazyCon m) [[a]]
segs) = do
    Thunk m (MLazyCon m) [[a]]
segs' <- MLazyCon m [[a]] -> m (Thunk m (MLazyCon m) [[a]])
forall (m :: * -> *) (t :: * -> *) a.
MonadLazy m =>
t a -> m (Thunk m t a)
forall (t :: * -> *) a. t a -> m (Thunk m t a)
delay (MLazyCon m [[a]] -> m (Thunk m (MLazyCon m) [[a]]))
-> MLazyCon m [[a]] -> m (Thunk m (MLazyCon m) [[a]])
forall a b. (a -> b) -> a -> b
$ [a] -> Thunk m (MLazyCon m) [[a]] -> Size -> MLazyCon m [[a]]
forall a (m :: * -> *).
Ord a =>
[a] -> Thunk m (MLazyCon m) [[a]] -> Size -> MLazyCon m [[a]]
AddSeg [a
x] Thunk m (MLazyCon m) [[a]]
segs Size
size
    Thunk m (MLazyCon m) [[a]] -> Credit -> m ()
forall (m :: * -> *) (t :: * -> *) a.
MonadCredit m =>
Thunk m t a -> Credit -> m ()
forall (t :: * -> *) a. Thunk m t a -> Credit -> m ()
creditWith Thunk m (MLazyCon m) [[a]]
segs' (Credit
2 Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
* Size -> Credit
log2 Size
size Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
+ Credit
1)
    MergeSort a m -> m (MergeSort a m)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MergeSort a m -> m (MergeSort a m))
-> MergeSort a m -> m (MergeSort a m)
forall a b. (a -> b) -> a -> b
$ Size -> Thunk m (MLazyCon m) [[a]] -> MergeSort a m
forall a (m :: * -> *).
Size -> Thunk m (MLazyCon m) [[a]] -> MergeSort a m
MergeSort (Size
size Size -> Size -> Size
forall a. Num a => a -> a -> a
+ Size
1) Thunk m (MLazyCon m) [[a]]
segs'
  sort :: forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
MergeSort a m -> m [a]
sort (MergeSort Size
size Thunk m (MLazyCon m) [[a]]
segs) = do
    Thunk m (MLazyCon m) [[a]] -> Credit -> m ()
forall (m :: * -> *) (t :: * -> *) a.
MonadCredit m =>
Thunk m t a -> Credit -> m ()
forall (t :: * -> *) a. Thunk m t a -> Credit -> m ()
creditWith Thunk m (MLazyCon m) [[a]]
segs (Size -> Credit
psi Size
size)
    [[a]]
segs' <- Thunk m (MLazyCon m) [[a]] -> m [[a]]
forall (m :: * -> *) (t :: * -> *) a.
(MonadLazy m, HasStep t m) =>
Thunk m t a -> m a
forall (t :: * -> *) a. HasStep t m => Thunk m t a -> m a
force Thunk m (MLazyCon m) [[a]]
segs
    [a] -> [[a]] -> m [a]
forall (m :: * -> *) a.
(MonadCredit m, Ord a) =>
[a] -> [[a]] -> m [a]
mrgAll [] [[a]]
segs'

instance BoundedSortable MergeSort where
  scost :: forall a. Size -> SortableOp a -> Credit
scost Size
n (Add a
_) = Credit
2 Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
* Size -> Credit
log2 Size
n Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
+ Credit
1
  scost Size
n SortableOp a
Sort = Size -> Credit
psi Size
n Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
+  Credit
2 Credit -> Credit -> Credit
forall a. Num a => a -> a -> a
* Size -> Credit
linear Size
n

instance (MonadMemory m, MemoryCell m a) => MemoryCell m (MLazyCon m a) where
  prettyCell :: MLazyCon m a -> m Memory
prettyCell MLazyCon m a
Empty = Memory -> m Memory
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Memory -> m Memory) -> Memory -> m Memory
forall a b. (a -> b) -> a -> b
$ String -> [Memory] -> Memory
mkMCell String
"Empty" []
  prettyCell (AddSeg [a]
seg Thunk m (MLazyCon m) [[a]]
segs Size
size) = do
    -- seg' <- prettyCell seg
    Memory
segs' <- Thunk m (MLazyCon m) [[a]] -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Thunk m (MLazyCon m) [[a]]
segs
    Memory
size' <- Size -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Size
size
    Memory -> m Memory
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Memory -> m Memory) -> Memory -> m Memory
forall a b. (a -> b) -> a -> b
$ String -> [Memory] -> Memory
mkMCell String
"AddSeg" [Memory
segs', Memory
size']

instance (MonadMemory m, MemoryCell m a) => MemoryCell m (MergeSort a m) where
  prettyCell :: MergeSort a m -> m Memory
prettyCell (MergeSort Size
size Thunk m (MLazyCon m) [[a]]
segs) = do
    Memory
size' <- Size -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Size
size
    Memory
segs' <- Thunk m (MLazyCon m) [[a]] -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Thunk m (MLazyCon m) [[a]]
segs
    Memory -> m Memory
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Memory -> m Memory) -> Memory -> m Memory
forall a b. (a -> b) -> a -> b
$ String -> [Memory] -> Memory
mkMCell String
"MergeSort" [Memory
size', Memory
segs']

instance Pretty a => MemoryStructure (MergeSort (PrettyCell a)) where
  prettyStructure :: forall (m :: * -> *).
MonadMemory m =>
MergeSort (PrettyCell a) m -> m Memory
prettyStructure = MergeSort (PrettyCell a) m -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell