    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 (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 :: forall a. Int -> (a -> a) -> a -> a
    nest1 Int
    n a -> a
    f = ((a -> a) -> (a -> a) -> a -> a) -> (a -> a) -> [a -> a] -> a -> a
    forall a b. (a -> b -> b) -> b -> [a] -> b
    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 :: forall a. 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. HasCallStack => [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 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 :: forall a. (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 b a. (b -> a -> b) -> b -> [a] -> b
    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 :: forall a. (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 a b. (a -> b -> b) -> b -> [a] -> b
    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
    $% :: forall a 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
    ($)