-- | This module implements the "Magic Do" optimization, which inlines calls to return
-- and bind for the Eff monad, as well as some of its actions.
module Language.PureScript.CoreImp.Optimizer.MagicDo (magicDoEffect, magicDoEff, magicDoST, inlineST) where

import Prelude
import Protolude (ordNub)

import Data.Maybe (fromJust, isJust)

import Language.PureScript.CoreImp.AST
import Language.PureScript.CoreImp.Optimizer.Common
import Language.PureScript.Names (ModuleName)
import Language.PureScript.PSString (mkString)
import qualified Language.PureScript.Constants.Libs as C

-- | Inline type class dictionaries for >>= and return for the Eff monad
--
-- E.g.
--
--  Prelude[">>="](dict)(m1)(function(x) {
--    return ...;
--  })
--
-- becomes
--
--  function __do {
--    var x = m1();
--    ...
--  }
magicDoEff :: (AST -> AST) -> AST -> AST
magicDoEff :: (AST -> AST) -> AST -> AST
magicDoEff = ModuleName -> EffectDictionaries -> (AST -> AST) -> AST -> AST
magicDo ModuleName
C.M_Control_Monad_Eff EffectDictionaries
C.effDictionaries

magicDoEffect :: (AST -> AST) -> AST -> AST
magicDoEffect :: (AST -> AST) -> AST -> AST
magicDoEffect = ModuleName -> EffectDictionaries -> (AST -> AST) -> AST -> AST
magicDo ModuleName
C.M_Effect EffectDictionaries
C.effectDictionaries

magicDoST :: (AST -> AST) -> AST -> AST
magicDoST :: (AST -> AST) -> AST -> AST
magicDoST = ModuleName -> EffectDictionaries -> (AST -> AST) -> AST -> AST
magicDo ModuleName
C.M_Control_Monad_ST_Internal EffectDictionaries
C.stDictionaries

magicDo :: ModuleName -> C.EffectDictionaries -> (AST -> AST) -> AST -> AST
magicDo :: ModuleName -> EffectDictionaries -> (AST -> AST) -> AST -> AST
magicDo ModuleName
effectModule C.EffectDictionaries{PSString
edUntil :: EffectDictionaries -> PSString
edWhile :: EffectDictionaries -> PSString
edMonadDict :: EffectDictionaries -> PSString
edBindDict :: EffectDictionaries -> PSString
edApplicativeDict :: EffectDictionaries -> PSString
edUntil :: PSString
edWhile :: PSString
edMonadDict :: PSString
edBindDict :: PSString
edApplicativeDict :: PSString
..} AST -> AST
expander = (AST -> AST) -> AST -> AST
everywhereTopDown AST -> AST
convert
  where
  -- The name of the function block which is added to denote a do block
  fnName :: Text
fnName = Text
"__do"
  -- Desugar monomorphic calls to >>= and return for the Eff monad
  convert :: AST -> AST
  -- Desugar pure
  convert :: AST -> AST
convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ AST
pure' [AST
val]) []) | AST -> Bool
isPure AST
pure' = AST
val
  -- Desugar discard
  convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ AST
bind [AST
m]) [Function Maybe SourceSpan
s1 Maybe Text
Nothing [] (Block Maybe SourceSpan
s2 [AST]
js)]) | AST -> Bool
isDiscard AST
bind =
    Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 (forall a. a -> Maybe a
Just Text
fnName) [] forall a b. (a -> b) -> a -> b
$ Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s2 (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s2 AST
m [] forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map AST -> AST
applyReturns [AST]
js )
  -- Desugar bind to wildcard
  convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ AST
bind [AST
m]) [Function Maybe SourceSpan
s1 Maybe Text
Nothing [] (Block Maybe SourceSpan
s2 [AST]
js)])
    | AST -> Bool
isBind AST
bind =
    Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 (forall a. a -> Maybe a
Just Text
fnName) [] forall a b. (a -> b) -> a -> b
$ Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s2 (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s2 AST
m [] forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map AST -> AST
applyReturns [AST]
js )
  -- Desugar bind
  convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ AST
