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 :: 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 (Int -> (a -> a) -> a -> a
forall a. Int -> (a -> a) -> a -> a
nest (Int
nInt -> Int -> Int
forall 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 :: Int -> (a -> a) -> a -> a
nest1 Int
n a -> a
f = ((a -> a) -> (a -> a) -> a -> a) -> (a -> a) -> [a -> a] -> a -> a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) a -> a
forall a. a -> a
id (Int -> (a -> a) -> [a -> a]
forall a. Int -> a -> [a]
replicate Int
n a -> a
f)
nest2 :: Int -> (a -> a) -> a -> a
nest2 Int
n a -> a
f a
x = (a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate a -> a
f a
x [a] -> Int -> a
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 :: (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 t -> Bool
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) (t -> t -> t
forall a. Integral a => a -> a -> a
div t
n t
2)
   in  a -> a -> Integer -> a
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 :: (a -> a -> a) -> a -> a -> Integer -> a
powerAssociativeList a -> a -> a
op a
a0 a
a Integer
n =
   (a -> (Integer, a) -> a) -> a -> [(Integer, a)] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\a
acc (Integer
bit,a
pow) -> if Integer
bitInteger -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
==Integer
0 then a
acc else a -> a -> a
op a
acc a
pow) a
a0 ([(Integer, a)] -> a) -> [(Integer, a)] -> a
forall a b. (a -> b) -> a -> b
$
   [Integer] -> [a] -> [(Integer, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip
      ((Integer -> Maybe (Integer, Integer)) -> Integer -> [Integer]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (\Integer
k -> Bool -> (Integer, Integer) -> Maybe (Integer, Integer)
forall a. Bool -> a -> Maybe a
toMaybe (Integer
kInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>Integer
0) ((Integer, Integer) -> Maybe (Integer, Integer))
-> (Integer, Integer) -> Maybe (Integer, Integer)
forall a b. (a -> b) -> a -> b
$ (Integer, Integer) -> (Integer, Integer)
forall a b. (a, b) -> (b, a)
swap ((Integer, Integer) -> (Integer, Integer))
-> (Integer, Integer) -> (Integer, Integer)
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
k Integer
2) Integer
n)
      ((a -> a) -> a -> [a]
forall a. (a -> a) -> a -> [a]
iterate (\a
pow -> a -> a -> a
op a
pow a
pow) a
a)

powerAssociativeNaive :: (a -> a -> a) -> a -> a -> Integer -> a
powerAssociativeNaive a -> a -> a
op a
a0 a
a Integer
n =
   (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> a -> a
op a
a0 (Integer -> a -> [a]
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
$% :: a -> (a -> b) -> b
($%) = ((a -> b) -> a -> b) -> a -> (a -> b) -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($)