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