module DDC.Core.Analysis.Arity
(
Arities
, emptyArities
, extendsArities
, getArity
, 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
type Arities n = (Map.Map n Int, [Int])
emptyArities :: Ord n => Arities n
emptyArities = (Map.empty, [])
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)
getArity :: Ord n => Arities n -> Bound n -> Maybe Int
getArity (named, anon) u
= case u of
UIx ix
-> let Just x = index anon ix
in Just x
UName n -> Map.lookup n named
UPrim _ t -> arityFromType t
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
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 []
aritiesOfPat :: Ord n => Pat n -> [(Bind n, Int)]
aritiesOfPat p
= case p of
PDefault -> []
(PData _b bs) -> zip bs (repeat 0)
arityOfExp :: Ord n => Exp a n -> Maybe Int
arityOfExp xx
= case xx of
XLam _ _ e -> liftM (+ 1) $ arityOfExp e
XLAM _ _ e -> liftM (+ 1) $ arityOfExp e
XCon _ (DaConPrim _ t)
-> arityFromType t
_ -> Just 0
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)