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