module Language.PureScript.CoreFn.Optimizer (optimizeCoreFn) where

import Protolude hiding (Type, moduleName)

import Control.Monad.Supply (Supply)
import Data.List (lookup)
import Language.PureScript.AST.Literals
import Language.PureScript.AST.SourcePos
import Language.PureScript.CoreFn.Ann
import Language.PureScript.CoreFn.CSE
import Language.PureScript.CoreFn.Expr
import Language.PureScript.CoreFn.Module
import Language.PureScript.CoreFn.Traversals
import Language.PureScript.Label
import Language.PureScript.Types
import qualified Language.PureScript.Constants.Libs as C
import qualified Language.PureScript.Constants.Prim as C

-- |
-- CoreFn optimization pass.
--
optimizeCoreFn :: Module Ann -> Supply (Module Ann)
optimizeCoreFn :: Module Ann -> Supply (Module Ann)
optimizeCoreFn Module Ann
m = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\[Bind Ann]
md -> Module Ann
m {moduleDecls :: [Bind Ann]
moduleDecls = [Bind Ann]
md}) forall b c a. (b -> c) -> (a -> b) -> a -> c
. ModuleName -> [Bind Ann] -> SupplyT Identity [Bind Ann]
optimizeCommonSubexpressions (forall a. Module a -> ModuleName
moduleName Module Ann
m) forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Bind Ann] -> [Bind Ann]
optimizeModuleDecls forall a b. (a -> b) -> a -> b
$ forall a. Module a -> [Bind a]
moduleDecls Module Ann
m

optimizeModuleDecls :: [Bind Ann] -> [Bind Ann]
optimizeModuleDecls :: [Bind Ann] -> [Bind Ann]
optimizeModuleDecls = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map Bind Ann -> Bind Ann
transformBinds
  where
  (Bind Ann -> Bind Ann
transformBinds, Expr Ann -> Expr Ann
_, Binder Ann -> Binder Ann
_) = forall a.
(Bind a -> Bind a)
-> (Expr a -> Expr a)
-> (Binder a -> Binder a)
-> (Bind a -> Bind a, Expr a -> Expr a, Binder a -> Binder a)
everywhereOnValues forall a. a -> a
identity Expr Ann -> Expr Ann
transformExprs forall a. a -> a
identity
  transformExprs :: Expr Ann -> Expr Ann
transformExprs
    = Expr Ann -> Expr Ann
optimizeClosedRecordUpdate
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Expr a -> Expr a
optimizeDataFunctionApply

optimizeClosedRecordUpdate :: Expr Ann -> Expr Ann
optimizeClosedRecordUpdate :: Expr Ann -> Expr Ann
optimizeClosedRecordUpdate ou :: Expr Ann
ou@(ObjectUpdate a :: Ann
a@(SourceSpan
_, [Comment]
_, Just SourceType
t, Maybe Meta
_) Expr Ann
r [(PSString, Expr Ann)]
updatedFields) =
  case forall a. Type a -> Maybe [Label]
closedRecordFields SourceType
t of
    Maybe [Label]
Nothing -> Expr Ann
ou
    Just [Label]
allFields -> forall a. a -> Literal (Expr a) -> Expr a
Literal Ann
a (forall a. [(PSString, a)] -> Literal a
ObjectLiteral (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
map Label -> (PSString, Expr Ann)
f [Label]
allFields))
      where f :: Label -> (PSString, Expr Ann)
f (Label PSString
l) = case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup PSString
l [(PSString, Expr Ann)]
updatedFields of
              Maybe (Expr Ann)
Nothing -> (PSString
l, forall a. a -> PSString -> Expr a -> Expr a
Accessor (SourceSpan
nullSourceSpan, [], forall a. Maybe a
Nothing, forall a. Maybe a
Nothing) PSString
l Expr Ann
r)
              Just Expr Ann
e -> (PSString
l, Expr Ann
e)
optimizeClosedRecordUpdate Expr Ann
e = Expr Ann
e

-- | Return the labels of a closed record, or Nothing for other types or open records.
closedRecordFields :: Type a -> Maybe [Label]
closedRecordFields :: forall a. Type a -> Maybe [Label]
closedRecordFields (TypeApp a
_ (TypeConstructor a
_ Qualified (ProperName 'TypeName)
C.Record) Type a
row) =
  forall a. Type a -> Maybe [Label]
collect Type a
row
  where
    collect :: Type a -> Maybe [Label]
    collect :: forall a. Type a -> Maybe [Label]
collect (REmptyKinded a
_ Maybe (Type a)
_) = forall a. a -> Maybe a
Just []
    collect (RCons a
_ Label
l Type a
_ Type a
r) = (Label
l forall a. a -> [a] -> [a]
:) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Type a -> Maybe [Label]
collect Type a
r
    collect Type a
_ = forall a. Maybe a
Nothing
closedRecordFields Type a
_ = forall a. Maybe a
Nothing

optimizeDataFunctionApply :: Expr a -> Expr a
optimizeDataFunctionApply :: forall a. Expr a -> Expr a
optimizeDataFunctionApply Expr a
e = case Expr a
e of
  (App a
a (App a
_ (Var a
_ Qualified Ident
fn) Expr a
x) Expr a
y)
    | Qualified Ident
C.I_functionApply <- Qualified Ident
fn -> forall a. a -> Expr a -> Expr a -> Expr a
App a
a Expr a
x Expr a
y
    | Qualified Ident
C.I_functionApplyFlipped <- Qualified Ident
fn -> forall a. a -> Expr a -> Expr a -> Expr a
App a
a Expr a
y Expr a
x
  Expr a
_ -> Expr a
e