Template Haskell: How to extract the number of arguments of a function?

Question

I have a function which given a Name of a function it augments it, yielding another function applied to some other stuff (details not very relevant):

mkSimple :: Name -> Int -> Q [Dec]
mkSimple adapteeName argsNum = do
  adapterName <- newName ("sfml" ++ (capitalize . nameBase $ adapteeName))
  adapteeFn <- varE adapteeName
  let args = mkArgs argsNum
  let wrapper = mkApply adapteeFn (map VarE args)
  -- generates something like SFML $ liftIO $ ((f a) b) c)
  fnBody <- [| SFML $ liftIO $ $(return wrapper) |]
  return [FunD adapterName [Clause (map VarP args) (NormalB fnBody) []]]

  where
    mkArgs :: Int -> [Name]
    mkArgs n = map (mkName . (:[])) . take n $ ['a' .. 'z']


-- Given f and its args (e.g. x y z) builds ((f x) y) z)
mkApply :: Exp -> [Exp] -> Exp
mkApply fn [] = fn
mkApply fn (x:xs) = foldr (\ e acc -> AppE acc e) (AppE fn x) xs

This works, but it's tedious to pass externally the number of args the adaptee function has. There exists some TH function to extract the number of args? I suspect it can be achieved with reify but I don't know how.

Thanks!

jberryman · Accepted Answer · 2013-12-07T19:27:38.753

6

Sure, you should be able to do

do (VarI _ t _ _) <- reify adapteeName
   -- t :: Type
   -- e.g. AppT (AppT ArrowT (VarT a)) (VarT b)
   let argsNum = countTheTopLevelArrowTs t
   ...

   where
     countTheTopLevelArrowTs (AppT (AppT ArrowT _) ts) = 1 + countTheTopLevelArrowTs
     countTheTopLevelArrowTs _ = 0

The above is just from my head and may not be quite right.

edited Dec 07 '13 at 19:27

answered Dec 07 '13 at 18:17

jberryman

16,334
5
42
83

Simple enough, even though extracting the number of arrows seems cumbersome and involving a lot of pattern matches. It can be made recursive, but I can't see a simple base case for the recursion. Any idea? Obviously something like T.count "ArrowT" (show t) is error prone (e.g. if you have :: (a -> b) -> c -> d the count would be wrong) – Alfredo Di Napoli Dec 07 '13 at 19:06
@AlfredoDiNapoli sorry to be lazy with `countTheTopLevelArrowTs` in my answer. It is indeed cumbersome, but shouldn't be ambiguous: notice you don't need to recurse down the first arguments of any of the `ArrowT`s I'll update my answer – jberryman Dec 07 '13 at 19:26
2

It is ambiguous.. How many arguments are there in a function with the signature `Arrow a => Int -> a Int Int`? Polymorphism means that there will always be some ambiguity. – Carl Dec 07 '13 at 20:20
@Carl There is one argument, but you can start to argue about currying and how the answer is always one. I'd say the only sane thing to do is count the number of explicit arrows and ignore the fact that `a Int Int` can be instantiated as `Int -> Int` since the body of the function won't treat it this way in any case. – daniel gratzer Dec 07 '13 at 20:29
3

You might also need to expand type synonyms. – augustss Dec 07 '13 at 21:14

score 1 · Answer 2 · answered Sep 27 '14 at 12:18

A slight improvement on jberryman's answer that deals with type constraints such as (Ord a) -> a -> a is:

arity :: Type -> Integer
arity = \case
    ForallT _ _ rest -> arity rest
    AppT (AppT ArrowT _) rest -> arity rest +1
    _ -> 0

usage:

do (VarI _ t _ _) <- reify adapteeName
    let argsNum = arity t

Template Haskell: How to extract the number of arguments of a function?

2 Answers2