module Acc
  ( Acc,
    fromReverseList,
    cons,
    snoc,
    uncons,
    unsnoc,
    toNonEmpty,
    toNeAcc,
    enumFromTo,
  )
where

import qualified Acc.NeAcc as NeAcc
import qualified Acc.NeAcc.Def as NeAcc
import Acc.Prelude hiding (enumFromTo, toNonEmpty, unsnoc)
import qualified Data.Semigroup.Foldable as Foldable1

-- |
-- Data structure intended for accumulating a sequence of elements
-- for later traversal or folding.
-- Useful for implementing all kinds of builders on top.
--
-- Appending and prepending is always \(\mathcal{O}(1)\).
--
-- Another way to think about this data-structure
-- is as of a strict list with fast append and snoc.
--
-- To produce a single element 'Acc' use 'pure'.
-- To produce a multielement 'Acc' use 'fromList'.
-- To combine use '<|>' or '<>' and other 'Alternative' and 'Monoid'-related utils.
-- To extract elements use 'Foldable' API.
--
-- The benchmarks show that for the described use-case this data-structure
-- is on average 2 times faster than 'Data.DList.DList' and 'Data.Sequence.Seq',
-- is on par with list when you always prepend elements and
-- is exponentially faster than list when you append.
--
-- Internally it is implemented as a simple binary tree
-- with all functions optimized to use tail recursion,
-- ensuring that you don\'t get stack overflow.
data Acc a
  = EmptyAcc
  | TreeAcc !(NeAcc.NeAcc a)

instance (NFData a) => NFData (Acc a) where
  rnf :: Acc a -> ()
rnf = \case
    TreeAcc NeAcc a
tree -> forall a. NFData a => a -> ()
rnf NeAcc a
tree
    Acc a
EmptyAcc -> ()

instance NFData1 Acc where
  liftRnf :: forall a. (a -> ()) -> Acc a -> ()
liftRnf a -> ()
rnfLeaf = \case
    TreeAcc NeAcc a
tree -> forall (f :: * -> *) a. NFData1 f => (a -> ()) -> f a -> ()
liftRnf a -> ()
rnfLeaf NeAcc a
tree
    Acc a
EmptyAcc -> ()

deriving instance Functor Acc

