-- | Snip out nested applications. module DDC.Core.Transform.Snip (Snip(..)) where import DDC.Core.Analysis.Arity import DDC.Core.Module import DDC.Core.Exp import DDC.Core.Compounds import qualified DDC.Core.Transform.LiftX as L import qualified DDC.Type.Compounds as T class Snip (c :: * -> *) where -- | Snip out nested applications as anonymous bindings. -- -- @ -- f (g x) (h y) -- ==> let ^ = g x in ^ = h y in f ^1 ^0 -- @ snip :: Ord n => Bool -- ^ Introduce extra bindings for over-applied functions. -> c n -> c n instance Snip (Module a) where snip bOver mm = {-# SCC "snip[Module]" #-} let arities = aritiesOfModule mm body' = snipX bOver arities (moduleBody mm) [] in mm { moduleBody = body' } instance Snip (Exp a) where snip bOver x = {-# SCC "snip[Exp]" #-} snipX bOver emptyArities x [] -- | Convert an expression into A-normal form. snipX :: Ord n => Bool -- ^ Introduce extra bindings for over-applied functions. -> Arities n -- ^ Arities of functions in environment. -> Exp a n -- ^ Expression to transform. -> [(Exp a n, a)] -- ^ Arguments being applied to current expression. -> Exp a n snipX bOver arities x args -- For applications, remember the argument that the function is being -- applied to, and decend into the function part. -- This unzips application nodes as we decend into the tree. | XApp a fun arg <- x = snipX bOver arities fun $ (snipX bOver arities arg [], a) : args -- Some non-application node with no arguments. | null args = enterX bOver arities x -- Some non-application node being applied to arguments. | otherwise = buildNormalisedApp bOver arities (enterX bOver arities x) args -- Enter into a non-application. enterX bOver arities xx = let down ars e = snipX bOver (extendsArities arities ars) e [] in case xx of -- The snipX function shouldn't have called us with an XApp. XApp{} -> error "DDC.Core.Transform.Snip: snipX shouldn't give us an XApp" -- leafy constructors XVar{} -> xx XCon{} -> xx XType{} -> xx XWitness{} -> xx -- lambdas XLAM a b e -> XLAM a b (down [(b,0)] e) XLam a b e -> XLam a b (down [(b,0)] e) -- non-recursive let XLet a (LLet m b x1) x2 -> let x1' = down [] x1 x2' = down [(b, arityOfExp' x1')] x2 in XLet a (LLet m b x1') x2' -- recursive let XLet a (LRec lets) x2 -> let bs = map fst lets xs = map snd lets ars = zip bs (map arityOfExp' xs) xs' = map (down ars) xs x2' = down ars x2 in XLet a (LRec $ zip bs xs') x2' -- letregion, just make sure we record bindings with dummy val. XLet a (LLetRegions b bs) x2 -> let ars = zip bs (repeat 0) in XLet a (LLetRegions b bs) (down ars x2) -- withregion XLet a (LWithRegion b) z2 -> XLet a (LWithRegion b) (down [] z2) -- case -- Split out non-atomic discriminants into their own bindings. XCase a e alts | isAtom e -> let e' = down [] e alts' = map (\(AAlt pat ae) -> AAlt pat (down (aritiesOfPat pat) ae)) alts in XCase a e' alts' | otherwise -> let e' = down [] e alts' = [AAlt pat (down (aritiesOfPat pat) ae) | AAlt pat ae <- alts] in XLet a (LLet LetStrict (BAnon (T.tBot T.kData)) e') (XCase a (XVar a $ UIx 0) (map (L.liftX 1) alts')) -- cast XCast a c e -> XCast a c (down [] e) -- | Build an A-normalised application of some functional expression to -- its arguments. Atomic arguments are applied directly, while -- on-atomic arguments are bound via let-expressions, then the -- associated let-bound variable is passed to the function. buildNormalisedApp :: Ord n => Bool -- ^ Introduce extra bindings for over-applied functions. -> Arities n -- ^ environment, arities of bound variables -> Exp a n -- ^ function -> [(Exp a n,a)] -- ^ arguments being applied to current expression -> Exp a n buildNormalisedApp _bOver _ f0 [] = f0 buildNormalisedApp bOver arities f0 args@( (_, annot) : _) = make annot f0 args where tBot' = T.tBot T.kData -- Lookup the arity of the function. f0Arity = case f0 of XVar _ b | Just arity <- getArity arities b -> max arity 1 _ -> max (arityOfExp' f0) 1 -- Make a normalised function application. make a xFun xsArgs -- The function part is already atomic. | isAtom xFun = buildNormalisedFunApp bOver a f0Arity xFun xsArgs -- The function part is not atomic, -- so we need to add an outer-most let-binding for it. | otherwise = XLet a (LLet LetStrict (BAnon tBot') xFun) (buildNormalisedFunApp bOver a f0Arity (XVar a (UIx 0)) [ (L.liftX 1 x, a') | (x, a') <- xsArgs]) -- | Build an A-normalised application of some functional expression to -- its arguments. Atomic arguments are applied directly, while -- on-atomic arguments are bound via let-expressions, then the -- associated let-bound variable is passed to the function. -- -- Unlike the `buildNormalisedFunApp` function above, this one -- wants the function part to be normalised as well. buildNormalisedFunApp :: Ord n => Bool -- ^ Introduce extra bindings for over-applied functions. -> a -- ^ Annotation to use. -> Int -- ^ Arity of the function part. -> Exp a n -- ^ Function part. -> [(Exp a n, a)] -- ^ Arguments to apply -> Exp a n buildNormalisedFunApp bOver an funArity xFun xsArgs = let tBot' = T.tBot T.kData -- Split arguments into the already atomic ones, -- and the ones we need to introduce let-expressions for. argss = splitArgs xsArgs -- Collect up the new let-bindings. xsLets = [ (x, a) | (_, a, _, Just x) <- argss] -- The total number of new let-bindings. nLets = length xsLets -- Lift indices in each binding over the bindings before it. xsLets' = [ (L.liftX n x, a) | (x, a) <- xsLets | (n :: Int) <- [0..] ] -- Lift indices in the function over the bindings before it. xFun' = L.liftX nLets xFun -- Collect up the new function arguments. -- If the argument was already atomic then we have to lift -- its indices past the new let bindings we're about to add. -- Otherwise it's a reference to one of the bindings directly. xsArgs' = [if liftMe then (L.liftX nLets xArg, a) else (xArg, a) | (xArg, a, liftMe, _) <- argss] -- Construct the new function application. xFunApps -- If the function is over-applied then create an intermediate -- binding that saturates it, then apply the extra arguments -- separately. | bOver , length xsArgs' > funArity , (xsSat, xsOver) <- splitAt funArity xsArgs' = XLet an (LLet LetStrict (BAnon tBot') (makeXAppsWithAnnots xFun' xsSat)) (makeXAppsWithAnnots (XVar an (UIx 0)) [ (L.liftX 1 x, a) | (x, a) <- xsOver ]) -- Function has the correct number of arguments, -- or is partially applied. | otherwise = makeXAppsWithAnnots xFun' xsArgs' -- Wrap the function application in the let-bindings -- for its arguments. in foldr (\(x, a) x' -> XLet a x x') xFunApps [ (LLet LetStrict (BAnon tBot') x, a) | (x, a) <- xsLets' ] -- | Sort function arguments into either the atomic ones, -- or compound ones. splitArgs :: Ord n => [(Exp a n, a)] -> [( Exp a n -- Expression to use as the new argument. , a -- Annoation for the argument application. , Bool -- Whether this argument was already atomic. , Maybe (Exp a n))] -- New expression to let-bind. splitArgs args = reverse $ go 0 $ reverse args where go _n [] = [] go n ((xArg, a) : xsArgs) | isAtom xArg = (xArg, a, True, Nothing) : go n xsArgs | otherwise = (XVar a (UIx n), a, False, Just xArg) : go (n + 1) xsArgs -- | Check if an expression needs a binding, or if it's simple enough to be -- applied as-is. isAtom :: Ord n => Exp a n -> Bool isAtom xx = case xx of XVar{} -> True XCon{} -> True XType{} -> True XWitness{} -> True -- Casts are ignored by code generator, so we can leave them in if -- their subexpression is normal XCast _ _ x -> isAtom x _ -> False -- | Take the arity of an expression, -- returning 0 for XType and XWitness. arityOfExp' :: Ord n => Exp a n -> Int arityOfExp' xx = case arityOfExp xx of Nothing -> 0 Just a -> a