Internal.hs

{-# language DeriveTraversable #-}
{-# language ScopedTypeVariables #-}
{-# language BangPatterns #-}
{-# language MagicHash #-}
{-# language UnboxedTuples #-}
{-# language PatternSynonyms #-}
{-# language ViewPatterns #-}
{-# language Trustworthy #-}
{-# language TypeFamilies #-}
{-# language FlexibleContexts #-}
{- OPTIONS_GHC -Wall #-}
{- OPTIONS_GHC -ddump-simpl #-}

{- |
Space-efficient queues with amortized \( O(\log n) \) operations.  These
directly use an underlying array-based implementation, without doing any
special optimization for the first few and last few elements of the queue.
-}

module Data.CompactSequence.Queue.Simple.Internal
  ( Queue (.., Empty, (:<))
  , (|>)
  , empty
  , snoc
  , uncons
  , take
  , fromList
  , fromListN
  , fromListNIncremental
  ) where

import qualified Data.CompactSequence.Queue.Internal as Q
import Data.CompactSequence.Internal.Size (Size, Twice)
import qualified Data.CompactSequence.Internal.Size as Sz
import qualified Data.CompactSequence.Internal.Array as A
import qualified Data.CompactSequence.Internal.Numbers as N
import qualified Data.Foldable as F
import qualified GHC.Exts as Exts
import Control.Monad.State.Strict
import qualified Control.Monad.State.Lazy as LS
import qualified Prelude as P
import Prelude hiding (take)

-- | A queue.
newtype Queue a = Queue (Q.Queue Sz.Sz1 a)
  deriving (Functor, Traversable, Eq, Ord)

-- | The empty queue.
empty :: Queue a
empty = Queue Q.empty

-- | Enqueue an element at the rear of a queue.
snoc :: Queue a -> a -> Queue a
snoc (Queue q) a = Queue $ Q.snocA Sz.one q (A.singleton a)

-- | An infix synonym for 'snoc'.
(|>) :: Queue a -> a -> Queue a
(|>) = snoc

-- | Dequeue an element from the front of a queue.
uncons :: Queue a -> Maybe (a, Queue a)
uncons (Queue q) = case Q.viewA Sz.one q of
  Q.EmptyA -> Nothing
  Q.ConsA sa q'
    | (# a #) <- A.getSingleton# sa
    -> Just (a, Queue q')

infixr 5 :<
infixl 4 `snoc`, |>

-- | A unidirectional pattern synonym for viewing the
-- front of a queue.
pattern (:<) :: a -> Queue a -> Queue a
pattern x :< xs <- (uncons -> Just (x, xs))

-- | A bidirectional pattern synonym for the empty queue.
pattern Empty :: Queue a
pattern Empty = Queue Q.Empty
{-# COMPLETE (:<), Empty #-}

instance Foldable Queue where
  -- TODO: Implement more methods?
  foldMap f (Queue q) = foldMap f q
  foldr c n (Queue q) = foldr c n q
  foldr' c n (Queue q) = F.foldr' c n q
  foldl f b (Queue q) = foldl f b q
  foldl' f b (Queue q) = F.foldl' f b q

  null (Queue Q.Empty) = True
  null _ = False
  -- Note: length only does O(log n) *unshared* work, but it does O(n) amortized
  -- work because it has to force the entire spine. We could avoid
  -- this, of course, by storing the size with the queue.
  length (Queue q) = go 0 Sz.one q
    where
      go :: Int -> Size m -> Q.Queue m a -> Int
      go !acc !_s Q.Empty = acc
      go !acc !s (Q.Node pr m sf) = go (acc + lpr + lsf) (Sz.twice s) m
        where
          lpr = case pr of
                  Q.FD1{} -> Sz.getSize s
                  Q.FD2{} -> 2*Sz.getSize s
                  Q.FD3{} -> 3*Sz.getSize s
          lsf = case sf of
                  Q.RD0 -> 0
                  Q.RD1{} -> Sz.getSize s
                  Q.RD2{} -> 2*Sz.getSize s

instance Show a => Show (Queue a) where
    showsPrec p xs = showParen (p > 10) $
        showString "fromList " . shows (F.toList xs)

instance Exts.IsList (Queue a) where
  type Item (Queue a) = a
  toList = F.toList
  fromList = fromList
  fromListN = fromListN

instance Semigroup (Queue a) where
  -- This gives us O(m + n) append. Can we do better?
  -- I suspect O(min(m,n)) might be possible.
  Empty <> q = q
  q <> Empty = q
  q <> r = fromListN (length q + length r) (F.toList q ++ F.toList r)

instance Monoid (Queue a) where
  mempty = empty

-- | Take up to the given number of elements from the front
-- of a queue to form a new queue. \( O(\min (k, n)) \), where
-- \( k \) is the integer argument and \( n \) is the size of
-- the queue.
take :: Int -> Queue a -> Queue a
take n s
  | n <= 0 = Empty
  | compareLength n s == LT
  = fromListN n (P.take n (F.toList s))
  | otherwise = s

-- | \( O(\min(m, n)) \). Compare an 'Int' to the length of a 'Queue'.
--
-- @compareLength n xs = compare n (length xs)@
compareLength :: Int -> Queue a -> Ordering
compareLength n0 (Queue que0) = go Sz.one n0 que0
  where
    go :: Size n -> Int -> Q.Queue n a -> Ordering
    go !_sz n Q.Empty = compare n 0
    go _sz n _ | n <= 0 = LT
    go sz n (Q.Node pr m sf)
      = go (Sz.twice sz) (n - frontLen sz pr - rearLen sz sf) m

frontLen :: Size n -> Q.FD n a -> Int
frontLen s Q.FD1{} = Sz.getSize s
frontLen s Q.FD2{} = 2 * Sz.getSize s
frontLen s Q.FD3{} = 3 * Sz.getSize s

rearLen :: Size n -> Q.RD n a -> Int
rearLen s Q.RD0{} = 0
rearLen s Q.RD1{} = Sz.getSize s
rearLen s Q.RD2{} = 2 * Sz.getSize s

-- | \( O(n \log n) \). Convert a list to a 'Queue', with the head of the
-- list at the front of the queue.
fromList :: [a] -> Queue a
fromList = F.foldl' snoc empty

-- | \( O(n) \). Convert a list of the given size to a 'Queue', with the
-- head of the list at the front of the queue.
fromListN :: Int -> [a] -> Queue a
fromListN n xs
  = Queue $ evalState (fromListQN Sz.one (N.toBin23 n)) xs

-- | \( O(n) \). Convert a list of the given size to a 'Queue', with the
-- head of the list at the front of the queue. Unlike 'fromListN',
-- the conversion is performed incrementally. This is generally
-- beneficial if the list is represented compactly (e.g., an enumeration)
-- or when it's otherwise not important to consume the entire list
-- immediately.
fromListNIncremental :: Int -> [a] -> Queue a
fromListNIncremental n xs
  = Queue $ LS.evalState (fromListQN Sz.one (N.toBin23 n)) xs

-- We use a similar approach to the one we use for stacks.  Every node of the
-- resulting queue will be safe, except possibly the last one. This should make
-- the resulting queue cheap to work with initially. In particular, each front
-- digit (except possibly the last) will be 2 or 3, and each rear digit will be
-- 0. This arrangement also lets us offer an incremental version of fromListN.

-- Without these SPECIALIZE pragmas, this doesn't get specialized
-- for some reason. Bleh!
{-# SPECIALIZE
  fromListQN :: Size n -> N.Bin23 -> State [a] (Q.Queue n a)
 #-}
{-# SPECIALIZE
  fromListQN :: Size n -> N.Bin23 -> LS.State [a] (Q.Queue n a)
 #-}
fromListQN :: MonadState [a] m => Size n -> N.Bin23 -> m (Q.Queue n a)
fromListQN !_ N.End23 = do
  remains <- get
  if null remains
    then pure Q.empty
    else error "Data.CompactSequence.Queue.Simple.fromListQN: List too long"
fromListQN !sz N.OneEnd23 = do
  sa <- state (A.arraySplitListN sz)
  remains <- get
  if null remains
    then pure $! Q.Node (Q.FD1 sa) Q.Empty Q.RD0
    else error "Data.CompactSequence.Queue.Simple.fromListQN: List too long"
fromListQN !sz (N.Two23 mn) = do
  sa1 <- state (A.arraySplitListN sz)
  sa2 <- state (A.arraySplitListN sz)
  m <- fromListQN (Sz.twice sz) mn
  pure $! Q.Node (Q.FD2 sa1 sa2) m Q.RD0
fromListQN !sz (N.Three23 mn) = do
  sa1 <- state (A.arraySplitListN sz)
  sa2 <- state (A.arraySplitListN sz)
  sa3 <- state (A.arraySplitListN sz)
  m <- fromListQN (Sz.twice sz) mn
  pure $! Q.Node (Q.FD3 sa1 sa2 sa3) m Q.RD0