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!

Answer 1

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.

Answer 2

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