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.Core.Transform.AnonymizeX as A 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 -- | Arities of known bound variables. -- We need to track everything even if it's not a function to keep indices correct. -- Just use zero for unknown/irrelevant type Arities n = (Map.Map n Int, [Int]) -- | Empty arities context arEmpty :: Ord n => Arities n arEmpty = (Map.empty, []) -- | Extend map with multiple bindings and their arities 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) -- | Look up a binder's arity arGet :: Ord n => Arities n -> Bound n -> Int arGet (_named, anon) (UIx ix _) = anon !! ix arGet (named, _anon) (UName n _) = named Map.! n -- Get a primitive's arity from its type arGet (_named,_anon) (UPrim _ t) = arityOfType t -- **** Finding arities of expressions etc -- | Count all the arrows and foralls, ignoring any effects -- We can be sure that primitives don't effect until they're fully applied arityOfType :: Ord n => Type n -> Int arityOfType (T.TForall _ t) = 1 + arityOfType t arityOfType t = let (args, _) = T.takeTFunArgResult t in length args -- | Find arity of an expression. Count lambdas, use type for primitives arityOfExp :: Ord n => Exp a n -> Int -- Counting all binders, because they all correspond to XApps. arityOfExp (XLam _ _ e) = 1 + arityOfExp e arityOfExp (XLAM _ _ e) = 1 + arityOfExp e -- Find primitive's constructor's arities from type, -- we might need to do this for user defined constructors too. arityOfExp (XCon _ (UPrim _ t)) = arityOfType t -- Anything else we'll need to apply one at a time arityOfExp _ = 0 -- | 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 :: Ord n => Pat n -> [(Bind n, Int)] aritiesOfPat PDefault = [] aritiesOfPat (PData _b bs) = zip bs (repeat 0) -- **** Actually converting to a-normal form -- | Recursively transform expression into a-normal anormal :: Ord n => Arities n -- ^ environment, arities of bound variables -> Exp a n -- ^ expression to transform -> [(Exp a n,a)]-- ^ arguments being applied to current expression -> Exp a n -- Application: just record argument and descend into function anormal ar (XApp a lhs rhs) args = -- normalise rhs and add to arguments let args' = (anormal ar rhs [], a) : args in -- descend into lhs, remembering all args anormal ar lhs args' -- Anything other than application: if we're applied to arguments add bindings, -- otherwise just recurse. anormal ar x args = let x' = go x in case args of -- if there are no args, we're done [] -> x' -- there are arguments. we must apply them. _ -> flattenLets $ 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 "DDC.Core.Transform.ANormal.anormal: impossible XApp!" -- leafy ones go (XVar{}) = x go (XCon{}) = x go (XType{}) = x go (XWitness{}) = x -- lambdas 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) -- withregion: I don't think this should ever show up. go (XLet a (LWithRegion b) re) = XLet a (LWithRegion b) (down [] re) -- case 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' -- cast go (XCast a c e) = XCast a c (down [] e) -- | Convert an expression into a-normal form 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 be applied as-is isNormal :: Ord n => Exp a n -> Bool -- Trivial expressions isNormal (XVar{}) = True isNormal (XCon{}) = True isNormal (XType{}) = True isNormal (XWitness{}) = True -- Casts are ignored by code generator, so we can leave them in if their subexpression is normal isNormal (XCast _ _ x) = isNormal x isNormal _ = False -- | Create lets for any non-trivial arguments makeLets :: Ord n => Arities n -- ^ environment, arities of bound variables -> Exp a n -- ^ function -> [(Exp a n,a)]-- ^ arguments being applied to current expression -> Exp a n makeLets _ f0 [] = f0 makeLets ar f0 args@((_,annot):_) = go 0 (findArity f0) ((f0,annot):args) [] where tBot = T.tBot T.kData -- out of arguments, create XApps out of leftovers go i _arf [] acc = mkApps i 0 acc -- f is fully applied and we have arguments left to add: -- create let for intermediate result go i arf ((x,a):xs) acc | length acc > arf = XLet a (LLet LetStrict (BAnon tBot) (mkApps i 0 acc)) (go i 1 ((x,a):xs) [(XVar a $ UIx 0 tBot,a)]) -- application to variable, don't bother binding go i arf ((x,a):xs) acc | isNormal x = go i arf xs ((x,a):acc) -- non-trivial argument, create binding go i arf ((x,a):xs) acc = XLet a (LLet LetStrict (BAnon tBot) (L.liftX i x)) (go (i+1) arf xs ((x,a):acc)) -- fold list into applications -- can't create empty app mkApps _ _ [] = error "DDC.Core.Transform.ANormal.makeLets.mkApps: impossible empty list" -- single element - this is the function mkApps l _ [(x,_)] | isNormal x = L.liftX l x mkApps _ i [(_,a)] = XVar a $ UIx i tBot -- apply this argument and recurse mkApps l i ((x,a):xs) | isNormal x = XApp a (mkApps l i xs) (L.liftX l x) mkApps l i ((_,a):xs) = XApp a (mkApps l (i+1) xs) (XVar a $ UIx i tBot) findArity (XVar _ b) = max (arGet ar b) 1 findArity x = max (arityOfExp x) 1 -- | Perform let-floating on strict non-recursive lets -- Only does the top level, to clean up the ones directly produced by makeLets. -- let b1 = (let b2 = def2 in x2) -- in x1 -- ==> -- let b2 = def2 -- in let b1 = x2 -- in x1 flattenLets :: Ord n => Exp a n -> Exp a n -- We only do this if b2 is anonymous (ones generated by makeLets are). -- If we tried to wrap x1 in b2 when b2's name is already used, -- we'd be in trouble. flattenLets (XLet a1 (LLet LetStrict b1 (XLet a2 (LLet LetStrict b2@(BAnon _) def2) x2)) x1) = -- If b1 is anon, we don't want to lift references to it let liftDepth = case b1 of { BAnon _ -> 1; _ -> 0 } in let x1' = L.liftAtDepthX 1 liftDepth x1 in XLet a2 (LLet LetStrict b2 def2) $ flattenLets $ XLet a1 (LLet LetStrict b1 x2) x1' -- Same as above but b2 isn't anonymous - anonymize inner let & re-flatten. flattenLets (XLet a1 (LLet LetStrict b1 inner@(XLet _ (LLet LetStrict _ _) _)) x1) = flattenLets $ XLet a1 (LLet LetStrict b1 (A.anonymizeX inner)) x1 -- Any let, its bound expression doesn't contain a strict non-recursive let so just flatten the body flattenLets (XLet a1 llet1 x1) = XLet a1 llet1 (flattenLets x1) -- Anything else we can ignore. We don't need to recurse, because this is always called immediately after makeLets. flattenLets x = x