{-# 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)