Writing recursive template haskell functions

Question

Is it possible to transform a recursive TH function into an equivalent form which will compile? The following definition doesn't work, because in order to compile fact you must first compile fact.

fact :: ExpQ -> ExpQ
fact n = [|
    case $n of 
      1 -> 1
      _ -> $n * $(fact [| $n - 1 |]) |]

This simple example is easily solved (fact n = [| product [ 1.. $n] |]) but in the general case, if it is not possible to rewrite a given function as a loop, can a recursive TH function be defined? Is there even a single example for which this doable?

To clarify for future readers: this question is specifically about writing recursive TH functions - not about 'how do I splice the factorial function'.

The answer to my question turned out to be remarkably simple:

{-# LANGUAGE TemplateHaskell #-}

import Control.Monad.Fix (fix)
import Language.Haskell.TH

fact = [| \f x -> $([| 
     case x of 
       1 -> 1 
       _ -> f $([| x - 1 |]) * x  |]) |]

factorial = [| fix $fact |]

fact can be compiled because it is no longer recursive, and [| fix $fact |] is compiled at a later time, so no more infinite recursive definitions.

This version of fact looks slightly different than the original, but you can write the new fact exactly as the old one and transform it later:

fact' recurse n = [|
        case $n of 
          1 -> 1
          _ -> $n * $(recurse [| $n - 1 |]) |]

fact = [| \x -> $((\f -> [| \x -> $(fact (\x -> [| $f $x |]) [| x |]) |]) [| x |]) |]

Answer 1

The fundamental problem with your code is not that it is recursive, but that it doesn't terminate. The n argument to fact just keeps getting bigger and bigger because [| $n - 1 ]| is an expression tree representing the operation (-) applied to n and 1.

Any non-terminating splice will hang the compiler in just the same way, for example the following behaves just like your fact when spliced:

broken :: ExpQ -> ExpQ
broken n = return $ LitE (IntegerL (fromIntegral (length [1..])))

A recursive function where the recursion is guaranteed to bottom out is guaranteed to terminate and works fine for appropriate inputs:

fact1 :: ExpQ -> ExpQ
fact1 n = do
    nBody <- n
    case nBody of
      LitE (IntegerL 1) -> [|1|]
      LitE (IntegerL nInt) | nInt > 1 ->
          let nMinusOne = return $ LitE (IntegerL (nInt-1))
          in [| $n * $(fact1 nMinusOne) |]

but of course it fails if the input isn't an appropriate integer literal.

You can also shift the recursion to runtime, so that instead of the recursive call being with an ever-bigger expression tree, it's with the runtime evaluated and shrinking Int:

fact2 :: ExpQ -> ExpQ
fact2 n =
   [|
    let factImpl n =
         case n of
          1 -> 1
          _ -> n * factImpl (n-1)
    in factImpl $n
   |]

Of course in this code we're not doing any analysis of the structure of n. But we can put it together with fact1 to get a version that is compile-time executed in some cases and defers others to runtime:

fact3 :: ExpQ -> ExpQ
fact3 n = do
    nBody <- n
    case nBody of
      LitE (IntegerL 1) -> [|1|]
      LitE (IntegerL nInt) ->
          let nMinusOne = return $ LitE (IntegerL (nInt-1))
          in [| $n * $(fact3 nMinusOne) |]
      _ -> [|
            let factImpl n =
                  case n of
                    1 -> 1
                    _ -> n * factImpl (n-1)
            in factImpl $n
           |]

Ultimately in your real code, you will need to apply some combination of these techniques - make sure that your compile-time recursion terminates and defer any remaining cases to runtime evaluation somehow.

Answer 2

Yes, you can by using the following:

fact :: Int -> ExpQ
fact 0 = [| 1 |]
fact n = [| $(lift n) * $(fact $ n - 1) |]

lift is a function inside Language.Haskell.TH.Lift which converts a basic haskell values into template haskell values (eg Int to ExpQ).

Note that you don't need the case code to be generated, as you know the number at compile time. The above macro will expand to a series of multiplications. Eg $(fact 4) will expand to 4*3*2*1.

Note that in this case, you can do much better though. A template haskell expression is run at compile time, so a template haskell fact function can just return the literal value it represents. Eg $(fact 4) can return 24 (instead of 4*3*2*1). This can be done with the following code:

fact2 :: Int -> ExpQ
fact2 n = lift (product [1..n])