{-# language BangPatterns #-}
{-# language PatternSynonyms #-}
{-# language ViewPatterns #-}
{-# language TypeFamilies #-}
{-# language DeriveTraversable #-}
-- We need Trustworthy for the IsList instance. *sigh*
{-# language Trustworthy #-}

{- |
Space-efficient stacks with amortized \( O(\log n) \) operations.
These directly use an underlying array-based implementation,
without doing any special optimization for the very top of the
stack.
-}

module Data.CompactSequence.Stack.Simple.Internal
  ( Stack (.., Empty, (:<))
  , empty
  , cons
  , (<|)
  , uncons
  , compareLength
  , take
  , fromList
  , fromListN
  ) where

import qualified Data.CompactSequence.Stack.Internal as S
import Data.CompactSequence.Stack.Internal (consA, unconsA, ViewA (..))
import Data.CompactSequence.Internal.Size (Size, Twice)
import qualified Data.CompactSequence.Internal.Size as Sz
import qualified Data.CompactSequence.Internal.Array.Safe as A
import qualified Data.CompactSequence.Internal.Numbers as N
import qualified Data.Foldable as F
import qualified GHC.Exts as Exts
import qualified Prelude as P
import Prelude hiding (take)

-- | A stack.
newtype Stack a = Stack {unStack :: S.Stack Sz.Sz1 a}
  deriving (Functor, Traversable, Eq, Ord)
  -- TODO: Write a custom Traversable instance to avoid
  -- an extra fmap at the top.

-- | The empty stack.
empty :: Stack a
empty = Stack S.empty

infixr 5 `cons`, :<, <|

-- | Push an element onto the front of a stack.
--
-- \( O(\log n) \)
cons :: a -> Stack a -> Stack a
cons a (Stack s) = Stack $ consA Sz.one (A.singleton a) s

-- | Pop an element off the front of a stack.
--
-- Accessing the first element is \( O(1) \). Accessing the rest is
-- \( O(\log n) \).
uncons :: Stack a -> Maybe (a, Stack a)
uncons (Stack stk) = do
  ConsA sa stk' <- pure $ unconsA Sz.one stk
  hd <- A.getSingletonA sa
  Just (hd, Stack stk')

-- | An infix synonym for 'cons'.
(<|) :: a -> Stack a -> Stack a
(<|) = cons

-- | A bidirectional pattern synonym for working with
-- the front of a stack.
pattern (:<) :: a -> Stack a -> Stack a
pattern x :< xs <- (uncons -> Just (x, xs))
  where
    (:<) = cons

-- | A bidirectional pattern synonym for the empty stack.
pattern Empty :: Stack a
pattern Empty = Stack S.Empty

{-# COMPLETE (:<), Empty #-}

instance Foldable Stack where
  -- TODO: implement more methods.
  foldMap f (Stack s) = foldMap f s
  foldr c n (Stack s) = foldr c n s
  foldl' f b (Stack s) = F.foldl' f b s
  null (Stack s) = null s

  -- length does O(log n) *unshared* work, but since
  -- it forces the spine it does O(n) *amortized* work.
  -- The right way to get stack sizes efficiently is to track
  -- them separately.
  length (Stack xs) = go 1 0 xs
    where
      go :: Int -> Int -> S.Stack m a -> Int
      go !_s acc S.Empty = acc
      go s acc (S.One _ more) = go (2*s) (acc + s) more
      go s acc (S.Two _ _ more) = go (2*s) (acc + 2*s) more
      go s acc (S.Three _ _ _ more) = go (2*s) (acc + 3*s) more

-- | \( O(\min(m, n)) \). Compare an 'Int' to the length of a 'Stack'.
--
-- @compareLength n xs = compare n (length xs)@
compareLength :: Int -> Stack a -> Ordering
compareLength n0 (Stack stk0) = go Sz.one n0 stk0
  where
    go :: Size n -> Int -> S.Stack n a -> Ordering
    go !_sz n S.Empty = compare n 0
    go _sz n _ | n <= 0 = LT
    go sz n (S.One _ more) = go (Sz.twice sz) (n - Sz.getSize sz) more
    go sz n (S.Two _ _ more) = go (Sz.twice sz) (n - 2*Sz.getSize sz) more
    go sz n (S.Three _ _ _ more) = go (Sz.twice sz) (n - 3*Sz.getSize sz) more

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

instance Semigroup (Stack a) where
  -- This gives us O(m + n) append. I believe it's possible to
  -- achieve O(m). See #12 for a sketch.
  Empty <> s = s
  s <> Empty = s
  s <> t = fromListN (length s + length t) (F.toList s ++ F.toList t)

instance Monoid (Stack a) where
  mempty = empty

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

-- | \( O(n \log n) \). Convert a list to a stack, with the
-- first element of the list as the top of the stack.
fromList :: [a] -> Stack a
fromList = foldr cons empty

-- | \( O(n) \). Convert a list of known length to a stack,
-- with the first element of the list as the top of the stack.
fromListN :: Int -> [a] -> Stack a
fromListN n !_
  | n < 0 = error "Data.CompactSequence.Stack.fromListN: Negative argument."
fromListN s xs = Stack $ fromListSN Sz.one (N.toDyadic s) xs

-- We convert the size to a dyadic representation
-- (1-2 binary) and use that as the shape of the stack.
fromListSN :: Size n -> N.Dyadic -> [a] -> S.Stack n a
fromListSN !_ N.DEnd xs
  | F.null xs = S.Empty
  | otherwise = error "Data.CompactSequence.Stack.fromListN: List too long."
fromListSN s (N.DOne n') xs
  | (ar, xs') <- A.arraySplitListN s xs
  = S.One ar (fromListSN (Sz.twice s) n' xs')
fromListSN s (N.DTwo n') xs
  | (ar1, xs') <- A.arraySplitListN s xs
  , (ar2, xs'') <- A.arraySplitListN s xs'
    -- We build eagerly to dispose of the list as soon as
    -- possible.
  = S.Two ar1 ar2 $! fromListSN (Sz.twice s) n' xs''

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