{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}

-- |
-- Module      : Jikka.Core.Convert.ConstantPropagation
-- Description : propagates something constants, for exprs which are computable with constant time. / 定数時間で計算できるような式についての、ある種の定数伝播をします。
-- Copyright   : (c) Kimiyuki Onaka, 2021
-- License     : Apache License 2.0
-- Maintainer  : kimiyuki95@gmail.com
-- Stability   : experimental
-- Portability : portable
module Jikka.Core.Convert.ConstantPropagation
  ( run,
    run',
  )
where

import qualified Data.Map as M
import Data.Maybe (fromMaybe)
import Jikka.Common.Error
import Jikka.Core.Language.Expr
import Jikka.Core.Language.Lint
import Jikka.Core.Language.Util

type Env = M.Map VarName Expr

runExpr :: Env -> Expr -> Expr
runExpr :: Env -> Expr -> Expr
runExpr Env
env = \case
  Var VarName
x -> Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe (VarName -> Expr
Var VarName
x) (VarName -> Env -> Maybe Expr
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup VarName
x Env
env)
  Lit Literal
lit -> Literal -> Expr
Lit Literal
lit
  App Expr
f Expr
e -> Expr -> Expr -> Expr
App (Env -> Expr -> Expr
runExpr Env
env Expr
f) (Env -> Expr -> Expr
runExpr Env
env Expr
e)
  Lam VarName
x Type
t Expr
body -> VarName -> Type -> Expr -> Expr
Lam VarName
x Type
t (Env -> Expr -> Expr
runExpr Env
env Expr
body)
  Let VarName
x Type
t Expr
e1 Expr
e2 ->
    let e1' :: Expr
e1' = Env -> Expr -> Expr
runExpr Env
env Expr
e1
     in if Expr -> Bool
isConstantTimeExpr Expr
e1'
          then Env -> Expr -> Expr
runExpr (VarName -> Expr -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VarName
x Expr
e1' Env
env) Expr
e2
          else VarName -> Type -> Expr -> Expr -> Expr
Let VarName
x Type
t Expr
e1' (Env -> Expr -> Expr
runExpr Env
env Expr
e2)

runToplevelExpr :: Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr :: Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr Env
env = \case
  ResultExpr Expr
e -> Expr -> ToplevelExpr
ResultExpr (Env -> Expr -> Expr
runExpr Env
env Expr
e)
  ToplevelLet VarName
x Type
t Expr
e ToplevelExpr
cont ->
    let e' :: Expr
e' = Env -> Expr -> Expr
runExpr Env
env Expr
e
     in if Expr -> Bool
isConstantTimeExpr Expr
e'
          then Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr (VarName -> Expr -> Env -> Env
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert VarName
x Expr
e' Env
env) ToplevelExpr
cont
          else VarName -> Type -> Expr -> ToplevelExpr -> ToplevelExpr
ToplevelLet VarName
x Type
t Expr
e' (Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr Env
env ToplevelExpr
cont)
  ToplevelLetRec VarName
f [(VarName, Type)]
args Type
ret Expr
body ToplevelExpr
cont ->
    VarName
-> [(VarName, Type)]
-> Type
-> Expr
-> ToplevelExpr
-> ToplevelExpr
ToplevelLetRec VarName
f [(VarName, Type)]
args Type
ret (Env -> Expr -> Expr
runExpr Env
env Expr
body) (Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr Env
env ToplevelExpr
cont)

run' :: Program -> Program
run' :: ToplevelExpr -> ToplevelExpr
run' = Env -> ToplevelExpr -> ToplevelExpr
runToplevelExpr Env
forall k a. Map k a
M.empty

-- | `run` does constant propagation.
-- This assumes that the program is alpha-converted.
--
-- For example, this converts the following:
--
-- > let x = 1
-- > in let f = fun y -> y
-- > in x + x + f(x)
--
-- to:
--
-- > let f = fun y -> y
-- > in 1 + 1 + f(1)
--
-- NOTE: this doesn't constant folding.
run :: MonadError Error m => Program -> m Program
run :: ToplevelExpr -> m ToplevelExpr
run ToplevelExpr
prog = String -> m ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. MonadError Error m => String -> m a -> m a
wrapError' String
"Jikka.Core.Convert.ConstantPropagation" (m ToplevelExpr -> m ToplevelExpr)
-> m ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ do
  m () -> m ()
forall (m :: * -> *) a. MonadError Error m => m a -> m a
precondition (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
ensureWellTyped ToplevelExpr
prog
  ToplevelExpr
prog <- ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return (ToplevelExpr -> m ToplevelExpr) -> ToplevelExpr -> m ToplevelExpr
forall a b. (a -> b) -> a -> b
$ ToplevelExpr -> ToplevelExpr
run' ToplevelExpr
prog
  m () -> m ()
forall (m :: * -> *) a. MonadError Error m => m a -> m a
postcondition (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    ToplevelExpr -> m ()
forall (m :: * -> *). MonadError Error m => ToplevelExpr -> m ()
ensureWellTyped ToplevelExpr
prog
  ToplevelExpr -> m ToplevelExpr
forall (m :: * -> *) a. Monad m => a -> m a
return ToplevelExpr
prog