{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Concatenation
  ( concatSized
  , concatSized1
  ) where

import Data.List.NonEmpty (NonEmpty((:|)))

import qualified Data.List as L
import qualified Data.List.NonEmpty as NE

-- | Concatenate all the values in the list in the order they
-- are given. This function attempts to perform smaller concatenations
-- together. This is good for data structures that do not take
-- advantage of sharing.
concatSized :: forall m.
     (m -> Int) -- size function 
  -> m
  -> (m -> m -> m)
  -> [m]
  -> m
concatSized :: forall m. (m -> Int) -> m -> (m -> m -> m) -> [m] -> m
concatSized m -> Int
size m
empty m -> m -> m
combine = [m] -> [m] -> m
go [] where
  go :: [m] -> [m] -> m
  go :: [m] -> [m] -> m
go ![m]
stack [] = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' m -> m -> m
combine m
empty (forall a. [a] -> [a]
L.reverse [m]
stack)
  go ![m]
stack (m
x : [m]
xs) = [m] -> [m] -> m
go (m -> [m] -> [m]
pushStack m
x [m]
stack) [m]
xs
  pushStack :: m -> [m] -> [m]
  pushStack :: m -> [m] -> [m]
pushStack m
x [] = [m
x]
  pushStack m
x (m
s : [m]
ss) = if m -> Int
size m
x forall a. Ord a => a -> a -> Bool
>= m -> Int
size m
s
    then m -> [m] -> [m]
pushStack (m -> m -> m
combine m
s m
x) [m]
ss
    else m
x forall a. a -> [a] -> [a]
: m
s forall a. a -> [a] -> [a]
: [m]
ss

-- | This function is likely to be used for things like intersection
-- where the zero-sized element is not an identity but a zero.
concatSized1 :: forall m.
     (m -> Int) -- size function 
  -> (m -> m -> m)
  -> NonEmpty m
  -> m
concatSized1 :: forall m. (m -> Int) -> (m -> m -> m) -> NonEmpty m -> m
concatSized1 m -> Int
size m -> m -> m
combine (m
p :| [m]
ps) = NonEmpty m -> [m] -> m
go (m
p forall a. a -> [a] -> NonEmpty a
:| []) [m]
ps where
  go :: NonEmpty m -> [m] -> m
  go :: NonEmpty m -> [m] -> m
go !NonEmpty m
stack [] = forall a. (a -> a -> a) -> NonEmpty a -> a
safeFoldl1' m -> m -> m
combine (forall a. NonEmpty a -> NonEmpty a
NE.reverse NonEmpty m
stack)
  go !NonEmpty m
stack (m
x : [m]
xs) = NonEmpty m -> [m] -> m
go (m -> NonEmpty m -> NonEmpty m
pushStack m
x NonEmpty m
stack) [m]
xs
  pushStack :: m -> NonEmpty m -> NonEmpty m
  pushStack :: m -> NonEmpty m -> NonEmpty m
pushStack m
x (m
s :| [m]
ss) = if m -> Int
size m
x forall a. Ord a => a -> a -> Bool
>= m -> Int
size m
s
    then case [m]
ss of
      [] -> m -> m -> m
combine m
s m
x forall a. a -> [a] -> NonEmpty a
:| []
      m
r : [m]
rs -> m -> NonEmpty m -> NonEmpty m
pushStack (m -> m -> m
combine m
s m
x) (m
r forall a. a -> [a] -> NonEmpty a
:| [m]
rs)
    else m
x forall a. a -> [a] -> NonEmpty a
:| (m
s forall a. a -> [a] -> [a]
: [m]
ss)

safeFoldl1' :: (a -> a -> a) -> NonEmpty a -> a
safeFoldl1' :: forall a. (a -> a -> a) -> NonEmpty a -> a
safeFoldl1' a -> a -> a
f (a
a :| [a]
as) = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' a -> a -> a
f a
a [a]
as