Execute a function n times, where n is known at compile time

Question

Motivation

In this question, a user asked whether it is possible to inline the following function:

-- simplified version
{-# INLINE nTimes #-}
nTimes :: Int -> (a -> a) -> a -> a    
nTimes 0 f x = x
nTimes n f x = nTimes (n-1) f (f x)

Unfortunately, the answer seems no, since GHC sees a recursive function and gives up. Even if you use a compile-time constant, e.g. nTimes 1 (+1) x, you don't end up with x + 1, but with nTimes 1 (+1) x.

While it's of course fine to refuse inlining if the number of loops is unknown, it's also a hassle if it is known.

Code

As you can see in the question above, I've proposed the following solution:

{-# LANGUAGE TemplateHaskell #-}
module Times where

import Control.Monad (when)
import Language.Haskell.TH

-- > item under review begins here

nTimesTH :: Int -> Q Exp
nTimesTH n = do
  f <- newName "f"
  x <- newName "x"

  when (n <= 0)    (reportWarning "nTimesTH: argument non-positive")
  when (n >= 1000) (reportWarning "nTimesTH: argument large, can lead to memory exhaustion")

  let go k | k <= 0 = VarE x
      go k          = AppE (VarE f) (go (k - 1))
  return $ LamE [VarP f,VarP x] (go n)

-- < item under review ends here

This should, for any n, create a function with patterns named f and x, and apply f to x with AppE n times:

$(nTimesTH 0) = \f x -> x
$(nTimesTH 1) = \f x -> f x
$(nTimesTH 2) = \f x -> f (f x)
$(nTimesTH 3) = \f x -> f (f (f x))

I can verify that the created function has the correct type:

$ ghci -XTemplateHaskell Times.sh 
ghci> :t $(nTimesTH 0)
$(nTimesTH 0) :: r -> r1 -> r1

ghci> :t $(nTimesTH 1)
$(nTimesTH 1) :: (r1 -> r) -> r1 -> r

ghci> :t $(nTimesTH 2)
$(nTimesTH 2) :: (r -> r) -> r -> r

ghci> :t $(nTimesTH 3)
$(nTimesTH 3) :: (r -> r) -> r -> r

To all my knowledge, nTimesTH works exactly as expected.

Given that this is my first time dabbling with Template Haskell, does this follow best practices? Also, I'm using VarE, AppE and so on. Language.Haskell.TH also provides some combinators, so that one can write

  let go k | k <= 0 = varE x
      go k          = appE (varE f) (go (k - 1))
  lamE (map varP [f,x]) (go n)

Is this just personal preference, or is lamE preferred? As far as I can see, the expression combinators use the canonic implementation, e.g. varE = return . VarE, appE f x = liftA2 AppE f x and so on.

Tests

Just for completeness the QuickCheck tests. Those aren't really part of the review, but here fore completeness:

module Times where
import Test.QuickCheck

-- .. rest of module as above

genTestSingle :: Int -> Q Exp
genTestSingle n = do
  f <- newName "gTSf"
  x <- newName "gTSx"

  lamE [varP f, varP x] $ 
    appsE [ [| (==) |], appsE [nTimesTH n,           varE f, varE x]
                      , appsE [ [| nTimesFoldr n |], varE f, varE x]]

genAllTest :: Int -> Q Exp
genAllTest n = do
  f <- newName "gATf"
  x <- newName "gATv"

  lamE [varP f, varP x] $ doE $ (flip map) [1..n] $ \i ->
    noBindS $ appE [| quickCheck |] $ appsE [genTestSingle i, varE f, varE x]

module Main where

main = $(genAllTest 100) sin 40

Answer 1

Your implementation is fine. But it is a little bit ironic: In your SO answer, you start of with "you can use better functions", yet you rewrite nTimes via go. If this was a real module, you would probably want to export both functions nTimesTH and nTimes, and implement the former via the latter:

{-# LANGUAGE TemplateHaskell #-}
module Times (nTimes, nTimesTH) where

import Control.Monad (when)
import Language.Haskell.TH

-- example implementation
nTimes :: Int -> (a -> a) -> a -> a
nTimes n f x = iterate f x !! max 0 n

nTimesTH :: Int -> Q Exp
nTimesTH n = do
  when (n <= 0)    (reportWarning "nTimesTH: argument non-positive")
  when (n >= 1000) (reportWarning "nTimesTH: argument large, can lead to memory exhaustion")

  f <- newName "f"
  x <- newName "x"

  lamE (map varP [f,x]) $ nTimes n (appE (varE f)) (varE x)

This removes the need for code duplication, and any future optimization you find for nTimes gets automatically applied in nTimesTH, although that only matter during the compile time. After all $(nTimesTH n) is completely evaluated during compilation.

Other than that, I moved the warnings on top of the function, since this brings the names f and x closer to their actual use. I also replaced the data constructors with their respective functions.

Note that for a real module you still want to add documentation:

-- | @'nTimes' n f x@ applies @f@ @n@ times over @x@. 
-- It does not apply @f@ if @n@ is negative. In this case @x@ is returned.
--
-- prop> nTimes n f x == iterate f x !! max 0
nTimes :: Int -> (a -> a) -> a -> a

-- | @'nTimesTH' n@ returns a function that applies the first 
--   argument @n@ times to the second argument. 
--   @n@ must be known at compile time. Needs @-XTemplateHaskell@.
--
-- prop> $(nTimesTH n) f x == nTimes n f x
nTimesTH :: Int -> Q Exp