1-- | This module implements the "Magic Do" optimization, which inlines calls to return 2-- and bind for the Eff monad, as well as some of its actions. 3module Language.PureScript.CoreImp.Optimizer.MagicDo (magicDoEffect, magicDoEff, magicDoST, inlineST) where 4 5import Prelude.Compat 6import Protolude (ordNub) 7 8import Data.Maybe (fromJust, isJust) 9import Data.Text (Text) 10 11import Language.PureScript.CoreImp.AST 12import Language.PureScript.CoreImp.Optimizer.Common 13import Language.PureScript.PSString (mkString) 14import qualified Language.PureScript.Constants.Prelude as C 15 16-- | Inline type class dictionaries for >>= and return for the Eff monad 17-- 18-- E.g. 19-- 20-- Prelude[">>="](dict)(m1)(function(x) { 21-- return ...; 22-- }) 23-- 24-- becomes 25-- 26-- function __do { 27-- var x = m1(); 28-- ... 29-- } 30magicDoEff :: AST -> AST 31magicDoEff = magicDo C.eff C.effDictionaries 32 33magicDoEffect :: AST -> AST 34magicDoEffect = magicDo C.effect C.effectDictionaries 35 36magicDoST :: AST -> AST 37magicDoST = magicDo C.st C.stDictionaries 38 39magicDo :: Text -> C.EffectDictionaries -> AST -> AST 40magicDo effectModule C.EffectDictionaries{..} = everywhereTopDown convert 41 where 42 -- The name of the function block which is added to denote a do block 43 fnName = "__do" 44 -- Desugar monomorphic calls to >>= and return for the Eff monad 45 convert :: AST -> AST 46 -- Desugar pure 47 convert (App _ (App _ pure' [val]) []) | isPure pure' = val 48 -- Desugar discard 49 convert (App _ (App _ bind [m]) [Function s1 Nothing [] (Block s2 js)]) | isDiscard bind = 50 Function s1 (Just fnName) [] $ Block s2 (App s2 m [] : map applyReturns js ) 51 -- Desugar bind to wildcard 52 convert (App _ (App _ bind [m]) [Function s1 Nothing [] (Block s2 js)]) 53 | isBind bind = 54 Function s1 (Just fnName) [] $ Block s2 (App s2 m [] : map applyReturns js ) 55 -- Desugar bind 56 convert (App _ (App _ bind [m]) [Function s1 Nothing [arg] (Block s2 js)]) | isBind bind = 57 Function s1 (Just fnName) [] $ Block s2 (VariableIntroduction s2 arg (Just (App s2 m [])) : map applyReturns js) 58 -- Desugar untilE 59 convert (App s1 (App _ f [arg]) []) | isEffFunc edUntil f = 60 App s1 (Function s1 Nothing [] (Block s1 [ While s1 (Unary s1 Not (App s1 arg [])) (Block s1 []), Return s1 $ ObjectLiteral s1 []])) [] 61 -- Desugar whileE 62 convert (App _ (App _ (App s1 f [arg1]) [arg2]) []) | isEffFunc edWhile f = 63 App s1 (Function s1 Nothing [] (Block s1 [ While s1 (App s1 arg1 []) (Block s1 [ App s1 arg2 [] ]), Return s1 $ ObjectLiteral s1 []])) [] 64 -- Inline __do returns 65 convert (Return _ (App _ (Function _ (Just ident) [] body) [])) | ident == fnName = body 66 -- Inline double applications 67 convert (App _ (App s1 (Function s2 Nothing [] (Block ss body)) []) []) = 68 App s1 (Function s2 Nothing [] (Block ss (applyReturns `fmap` body))) [] 69 convert other = other 70 -- Check if an expression represents a monomorphic call to >>= for the Eff monad 71 isBind (App _ fn [dict]) | isDict (effectModule, edBindDict) dict && isBindPoly fn = True 72 isBind _ = False 73 -- Check if an expression represents a call to @discard@ 74 isDiscard (App _ (App _ fn [dict1]) [dict2]) 75 | isDict (C.controlBind, C.discardUnitDictionary) dict1 && 76 isDict (effectModule, edBindDict) dict2 && 77 isDiscardPoly fn = True 78 isDiscard _ = False 79 -- Check if an expression represents a monomorphic call to pure or return for the Eff applicative 80 isPure (App _ fn [dict]) | isDict (effectModule, edApplicativeDict) dict && isPurePoly fn = True 81 isPure _ = False 82 -- Check if an expression represents the polymorphic >>= function 83 isBindPoly = isDict (C.controlBind, C.bind) 84 -- Check if an expression represents the polymorphic pure function 85 isPurePoly = isDict (C.controlApplicative, C.pure') 86 -- Check if an expression represents the polymorphic discard function 87 isDiscardPoly = isDict (C.controlBind, C.discard) 88 -- Check if an expression represents a function in the Effect module 89 isEffFunc name (Indexer _ (StringLiteral _ name') (Var _ eff)) = eff == effectModule && name == name' 90 isEffFunc _ _ = False 91 92 applyReturns :: AST -> AST 93 applyReturns (Return ss ret) = Return ss (App ss ret []) 94 applyReturns (Block ss jss) = Block ss (map applyReturns jss) 95 applyReturns (While ss cond js) = While ss cond (applyReturns js) 96 applyReturns (For ss v lo hi js) = For ss v lo hi (applyReturns js) 97 applyReturns (ForIn ss v xs js) = ForIn ss v xs (applyReturns js) 98 applyReturns (IfElse ss cond t f) = IfElse ss cond (applyReturns t) (applyReturns `fmap` f) 99 applyReturns other = other 100 101-- | Inline functions in the ST module 102inlineST :: AST -> AST 103inlineST = everywhere convertBlock 104 where 105 -- Look for runST blocks and inline the STRefs there. 106 -- If all STRefs are used in the scope of the same runST, only using { read, write, modify }STRef then 107 -- we can be more aggressive about inlining, and actually turn STRefs into local variables. 108 convertBlock (App s1 f [arg]) | isSTFunc C.runST f = 109 let refs = ordNub . findSTRefsIn $ arg 110 usages = findAllSTUsagesIn arg 111 allUsagesAreLocalVars = all (\u -> let v = toVar u in isJust v && fromJust v `elem` refs) usages 112 localVarsDoNotEscape = all (\r -> length (r `appearingIn` arg) == length (filter (\u -> let v = toVar u in v == Just r) usages)) refs 113 in App s1 (everywhere (convert (allUsagesAreLocalVars && localVarsDoNotEscape)) arg) [] 114 convertBlock other = other 115 -- Convert a block in a safe way, preserving object wrappers of references, 116 -- or in a more aggressive way, turning wrappers into local variables depending on the 117 -- agg(ressive) parameter. 118 convert agg (App s1 f [arg]) | isSTFunc C.newSTRef f = 119 Function s1 Nothing [] (Block s1 [Return s1 $ if agg then arg else ObjectLiteral s1 [(mkString C.stRefValue, arg)]]) 120 convert agg (App _ (App s1 f [ref]) []) | isSTFunc C.readSTRef f = 121 if agg then ref else Indexer s1 (StringLiteral s1 C.stRefValue) ref 122 convert agg (App _ (App _ (App s1 f [arg]) [ref]) []) | isSTFunc C.writeSTRef f = 123 if agg then Assignment s1 ref arg else Assignment s1 (Indexer s1 (StringLiteral s1 C.stRefValue) ref) arg 124 convert agg (App _ (App _ (App s1 f [func]) [ref]) []) | isSTFunc C.modifySTRef f = 125 if agg then Assignment s1 ref (App s1 func [ref]) else Assignment s1 (Indexer s1 (StringLiteral s1 C.stRefValue) ref) (App s1 func [Indexer s1 (StringLiteral s1 C.stRefValue) ref]) 126 convert _ other = other 127 -- Check if an expression represents a function in the ST module 128 isSTFunc name (Indexer _ (StringLiteral _ name') (Var _ st)) = st == C.st && name == name' 129 isSTFunc _ _ = False 130 -- Find all ST Refs initialized in this block 131 findSTRefsIn = everything (++) isSTRef 132 where 133 isSTRef (VariableIntroduction _ ident (Just (App _ (App _ f [_]) []))) | isSTFunc C.newSTRef f = [ident] 134 isSTRef _ = [] 135 -- Find all STRefs used as arguments to readSTRef, writeSTRef, modifySTRef 136 findAllSTUsagesIn = everything (++) isSTUsage 137 where 138 isSTUsage (App _ (App _ f [ref]) []) | isSTFunc C.readSTRef f = [ref] 139 isSTUsage (App _ (App _ (App _ f [_]) [ref]) []) | isSTFunc C.writeSTRef f || isSTFunc C.modifySTRef f = [ref] 140 isSTUsage _ = [] 141 -- Find all uses of a variable 142 appearingIn ref = everything (++) isVar 143 where 144 isVar e@(Var _ v) | v == ref = [e] 145 isVar _ = [] 146 -- Convert a AST value to a String if it is a Var 147 toVar (Var _ v) = Just v 148 toVar _ = Nothing 149