instance Foldable Acc where
  {-# INLINE [0] foldMap #-}
  foldMap :: forall m a. Monoid m => (a -> m) -> Acc a -> m
foldMap a -> m
f =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f NeAcc a
a
      Acc a
EmptyAcc ->
        forall a. Monoid a => a
mempty
  {-# INLINE [0] foldMap' #-}
  foldMap' :: forall m a. Monoid m => (a -> m) -> Acc a -> m
foldMap' a -> m
f =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap' a -> m
f NeAcc a
a
      Acc a
EmptyAcc ->
        forall a. Monoid a => a
mempty
  {-# INLINE [0] foldr #-}
  foldr :: forall a b. (a -> b -> b) -> b -> Acc a -> b
foldr a -> b -> b
step b
acc =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
step b
acc NeAcc a
a
      Acc a
EmptyAcc ->
        b
acc
  {-# INLINE [0] foldr' #-}
  foldr' :: forall a b. (a -> b -> b) -> b -> Acc a -> b
foldr' a -> b -> b
step b
acc =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' a -> b -> b
step b
acc NeAcc a
a
      Acc a
EmptyAcc ->
        b
acc
  {-# INLINE [0] foldl #-}
  foldl :: forall b a. (b -> a -> b) -> b -> Acc a -> b
foldl b -> a -> b
step b
acc =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> a -> b
step b
acc NeAcc a
a
      Acc a
EmptyAcc ->
        b
acc
  {-# INLINE [0] foldl' #-}
  foldl' :: forall b a. (b -> a -> b) -> b -> Acc a -> b
foldl' b -> a -> b
step b
acc =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
step b
acc NeAcc a
a
      Acc a
EmptyAcc ->
        b
acc
  {-# INLINE [0] sum #-}
  sum :: forall a. Num a => Acc a -> a
sum =
    forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Num a => a -> a -> a
(+) a
0

instance Traversable Acc where
  {-# INLINE [0] traverse #-}
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Acc a -> f (Acc b)
traverse a -> f b
f =
    \case
      TreeAcc NeAcc a
a ->
        forall a. NeAcc a -> Acc a
TreeAcc forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse a -> f b
f NeAcc a
a
      Acc a
EmptyAcc ->
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Acc a
EmptyAcc

instance Applicative Acc where
  {-# INLINE [1] pure #-}
  pure :: forall a. a -> Acc a
pure =
    forall a. NeAcc a -> Acc a
TreeAcc forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. a -> NeAcc a
NeAcc.Leaf
  {-# INLINE [1] (<*>) #-}
  <*> :: forall a b. Acc (a -> b) -> Acc a -> Acc b
(<*>) =
    \case
      TreeAcc NeAcc (a -> b)
a ->
        \case
          TreeAcc NeAcc a
b ->
            forall a. NeAcc a -> Acc a
TreeAcc (NeAcc (a -> b)
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> NeAcc a
b)
          Acc a
EmptyAcc ->
            forall a. Acc a
EmptyAcc
      Acc (a -> b)
EmptyAcc ->
        forall a b. a -> b -> a
const forall a. Acc a
EmptyAcc

instance Alternative Acc where
  {-# INLINE [1] empty #-}
  empty :: forall a. Acc a
empty =
    forall a. Acc a
EmptyAcc
  {-# INLINE [1] (<|>) #-}
  <|> :: forall a. Acc a -> Acc a -> Acc a
(<|>) =
    \case
      TreeAcc NeAcc a
a ->
        \case
          TreeAcc NeAcc a
b ->
            forall a. NeAcc a -> Acc a
TreeAcc (forall a. NeAcc a -> NeAcc a -> NeAcc a
NeAcc.Branch NeAcc a
a NeAcc a
b)
          Acc a
EmptyAcc ->
            forall a. NeAcc a -> Acc a
TreeAcc NeAcc a
a
      Acc a
EmptyAcc ->
        forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id

instance Semigroup (Acc a) where
  {-# INLINE [1] (<>) #-}
  <> :: Acc a -> Acc a -> Acc a
(<>) =
    forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
(<|>)

instance Monoid (Acc a) where
  {-# INLINE [1] mempty #-}
  mempty :: Acc a
mempty =
    forall (f :: * -> *) a. Alternative f => f a
empty

instance IsList (Acc a) where
  type Item (Acc a) = a
  {-# INLINE [0] fromList #-}
  fromList :: [Item (Acc a)] -> Acc a
fromList = forall a. [a] -> Acc a
fromReverseList forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. [a] -> [a]
reverse
  {-# INLINE [0] toList #-}
  toList :: Acc a -> [Item (Acc a)]
toList =
    \case
      TreeAcc NeAcc a
a ->
        forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (:) [] NeAcc a
a
      Acc a
_ ->
        []

instance (Show a) => Show (Acc a) where
  show :: Acc a -> String
show =
    forall a. Show a => a -> String
show forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall l. IsList l => l -> [Item l]
toList

-- |
-- Prepend an element.
{-# INLINE [1] cons #-}
cons :: a -> Acc a -> Acc a
cons :: forall a. a -> Acc a -> Acc a
cons a
a =
  \case
    TreeAcc NeAcc a
tree ->
      forall a. NeAcc a -> Acc a
TreeAcc (forall a. NeAcc a -> NeAcc a -> NeAcc a
NeAcc.Branch (forall a. a -> NeAcc a
NeAcc.Leaf a
a) NeAcc a
tree)
    Acc a
EmptyAcc ->
      forall a. NeAcc a -> Acc a
TreeAcc (forall a. a -> NeAcc a
NeAcc.Leaf a
a)

-- |
-- Extract the first element.
--
-- The produced accumulator will lack the extracted element
-- and will have the underlying tree rebalanced towards the beginning.
-- This means that calling 'uncons' on it will be \(\mathcal{O}(1)\).
{-# INLINE uncons #-}
uncons :: Acc a -> Maybe (a, Acc a)
uncons :: forall a. Acc a -> Maybe (a, Acc a)
uncons =
  \case
    TreeAcc NeAcc a
tree ->
      case NeAcc a
tree of
        NeAcc.Branch NeAcc a
l NeAcc a
r ->
          case forall a. NeAcc a -> NeAcc a -> (a, NeAcc a)
NeAcc.unconsTo NeAcc a
r NeAcc a
l of
            (a
res, NeAcc a
newTree) ->
              forall a. a -> Maybe a
Just (a
res, forall a. NeAcc a -> Acc a
TreeAcc NeAcc a
newTree)
        NeAcc.Leaf a
res ->
          forall a. a -> Maybe a
Just (a
res, forall a. Acc a
EmptyAcc)
    Acc a
EmptyAcc ->
      forall a. Maybe a
Nothing

-- |
-- Append an element.
{-# INLINE [1] snoc #-}
snoc :: a -> Acc a -> Acc a
snoc :: forall a. a -> Acc a -> Acc a
snoc a
a =
  \case
    TreeAcc NeAcc a
tree ->
      forall a. NeAcc a -> Acc a
TreeAcc (forall a. NeAcc a -> NeAcc a -> NeAcc a
NeAcc.Branch NeAcc a
tree (forall a. a -> NeAcc a
NeAcc.Leaf a
a))
    Acc a
EmptyAcc ->
      forall a. NeAcc a -> Acc a
TreeAcc (forall a. a -> NeAcc a
NeAcc.Leaf a
a)

-- |
-- Extract the last element.
--
-- The produced accumulator will lack the extracted element
-- and will have the underlying tree rebalanced towards the end.
-- This means that calling 'unsnoc' on it will be \(\mathcal{O}(1)\) and
-- 'uncons' will be \(\mathcal{O}(n)\).
{-# INLINE unsnoc #-}
unsnoc :: Acc a -> Maybe (a, Acc a)
unsnoc :: forall a. Acc a -> Maybe (a, Acc a)
unsnoc =
  \case
    TreeAcc NeAcc a
tree ->
      case NeAcc a
tree of
        NeAcc.Branch NeAcc a
l NeAcc a
r ->
          case forall a. NeAcc a -> NeAcc a -> (a, NeAcc a)
NeAcc.unsnocTo NeAcc a
l NeAcc a
r of
            (a
res, NeAcc a
newTree) ->
              forall a. a -> Maybe a
Just (a
res, forall a. NeAcc a -> Acc a
TreeAcc NeAcc a
newTree)
        NeAcc.Leaf a
res ->
          forall a. a -> Maybe a
Just (a
res, forall a. Acc a
EmptyAcc)
    Acc a
EmptyAcc ->
      forall a. Maybe a
Nothing

-- |
-- Convert to non empty list if it's not empty.
{-# INLINE toNonEmpty #-}
toNonEmpty :: Acc a -> Maybe (NonEmpty a)
toNonEmpty :: forall a. Acc a -> Maybe (NonEmpty a)
toNonEmpty =
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: * -> *) a. Foldable1 t => t a -> NonEmpty a
Foldable1.toNonEmpty forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. forall a. Acc a -> Maybe (NeAcc a)
toNeAcc

-- |
-- Convert to non empty acc if it's not empty.
{-# INLINE toNeAcc #-}
toNeAcc :: Acc a -> Maybe (NeAcc.NeAcc a)
toNeAcc :: forall a. Acc a -> Maybe (NeAcc a)
toNeAcc =
  \case
    TreeAcc NeAcc a
tree ->
      forall a. a -> Maybe a
Just NeAcc a
tree
    Acc a
EmptyAcc ->
      forall a. Maybe a
Nothing

-- |
-- Enumerate in range, inclusively.
{-# INLINE [1] enumFromTo #-}
enumFromTo :: (Enum a, Ord a) => a -> a -> Acc a
enumFromTo :: forall a. (Enum a, Ord a) => a -> a -> Acc a
enumFromTo a
from a
to =
  if a
from forall a. Ord a => a -> a -> Bool
<= a
to
    then forall a. NeAcc a -> Acc a
TreeAcc (forall a. (Enum a, Ord a) => a -> a -> NeAcc a -> NeAcc a
NeAcc.appendEnumFromTo (forall a. Enum a => a -> a
succ a
from) a
to (forall a. a -> NeAcc a
NeAcc.Leaf a
from))
    else forall a. Acc a
EmptyAcc

-- |
-- Construct from list in reverse order.
--
-- This is more efficient than 'fromList',
-- which is actually defined as @fromReverseList . 'reverse'@.
{-# INLINE fromReverseList #-}
fromReverseList :: [a] -> Acc a
fromReverseList :: forall a. [a] -> Acc a
fromReverseList = \case
  a
a : [a]
b -> forall a. NeAcc a -> Acc a
TreeAcc (forall a. [a] -> NeAcc a -> NeAcc a
NeAcc.prependReverseList [a]
b (forall a. a -> NeAcc a
NeAcc.Leaf a
a))
  [a]
_ -> forall a. Acc a
EmptyAcc