module Data.Function.HT.Private where

import Data.List (genericReplicate, unfoldr)
import Data.Maybe.HT (toMaybe)
import Data.Tuple.HT (swap)

-- $setup
-- >>> import Test.QuickCheck (NonNegative(NonNegative))

{- |
Compositional power of a function,
i.e. apply the function @n@ times to a value.
It is rather the same as @iter@
in Simon Thompson: \"The Craft of Functional Programming\", page 172
-}
{-# INLINE nest #-}
nest :: Int -> (a -> a) -> a -> a
nest :: forall a. Int -> (a -> a) -> a -> a
nest Int
0 a -> a
_ a
x = a
x
nest Int
n a -> a
f a
x = a -> a
f (forall a. Int -> (a -> a) -> a -> a
nest (Int
nforall a. Num a => a -> a -> a
-Int
1) a -> a
f a
x)

{- |
prop> \(NonNegative n) x -> nest n succ x == nest1 n succ (x::Integer)
prop> \(NonNegative n) x -> nest n succ x == nest2 n succ (x::Integer)
-}
nest1, nest2 :: Int -> (a -> a) -> a -> a
nest1 :: forall a. Int -> (a -> a) -> a -> a
nest1 Int
n a -> a
f = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) forall a. a -> a
id (forall a. Int -> a -> [a]
replicate Int
n a -> a
f)
nest2 :: forall a. Int -> (a -> a) -> a -> a
nest2 Int
n a -> a
f a
x = forall a. (a -> a) -> a -> [a]
iterate a -> a
f a
x forall a. [a] -> Int -> a
!! Int
n


{- |
@powerAssociative@ is an auxiliary function that,
for an associative operation @op@,
computes the same value as

  @powerAssociative op a0 a n = foldr op a0 (genericReplicate n a)@

but applies "op" O(log n) times and works for large n.
-}

{-# INLINE powerAssociative #-}
powerAssociative :: (a -> a -> a) -> a -> a -> Integer -> a
powerAssociative :: forall a. (a -> a -> a) -> a -> a -> Integer -> a
powerAssociative a -> a -> a
op =
   let go :: a -> a -> t -> a
go a
acc a
_ t
0 = a
acc
       go a
acc a
a t
n = a -> a -> t -> a
go (if forall a. Integral a => a -> Bool
even t
n then a
acc else a -> a -> a
op a
acc a
a) (a -> a -> a
op a
a a
a) (forall a. Integral a => a -> a -> a
div t
n t
2)
   in  forall {t}. Integral t => a -> a -> t -> a
go

{- |
prop> \a0 a (NonNegative n) -> powerAssociative (+) a0 a n == (powerAssociativeList (+) a0 a n :: Integer)
prop> \a0 a (NonNegative n) -> powerAssociative (+) a0 a n == (powerAssociativeNaive (+) a0 a n :: Integer)
-}
powerAssociativeList, powerAssociativeNaive ::
   (a -> a -> a) -> a -> a -> Integer -> a
powerAssociativeList :: forall a. (a -> a -> a) -> a -> a -> Integer -> a
powerAssociativeList a -> a -> a
op a
a0 a
a Integer
n =
   forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\a
acc (Integer
bit,a
pow) -> if Integer
bitforall a. Eq a => a -> a -> Bool
==Integer
0 then a
acc else a -> a -> a
op a
acc a
pow) a
a0 forall a b. (a -> b) -> a -> b
$
   forall a b. [a] -> [b] -> [(a, b)]
zip
      (forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (\Integer
k -> forall a. Bool -> a -> Maybe a
toMaybe (Integer
kforall a. Ord a => a -> a -> Bool
>Integer
0) forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> (b, a)
swap forall a b. (a -> b) -> a -> b
$ forall a. Integral a => a -> a -> (a, a)
divMod Integer
k Integer
2) Integer
n)
      (forall a. (a -> a) -> a -> [a]
iterate (\a
pow -> a -> a -> a
op a
pow a
pow) a
a)

powerAssociativeNaive :: forall a. (a -> a -> a) -> a -> a -> Integer -> a
powerAssociativeNaive a -> a -> a
op a
a0 a
a Integer
n =
   forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> a -> a
op a
a0 (forall i a. Integral i => i -> a -> [a]
genericReplicate Integer
n a
a)


infixl 0 $%

{- |
Flipped version of '($)'.

It was discussed as (&) in
http://www.haskell.org/pipermail/libraries/2012-November/018832.html

I am not sure, that we need it.
It is not exported for now.
-}
($%) :: a -> (a -> b) -> b
$% :: forall a b. a -> (a -> b) -> b
($%) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
($)