{-# LANGUAGE Safe #-}
{-# LANGUAGE BangPatterns #-}

{- |
Module                  : Relude.Extra.Foldable
Copyright               : (c) 2018-2023 Kowainik
SPDX-License-Identifier : MIT
Maintainer              : Kowainik <xrom.xkov@gmail.com>
Stability               : Stable
Portability             : Portable

Contains utility functions for working with tuples.

@since 0.6.0.0
-}

module Relude.Extra.Foldable
    ( foldlSC
    , average
    ) where

import Relude


-- $setup

{- | A left-associative fold that's tail-recursive but can still short-circuit.
Returning a 'Left' short-circuits and immediately returns the value inside.
Returning a 'Right' continues the fold as usual with the value inside.

>>> foldlSC (\acc x -> if x == 0 then Left 0 else Right $! acc * x) 1 [1..6]
720
>>> foldlSC (\acc x -> if x == 0 then Left 0 else Right $! acc * x) 1 (0:error "Short-circuiting should keep this from happening")
0

@since 0.6.0.0
-}
foldlSC :: forall t b a. Foldable t => (b -> a -> Either b b) -> b -> t a -> b
foldlSC :: forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> Either b b) -> b -> t a -> b
foldlSC b -> a -> Either b b
f = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (b -> b) -> b -> b
go forall a. a -> a
id
  where
    go :: a -> (b -> b) -> b -> b
    go :: a -> (b -> b) -> b -> b
go a
x b -> b
k b
z = case b -> a -> Either b b
f b
z a
x of
        Left b
l  -> b
l
        Right b
r -> b -> b
k b
r
{-# INLINE foldlSC #-}

{- | Given a 'Foldable' of 'Fractional' elements, computes the average if
possible and returns 'Maybe' element.

>>> average [42]
Just 42.0
>>> average @Double [1, 2, 3, 4]
Just 2.5
>>> average @Float [1.5, 2.5, 3 ,4]
Just 2.75
>>> average []
Nothing

@since 1.0.0.0
-}
average :: forall a f . (Foldable f, Fractional a) => f a -> Maybe a
average :: forall a (f :: * -> *).
(Foldable f, Fractional a) =>
f a -> Maybe a
average f a
xs
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null f a
xs = forall a. Maybe a
Nothing
    | Bool
otherwise = forall a. a -> Maybe a
Just
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Fractional a => a -> a -> a
(/)
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\(!a
total, !a
count) a
x -> (a
total forall a. Num a => a -> a -> a
+ a
x, a
count forall a. Num a => a -> a -> a
+ a
1)) (a
0,a
0)
        forall a b. (a -> b) -> a -> b
$ f a
xs
{-# INLINE average #-}