content: 
  -- | Slurp out arities of function bindings.
  --   and infer arities for primitives based on their types.
  --
  --   For function bindings the arity is the number of outer-most lambdas
  --   in the definition.
  --
  --   For primitives, the arity is the number of function
  --   constructors in its type.
  --
  module DDC.Core.Analysis.Arity
          ( -- * Arities map
            Arities
          , emptyArities
          , extendsArities
          , getArity
  
            -- * Arity analysis
          , aritiesOfModule
          , aritiesOfLets
          , aritiesOfPat
          , arityFromType
          , arityOfExp)
  where
  import DDC.Core.Module
  import DDC.Core.Exp.Annot
  import DDC.Data.ListUtils
  import Control.Monad
  import Data.Maybe
  import qualified Data.Map       as Map
  
  
  -- | Arities of named and anonymous bindings.
  type Arities n = (Map.Map n Int, [Int])
  
  
  -- | Empty arities context.
  emptyArities :: Arities n
  emptyArities = (Map.empty, [])
  
  
  -- | Extend map with some binders and their arities.
  extendsArities :: Ord n => Arities n -> [(Bind n, Int)] -> Arities n
  extendsArities arity exts = foldl go arity exts
   where  go (named, anon) (BNone _t,   _) = (named,anon)
          go (named, anon) (BAnon _t,   a) = (named, a:anon)
          go (named, anon) (BName n _t, a) = (Map.insert n a named, anon)
  
  
  -- | Look up a binder's arity from the arity map
  --   or determine it from its type in the case of primops.
  getArity :: Ord n => Arities n -> Bound n -> Maybe Int
  getArity (named, anon) u
   = case u of
          -- Get arities of anonymous things from the stack.
          UIx ix
           -> let Just x  = index anon ix
              in  Just x
  
          -- Lookup arities of named things from the stack.
          UName n -> Map.lookup n named
  
  
  -- Slurp ----------------------------------------------------------------------
  -- | Slurp out arities of imports and top-level bindings from a module.
  aritiesOfModule :: Ord n => Module a n -> Arities n
  aritiesOfModule mm
    = let (lts, _)        = splitXLets $ moduleBody mm
  
          aritiesLets
           = concat $ catMaybes $ map aritiesOfLets  lts
  
          aritiesImports
           = catMaybes
           $ [ case arityFromType (typeOfImportValue isrc) of
                  Just a  -> Just (BName n (typeOfImportValue isrc), a)
                  Nothing -> Nothing
             | (n, isrc) <- moduleImportValues mm ]
  
      in  emptyArities
          `extendsArities` aritiesImports
          `extendsArities` aritiesLets
  
  
  -- | Get the arities of a `Lets`
  aritiesOfLets :: Ord n => Lets a n -> Maybe [(Bind n, Int)]
  aritiesOfLets ll
   = let get (b, x)
          = case arityOfExp x of
                  Nothing         -> Nothing
                  Just a          -> Just (b, a)
     in case ll of
          LLet b x        -> sequence $ map get [(b, x)]
          LRec bxs        -> sequence $ map get bxs
          _               -> Just []
  
  
  -- | Retrieve binders from case pattern, so we can extend the arity context.
  --   We don't know anything about their values, so record as 0.
  aritiesOfPat :: Pat n -> [(Bind n, Int)]
  aritiesOfPat p
   = case p of
          PDefault        -> []
          (PData _b bs)   -> zip bs (repeat 0)
  
  
  -- | Get the arity of an expression.
  arityOfExp :: Ord n => Exp a n -> Maybe Int
  arityOfExp xx
   = case xx of
          -- Counting all binders, because they all correspond to XApps.
          XLam _ _ e      -> liftM (+ 1) $ arityOfExp e
          XLAM _ _ e      -> liftM (+ 1) $ arityOfExp e
  
          -- Determine a data constructor's arity from its type.
          -- FIXME: do we need dacon arities any more?
  --        XCon _ (DaConPrim _)
  --                -> arityFromType t
  
          -- Anything else we'll need to apply one at a time
          _               -> Just 0
  
  
  -- | Determine the arity of an expression by looking at its type.
  --   Count all the function arrows, and foralls.
  arityFromType :: Ord n => Type n -> Maybe Int
  arityFromType tt
          | TForall _ t   <- tt
          = case arityFromType t of
                  Nothing -> Nothing
                  Just a  -> Just (1 + a)
  
          | isBot tt
          = Nothing
  
          | (args, _)     <- takeTFunArgResult tt
          = Just (length args)
  
  