bind [AST
m]) [Function Maybe SourceSpan
s1 Maybe Text
Nothing [Text
arg] (Block Maybe SourceSpan
s2 [AST]
js)]) | AST -> Bool
isBind AST
bind =
    Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 (forall a. a -> Maybe a
Just Text
fnName) [] forall a b. (a -> b) -> a -> b
$ Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s2 (Maybe SourceSpan -> Text -> Maybe (InitializerEffects, AST) -> AST
VariableIntroduction Maybe SourceSpan
s2 Text
arg (forall a. a -> Maybe a
Just (InitializerEffects
UnknownEffects, Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s2 AST
m [])) forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map AST -> AST
applyReturns [AST]
js)
  -- Desugar untilE
  convert (App Maybe SourceSpan
s1 (App Maybe SourceSpan
_ AST
f [AST
arg]) []) | PSString -> AST -> Bool
isEffFunc PSString
edUntil AST
f =
    Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 (Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 forall a. Maybe a
Nothing [] (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s1 [ Maybe SourceSpan -> AST -> AST -> AST
While Maybe SourceSpan
s1 (Maybe SourceSpan -> UnaryOperator -> AST -> AST
Unary Maybe SourceSpan
s1 UnaryOperator
Not (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 AST
arg [])) (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s1 []), Maybe SourceSpan -> AST -> AST
Return Maybe SourceSpan
s1 forall a b. (a -> b) -> a -> b
$ Maybe SourceSpan -> [(PSString, AST)] -> AST
ObjectLiteral Maybe SourceSpan
s1 []])) []
  -- Desugar whileE
  convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (App Maybe SourceSpan
s1 AST
f [AST
arg1]) [AST
arg2]) []) | PSString -> AST -> Bool
isEffFunc PSString
edWhile AST
f =
    Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 (Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 forall a. Maybe a
Nothing [] (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s1 [ Maybe SourceSpan -> AST -> AST -> AST
While Maybe SourceSpan
s1 (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 AST
arg1 []) (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s1 [ Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 AST
arg2 [] ]), Maybe SourceSpan -> AST -> AST
Return Maybe SourceSpan
s1 forall a b. (a -> b) -> a -> b
$ Maybe SourceSpan -> [(PSString, AST)] -> AST
ObjectLiteral Maybe SourceSpan
s1 []])) []
  -- Inline __do returns
  convert (Return Maybe SourceSpan
_ (App Maybe SourceSpan
_ (Function Maybe SourceSpan
_ (Just Text
ident) [] AST
body) [])) | Text
ident forall a. Eq a => a -> a -> Bool
== Text
fnName = AST
body
  -- Inline double applications
  convert (App Maybe SourceSpan
_ (App Maybe SourceSpan
s1 (Function Maybe SourceSpan
s2 Maybe Text
Nothing [] (Block Maybe SourceSpan
ss [AST]
body)) []) []) =
    Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 (Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s2 forall a. Maybe a
Nothing [] (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
ss (AST -> AST
applyReturns forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` [AST]
body))) []
  convert AST
other = AST
other
  -- Check if an expression represents a monomorphic call to >>= for the Eff monad
  isBind :: AST -> Bool
isBind (AST -> AST
expander -> App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
C.P_bind) [Ref (ModuleName, PSString)
dict]) = (ModuleName
effectModule, PSString
edBindDict) forall a. Eq a => a -> a -> Bool
== (ModuleName, PSString)
dict
  isBind AST
_ = Bool
False
  -- Check if an expression represents a call to @discard@
  isDiscard :: AST -> Bool
isDiscard (AST -> AST
expander -> App Maybe SourceSpan
_ (AST -> AST
expander -> App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
C.P_discard) [Ref (ModuleName, PSString)
C.P_discardUnit]) [Ref (ModuleName, PSString)
dict]) = (ModuleName
effectModule, PSString
edBindDict) forall a. Eq a => a -> a -> Bool
== (ModuleName, PSString)
dict
  isDiscard AST
_ = Bool
False
  -- Check if an expression represents a monomorphic call to pure or return for the Eff applicative
  isPure :: AST -> Bool
isPure (AST -> AST
expander -> App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
C.P_pure) [Ref (ModuleName, PSString)
dict]) = (ModuleName
effectModule, PSString
edApplicativeDict) forall a. Eq a => a -> a -> Bool
== (ModuleName, PSString)
dict
  isPure AST
_ = Bool
False
  -- Check if an expression represents a function in the Effect module
  isEffFunc :: PSString -> AST -> Bool
isEffFunc PSString
name (Ref (ModuleName, PSString)
fn) = (ModuleName
effectModule, PSString
name) forall a. Eq a => a -> a -> Bool
== (ModuleName, PSString)
fn
  isEffFunc PSString
_ AST
_ = Bool
False

  applyReturns :: AST -> AST
  applyReturns :: AST -> AST
applyReturns (Return Maybe SourceSpan
ss AST
ret) = Maybe SourceSpan -> AST -> AST
Return Maybe SourceSpan
ss (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
ss AST
ret [])
  applyReturns (Block Maybe SourceSpan
ss [AST]
jss) = Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
ss (forall a b. (a -> b) -> [a] -> [b]
map AST -> AST
applyReturns [AST]
jss)
  applyReturns (While Maybe SourceSpan
ss AST
cond AST
js) = Maybe SourceSpan -> AST -> AST -> AST
While Maybe SourceSpan
ss AST
cond (AST -> AST
applyReturns AST
js)
  applyReturns (For Maybe SourceSpan
ss Text
v AST
lo AST
hi AST
js) = Maybe SourceSpan -> Text -> AST -> AST -> AST -> AST
For Maybe SourceSpan
ss Text
v AST
lo AST
hi (AST -> AST
applyReturns AST
js)
  applyReturns (ForIn Maybe SourceSpan
ss Text
v AST
xs AST
js) = Maybe SourceSpan -> Text -> AST -> AST -> AST
ForIn Maybe SourceSpan
ss Text
v AST
xs (AST -> AST
applyReturns AST
js)
  applyReturns (IfElse Maybe SourceSpan
ss AST
cond AST
t Maybe AST
f) = Maybe SourceSpan -> AST -> AST -> Maybe AST -> AST
IfElse Maybe SourceSpan
ss AST
cond (AST -> AST
applyReturns AST
t) (AST -> AST
applyReturns forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Maybe AST
f)
  applyReturns AST
other = AST
other

-- | Inline functions in the ST module
inlineST :: AST -> AST
inlineST :: AST -> AST
inlineST = (AST -> AST) -> AST -> AST
everywhere AST -> AST
convertBlock
  where
  -- Look for run blocks and inline the STRefs there.
  -- If all STRefs are used in the scope of the same run, only using { read, write, modify } then
  -- we can be more aggressive about inlining, and actually turn STRefs into local variables.
  convertBlock :: AST -> AST
convertBlock (App Maybe SourceSpan
s1 (Ref (ModuleName, PSString)
C.P_run) [AST
arg]) =
    let refs :: [Text]
refs = forall a. Ord a => [a] -> [a]
ordNub forall b c a. (b -> c) -> (a -> b) -> a -> c
. AST -> [Text]
findSTRefsIn forall a b. (a -> b) -> a -> b
$ AST
arg
        usages :: [AST]
usages = AST -> [AST]
findAllSTUsagesIn AST
arg
        allUsagesAreLocalVars :: Bool
allUsagesAreLocalVars = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\AST
u -> let v :: Maybe Text
v = AST -> Maybe Text
toVar AST
u in forall a. Maybe a -> Bool
isJust Maybe Text
v Bool -> Bool -> Bool
&& forall a. HasCallStack => Maybe a -> a
fromJust Maybe Text
v forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
refs) [AST]
usages
        localVarsDoNotEscape :: Bool
localVarsDoNotEscape = forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Text
r -> forall (t :: * -> *) a. Foldable t => t a -> Int
length (Text
r Text -> AST -> [AST]
`appearingIn` AST
arg) forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall a. (a -> Bool) -> [a] -> [a]
filter (\AST
u -> let v :: Maybe Text
v = AST -> Maybe Text
toVar AST
u in Maybe Text
v forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just Text
r) [AST]
usages)) [Text]
refs
    in Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 ((AST -> AST) -> AST -> AST
everywhere (Bool -> AST -> AST
convert (Bool
allUsagesAreLocalVars Bool -> Bool -> Bool
&& Bool
localVarsDoNotEscape)) AST
arg) []
  convertBlock AST
other = AST
other
  -- Convert a block in a safe way, preserving object wrappers of references,
  -- or in a more aggressive way, turning wrappers into local variables depending on the
  -- agg(ressive) parameter.
  convert :: Bool -> AST -> AST
convert Bool
agg (App Maybe SourceSpan
s1 (Ref (ModuleName, PSString)
C.P_new) [AST
arg]) =
   Maybe SourceSpan -> Maybe Text -> [Text] -> AST -> AST
Function Maybe SourceSpan
s1 forall a. Maybe a
Nothing [] (Maybe SourceSpan -> [AST] -> AST
Block Maybe SourceSpan
s1 [Maybe SourceSpan -> AST -> AST
Return Maybe SourceSpan
s1 forall a b. (a -> b) -> a -> b
$ if Bool
agg then AST
arg else Maybe SourceSpan -> [(PSString, AST)] -> AST
ObjectLiteral Maybe SourceSpan
s1 [(Text -> PSString
mkString forall a. IsString a => a
C.stRefValue, AST
arg)]])
  convert Bool
agg (App Maybe SourceSpan
_ (App Maybe SourceSpan
s1 (Ref (ModuleName, PSString)
C.P_read) [AST
ref]) []) =
    if Bool
agg then AST
ref else Maybe SourceSpan -> AST -> AST -> AST
Indexer Maybe SourceSpan
s1 (Maybe SourceSpan -> PSString -> AST
StringLiteral Maybe SourceSpan
s1 forall a. IsString a => a
C.stRefValue) AST
ref
  convert Bool
agg (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (App Maybe SourceSpan
s1 (Ref (ModuleName, PSString)
C.P_write) [AST
arg]) [AST
ref]) []) =
    if Bool
agg then Maybe SourceSpan -> AST -> AST -> AST
Assignment Maybe SourceSpan
s1 AST
ref AST
arg else Maybe SourceSpan -> AST -> AST -> AST
Assignment Maybe SourceSpan
s1 (Maybe SourceSpan -> AST -> AST -> AST
Indexer Maybe SourceSpan
s1 (Maybe SourceSpan -> PSString -> AST
StringLiteral Maybe SourceSpan
s1 forall a. IsString a => a
C.stRefValue) AST
ref) AST
arg
  convert Bool
agg (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (App Maybe SourceSpan
s1 (Ref (ModuleName, PSString)
C.P_modify) [AST
func]) [AST
ref]) []) =
    if Bool
agg then Maybe SourceSpan -> AST -> AST -> AST
Assignment Maybe SourceSpan
s1 AST
ref (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 AST
func [AST
ref]) else Maybe SourceSpan -> AST -> AST -> AST
Assignment Maybe SourceSpan
s1 (Maybe SourceSpan -> AST -> AST -> AST
Indexer Maybe SourceSpan
s1 (Maybe SourceSpan -> PSString -> AST
StringLiteral Maybe SourceSpan
s1 forall a. IsString a => a
C.stRefValue) AST
ref) (Maybe SourceSpan -> AST -> [AST] -> AST
App Maybe SourceSpan
s1 AST
func [Maybe SourceSpan -> AST -> AST -> AST
Indexer Maybe SourceSpan
s1 (Maybe SourceSpan -> PSString -> AST
StringLiteral Maybe SourceSpan
s1 forall a. IsString a => a
C.stRefValue) AST
ref])
  convert Bool
_ AST
other = AST
other
  -- Find all ST Refs initialized in this block
  findSTRefsIn :: AST -> [Text]
findSTRefsIn = forall r. (r -> r -> r) -> (AST -> r) -> AST -> r
everything forall a. [a] -> [a] -> [a]
(++) AST -> [Text]
isSTRef
    where
    isSTRef :: AST -> [Text]
isSTRef (VariableIntroduction Maybe SourceSpan
_ Text
ident (Just (InitializerEffects
_, App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
C.P_new) [AST
_]) []))) = [Text
ident]
    isSTRef AST
_ = []
  -- Find all STRefs used as arguments to read, write, modify
  findAllSTUsagesIn :: AST -> [AST]
findAllSTUsagesIn = forall r. (r -> r -> r) -> (AST -> r) -> AST -> r
everything forall a. [a] -> [a] -> [a]
(++) AST -> [AST]
isSTUsage
    where
    isSTUsage :: AST -> [AST]
isSTUsage (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
C.P_read) [AST
ref]) []) = [AST
ref]
    isSTUsage (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (App Maybe SourceSpan
_ (Ref (ModuleName, PSString)
f) [AST
_]) [AST
ref]) []) | (ModuleName, PSString)
f forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [forall a. (Eq a, IsString a) => (ModuleName, a)
C.P_write, forall a. (Eq a, IsString a) => (ModuleName, a)
C.P_modify] = [AST
ref]
    isSTUsage AST
_ = []
  -- Find all uses of a variable
  appearingIn :: Text -> AST -> [AST]
appearingIn Text
ref = forall r. (r -> r -> r) -> (AST -> r) -> AST -> r
everything forall a. [a] -> [a] -> [a]
(++) AST -> [AST]
isVar
    where
    isVar :: AST -> [AST]
isVar e :: AST
e@(Var Maybe SourceSpan
_ Text
v) | Text
v forall a. Eq a => a -> a -> Bool
== Text
ref = [AST
e]
    isVar AST
_ = []
  -- Convert a AST value to a String if it is a Var
  toVar :: AST -> Maybe Text
toVar (Var Maybe SourceSpan
_ Text
v) = forall a. a -> Maybe a
Just Text
v
  toVar AST
_ = forall a. Maybe a
Nothing