{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE BangPatterns #-}


-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Logic.Sequence.Internal.ScheduledQueue
-- Copyright   :  (c) Atze van der Ploeg 2014
--                (c) David Feuer 2021
-- License     :  BSD-style
-- Maintainer  :  David.Feuer@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-- A sequence, a queue, with worst case constant time: '|>', and 'viewl'.
--
-- Based on: "Simple and Efficient Purely Functional Queues and Deques", Chris Okasaki,
-- Journal of Functional Programming 1995
--
-----------------------------------------------------------------------------

module Control.Monad.Logic.Sequence.Internal.ScheduledQueue
  ( Queue
  ) where
import Data.SequenceClass (Sequence, ViewL (..))
import qualified Data.SequenceClass as S
import Data.Foldable
import qualified Data.Traversable as T
import qualified Control.Applicative as A

#if !MIN_VERSION_base(4,8,0)
import Data.Functor (Functor (..))
import Data.Monoid (Monoid (..))
#endif

infixl 5 :>
-- | A strict-spined snoc-list
data SL a
  = SNil
  | !(SL a) :> a
  deriving forall a b. a -> SL b -> SL a
forall a b. (a -> b) -> SL a -> SL b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> SL b -> SL a
$c<$ :: forall a b. a -> SL b -> SL a
fmap :: forall a b. (a -> b) -> SL a -> SL b
$cfmap :: forall a b. (a -> b) -> SL a -> SL b
Functor

-- | Append a snoc list to a list.
--
-- Precondition: |f| = |r| - 1
appendSL :: [a] -> SL a -> [a]
appendSL :: forall a. [a] -> SL a -> [a]
appendSL [a]
f SL a
r = forall a. [a] -> SL a -> [a] -> [a]
rotate [a]
f SL a
r []

-- Precondition:
-- |f| = |r| - 1
rotate :: [a] -> SL a -> [a] -> [a]
rotate :: forall a. [a] -> SL a -> [a] -> [a]
rotate [] (SL a
_SNil :> a
y) [a]
a = a
y forall a. a -> [a] -> [a]
: [a]
a
rotate (a
x : [a]
f) (SL a
r :> a
y) [a]
a = a
x forall a. a -> [a] -> [a]
: forall a. [a] -> SL a -> [a] -> [a]
rotate [a]
f SL a
r (a
y forall a. a -> [a] -> [a]
: [a]
a)
rotate [a]
_f SL a
_a [a]
_r  = forall a. HasCallStack => [Char] -> a
error [Char]
"Invariant |f| = |r| + |a| - 1 broken"

-- | A scheduled banker's queue, as described by Okasaki. In theory, we only
-- need a queue supporting constant /amortized/ time operations. In practice,
-- once a queue gets large, linear-time pauses and cache effects relating to
-- rebuilding start to hurt.
data Queue a =
  RQ ![a]    -- front (f)
     !(SL a) -- rear (r)
     ![a]  -- schedule (a)
-- Invariant: |f| = |r| + |a|
  deriving forall a b. a -> Queue b -> Queue a
forall a b. (a -> b) -> Queue a -> Queue b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Queue b -> Queue a
$c<$ :: forall a b. a -> Queue b -> Queue a
fmap :: forall a b. (a -> b) -> Queue a -> Queue b
$cfmap :: forall a b. (a -> b) -> Queue a -> Queue b
Functor
  -- We would much rather write
  --
  --   data Queue a = forall x. RQ ![a] !(SL a) ![x]
  --
  -- to guarantee we don't accidentally look at elements in the schedule.
  -- Unfortuately, GHC can't currently unpack types with existentials, and
  -- we want to unpack into the catenable queue constructor. We used to use
  -- [Any], but the modern unsafeCoerce makes that produce rather messy core,
  -- and I'm a bit concerned about the term sizes for inlining and such.

queue :: [a] -> SL a -> [a] -> Queue a
-- precondition : |f| = |r| + |a| - 1
-- postcondition: |f| = |r| + |a|
queue :: forall a. [a] -> SL a -> [a] -> Queue a
queue [a]
f SL a
r [] =
  let
    f' :: [a]
f' = forall a. [a] -> SL a -> [a]
appendSL [a]
f SL a
r
    -- We NOINLINE f' to make sure that walking the schedule actually forces
    -- the front of the queue. GHC probably won't duplicate appendSL anyway,
    -- but let's be sure.
    {-# NOINLINE f' #-}
  in forall a. [a] -> SL a -> [a] -> Queue a
RQ [a]
f' forall a. SL a
SNil [a]
f'
queue [a]
f SL a
r (a
_h : [a]
t) = forall a. [a] -> SL a -> [a] -> Queue a
RQ [a]
f SL a
r [a]
t

instance Sequence Queue where
  empty :: forall c. Queue c
empty = forall a. [a] -> SL a -> [a] -> Queue a
RQ [] forall a. SL a
SNil []
  singleton :: forall c. c -> Queue c
singleton c
x =
    let
      c :: [c]
c = [c
x]
    in forall a. [a] -> SL a -> [a] -> Queue a
RQ [c]
c forall a. SL a
SNil [c]
c
  -- The special case for [] gives us better optimizations
  -- for singleton catenable queues.
  RQ [] SL c
_ [c]
_ |> :: forall c. Queue c -> c -> Queue c
|> c
x = forall (s :: * -> *) c. Sequence s => c -> s c
S.singleton c
x
  RQ [c]
f SL c
r [c]
a |> c
x = forall a. [a] -> SL a -> [a] -> Queue a
queue [c]
f (SL c
r forall a. SL a -> a -> SL a
:> c
x) [c]
a

  viewl :: forall c. Queue c -> ViewL Queue c
viewl (RQ [] SL c
_SNil [c]
_nil) = forall (s :: * -> *) c. ViewL s c
EmptyL
  viewl (RQ (c
h : [c]
t) SL c
f [c]
a) = c
h forall c (s :: * -> *). c -> s c -> ViewL s c
:< forall a. [a] -> SL a -> [a] -> Queue a
queue [c]
t SL c
f [c]
a

instance Foldable Queue where
  foldr :: forall a b. (a -> b -> b) -> b -> Queue a -> b
foldr a -> b -> b
c b
n = \Queue a
q -> forall {s :: * -> *}. Sequence s => s a -> b
go Queue a
q
    where
      go :: s a -> b
go s a
q = case forall (s :: * -> *) c. Sequence s => s c -> ViewL s c
S.viewl s a
q of
        ViewL s a
EmptyL -> b
n
        a
h :< s a
t -> a -> b -> b
c a
h (s a -> b
go s a
t)
  foldl' :: forall b a. (b -> a -> b) -> b -> Queue a -> b
foldl' b -> a -> b
f b
b0 = \Queue a
q -> forall {s :: * -> *}. Sequence s => s a -> b -> b
go Queue a
q b
b0
    where
      go :: s a -> b -> b
go s a
q !b
b = case forall (s :: * -> *) c. Sequence s => s c -> ViewL s c
S.viewl s a
q of
        ViewL s a
EmptyL -> b
b
        a
h :< s a
t -> s a -> b -> b
go s a
t (b -> a -> b
f b
b a
h)
#if MIN_VERSION_base(4,8,0)
  null :: forall a. Queue a -> Bool
null (RQ [] SL a
_SNil [a]
_nil) = Bool
True
  null Queue a
_ = Bool
False

  -- Invariant: |f| = |r| + |a|. The length of the queue is
  -- |f| + |r|
  -- We can calculate that as either 2 * |r| + |a|
  -- or 2 * |f| - a. I suspect the latter will give better
  -- cache utilization.
  length :: forall a. Queue a -> Int
length (RQ [a]
f SL a
_ [a]
a) = Int
2 forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
f forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
a
#endif

instance T.Traversable Queue where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Queue a -> f (Queue b)
traverse a -> f b
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall c. [c] -> Queue c
fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {s :: * -> *}. Sequence s => s a -> f [b]
go
    where
      go :: s a -> f [b]
go s a
q = case forall (s :: * -> *) c. Sequence s => s c -> ViewL s c
S.viewl s a
q of
        ViewL s a
EmptyL -> forall (f :: * -> *) a. Applicative f => a -> f a
A.pure []
        a
h :< s a
t -> forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
A.liftA2 (:) (a -> f b
f a
h) (s a -> f [b]
go s a
t)

fromList :: [a] -> Queue a
fromList :: forall c. [c] -> Queue c
fromList [a]
f = forall a. [a] -> SL a -> [a] -> Queue a
RQ [a]
f forall a. SL a
SNil [a]
f