module DDC.Core.Transform.ANormal (anormalise) where import DDC.Core.Exp import qualified DDC.Type.Exp as T import qualified DDC.Type.Compounds as T import qualified DDC.Type.Universe as U import qualified DDC.Core.Transform.LiftX as L import qualified Data.Map as Map -- **** Recording arities of known values -- So we can try to create apps to fully apply -- I did have these as Maybe Int, but I think for our purposes 0==Nothing is fine type Arities n = (Map.Map n Int, [Int]) arEmpty :: Ord n => Arities n arEmpty = (Map.empty, []) arExtends :: Ord n => Arities n -> [(Bind n, Int)] -> Arities n arExtends 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) arGet :: Ord n => Arities n -> Bound n -> Int -- TODO unsafe ix arGet (_named, anon) (UIx ix _) = anon !! ix arGet (named, _anon) (UName n _) = named Map.! n -- Get a primitive's arity from its type. -- Assuming all the primitives defer effects until fully applied. arGet (_named,_anon) (UPrim _ t) = arityOfType t -- **** Finding arities of expressions etc -- Count all the arrows, ignoring any effects arityOfType :: Ord n => Type n -> Int arityOfType (T.TForall _ t) = 1 + arityOfType t arityOfType t = let (args, _) = T.takeTFunArgResult t in length args arityOfExp :: Ord n => Exp a n -> Int arityOfExp (XLam _ b e) -- only count data binders | isBinderData b = 1 + arityOfExp e arityOfExp (XLam _ _ e) = 1 + arityOfExp e arityOfExp (XLAM _ _ e) = 1 + arityOfExp e arityOfExp (XCon _ (UPrim _ t)) = arityOfType t arityOfExp _ = 0 isBinderData :: Ord n => Bind n -> Bool isBinderData b | Just U.UniverseData <- U.universeFromType1 (T.typeOfBind b) = True isBinderData _ = False -- We don't know anything about their values, -- but we need to record them as 0 anyway (shadowing, de bruijn) aritiesOfPat :: Ord n => Pat n -> [(Bind n, Int)] aritiesOfPat PDefault = [] aritiesOfPat (PData _b bs) = zip bs (repeat 0) -- **** Actually converting to a-normal form anormal :: Ord n => Arities n -> Exp a n -> [Exp a n] -> Exp a n anormal ar (XApp _ lhs rhs) args = -- normalise applicand and record arguments let args' = anormal ar rhs [] : args in -- descend into lhs, remembering all args anormal ar lhs args' anormal ar x args = let x' = go x in -- if there are no args, we're done case args of [] -> x' _ -> -- there are arguments. we must apply them. makeLets ar x' args where -- helper for descent down ars e = anormal (arExtends ar ars) e [] -- we know x isn't an app. go (XApp{}) = error "ANormal.anormal: impossible XApp!" -- leafy ones go (XVar{}) = x go (XCon{}) = x go (XType{}) = x go (XWitness{}) = x go (XLAM a b e) = XLAM a b (down [(b,0)] e) go (XLam a b e) = XLam a b (down [(b,0)] e) -- non-recursive let go (XLet a (LLet m b le) re) = let le' = down [] le in let re' = down [(b, arityOfExp le')] re in XLet a (LLet m b le') re' -- recursive let go (XLet a (LRec lets) re) = let bs = map fst lets in let es = map snd lets in let ars= zip bs (map arityOfExp es) in let es'= map (down ars) es in let re'= down ars re in XLet a (LRec $ zip bs es') re' -- letregion, just make sure we record bindings with dummy val go (XLet a (LLetRegion b bs) re) = let ars = zip bs (repeat 0) in XLet a (LLetRegion b bs) (down ars re) -- I don't think a withregion should ever show up... go (XLet a (LWithRegion b) re) = XLet a (LWithRegion b) (down [] re) go (XCase a e alts) = let e' = down [] e in let alts' = map (\(AAlt pat ae) -> AAlt pat (down (aritiesOfPat pat) ae)) alts in XCase a e' alts' go (XCast a c e) = XCast a c (down [] e) -- | (under development) anormalise :: Ord n => Exp a n -> Exp a n anormalise x = anormal arEmpty x [] -- | Check if an expression needs a binding, or if it's simple enough to just be applied isNormal :: Ord n => Exp a n -> Bool isNormal (XVar{}) = True isNormal (XCon{}) = True isNormal (XType{}) = True isNormal (XWitness{}) = True isNormal (XCast _ _ x) = isNormal x isNormal _ = False makeLets ar f0 args = go 0 (findArity f0) (f0:args) [] where tBot = T.tBot T.kData -- sending arity of f to this is a hack because we should really be building up ar ctx? go i _arf [] acc = mkApps i 0 acc -- f is fully applied, and we *do* have arguments left to add go i arf (x:xs) acc | length acc > arf = XLet (annotOf x) (LLet LetStrict (BAnon tBot) (mkApps i 0 acc)) (go i 1 (x:xs) [XVar (annotOf x) $ UIx 0 tBot]) -- application to variable, don't bother binding go i arf (x:xs) acc | isNormal x = go i arf xs (x:acc) -- create binding go i arf (x:xs) acc = XLet (annotOf x) (LLet LetStrict (BAnon tBot) (L.liftX i x)) (go (i+1) arf xs (x:acc)) mkApps _ _ [] = error "ANormal.makeLets.mkApps: impossible empty list" mkApps l _ [x] | isNormal x = L.liftX l x mkApps _ i [x] = XVar (annotOf x) $ UIx i tBot mkApps l i (x:xs) | isNormal x = XApp (annotOf x) (mkApps l i xs) (L.liftX l x) mkApps l i (x:xs) = XApp (annotOf x) (mkApps l (i+1) xs) (XVar (annotOf x) $ UIx i tBot) findArity (XVar _ b) = max (arGet ar b) 1 findArity x = max (arityOfExp x) 1 -- does this exist elsewhere? ought it? annotOf :: Exp a n -> a annotOf (XVar a _) = a annotOf (XCon a _) = a annotOf (XApp a _ _) = a annotOf (XLAM a _ _) = a annotOf (XLam a _ _) = a annotOf (XLet a _ _) = a annotOf (XCase a _ _) = a annotOf (XCast a _ _) = a annotOf (XType{}) = error "DDC.Core.Transform.ANormal.annotOf: XType" annotOf (XWitness{}) = error "DDC.Core.Transform.ANormal.annotOf: XWitness"