<commit_before>{-# LANGUAGE TemplateHaskell #-}

module Koriel.Core.Optimize
  ( optimizeProgram,
  )
where

import Control.Monad.Except
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import qualified Data.List as List
import Koriel.Core.Alpha
import Koriel.Core.Flat
import Koriel.Core.Syntax
import Koriel.Core.Type
import Koriel.Id
import Koriel.MonadUniq
import Koriel.Prelude

data OptimizeEnv = OptimizeEnv {_optimizeUniqSupply :: UniqSupply, _inlineLevel :: Int}

makeLenses ''OptimizeEnv

instance HasUniqSupply OptimizeEnv where
  uniqSupply = optimizeUniqSupply

-- | Apply a monadic function n times.
times :: Monad m => Int -> (t -> m t) -> t -> m t
times 0 _ x = pure x
times n f x
  | n > 0 = times (n - 1) f =<< f x
  | otherwise = bug $ Unreachable $ tshow n <> " must be a natural number"

-- | 最適化を行う関数
--
-- * 自明な変数の付け替えの簡約 (optVarBind)
-- * 値コンストラクタを含む関数のインライン展開 (optPackInline, optCallInline)
-- * 不要なletの削除 (removeUnusedLet)
-- * 無意味なcastの削除（optIdCast）
optimizeProgram ::
  MonadIO m =>
  UniqSupply ->
  -- | インライン展開する関数のサイズ
  Int ->
  Program (Id Type) ->
  m (Program (Id Type))
optimizeProgram us level prog@Program {..} = runReaderT ?? OptimizeEnv {_optimizeUniqSupply = us, _inlineLevel = level} $ do
  state <- execStateT (traverse (checkInlineable . uncurry LocalDef . over _2 (uncurry Fun)) _topFuncs) mempty
  appProgram (optimizeExpr state) prog

optimizeExpr :: (MonadReader OptimizeEnv f, MonadIO f) => CallInlineMap -> Exp (Id Type) -> f (Exp (Id Type))
optimizeExpr state = 10 `times` opt
  where
    opt =
      pure
        >=> optVarBind
        >=> (flip runReaderT mempty . optPackInline)
        >=> removeUnusedLet
        >=> (flip evalStateT state . optCallInline)
        >=> optIdCast
        -- >=> optCast
        >=> pure
          . flat

type CallInlineMap = HashMap (Id Type) ([Id Type], Exp (Id Type))

optCallInline ::
  (MonadState CallInlineMap f, MonadReader OptimizeEnv f, MonadIO f) =>
  Exp (Id Type) ->
  f (Exp (Id Type))
optCallInline (Call (Var f) xs) = lookupCallInline (Call . Var) f xs
optCallInline (CallDirect f xs) = lookupCallInline CallDirect f xs
optCallInline (Match v cs) =
  Match <$> optCallInline v <*> traverseOf (traversed . appCase) optCallInline cs
optCallInline (Let ds e) = do
  ds' <- traverseOf (traversed . localDefObj . appObj) optCallInline ds
  traverse_ checkInlineable ds'
  Let ds' <$> optCallInline e
optCallInline e = pure e

checkInlineable ::
  (MonadState CallInlineMap m, MonadReader OptimizeEnv m) =>
  LocalDef (Id Type) ->
  m ()
checkInlineable (LocalDef f (Fun ps v)) = do
  level <- view inlineLevel
  -- ノードの数がlevel以下ならインライン展開する
  when (lengthOf cosmos v <= level {-  || f `notElem` freevars v -}) $ at f ?= (ps, v)
checkInlineable _ = pure ()<commit_msg>Add optimization pass to remove trivial applications<commit_after>{-# LANGUAGE TemplateHaskell #-}

module Koriel.Core.Optimize
  ( optimizeProgram,
  )
where

import Control.Monad.Except
import qualified Data.HashMap.Strict as HashMap
import qualified Data.HashSet as HashSet
import qualified Data.List as List
import Koriel.Core.Alpha
import Koriel.Core.Flat
import Koriel.Core.Syntax
import Koriel.Core.Type
import Koriel.Id
import Koriel.MonadUniq
import Koriel.Prelude

data OptimizeEnv = OptimizeEnv {_optimizeUniqSupply :: UniqSupply, _inlineLevel :: Int}

makeLenses ''OptimizeEnv

instance HasUniqSupply OptimizeEnv where
  uniqSupply = optimizeUniqSupply

-- | Apply a monadic function n times.
times :: Monad m => Int -> (t -> m t) -> t -> m t
times 0 _ x = pure x
times n f x
  | n > 0 = times (n - 1) f =<< f x
  | otherwise = bug $ Unreachable $ tshow n <> " must be a natural number"

-- | 最適化を行う関数
--
-- * 自明な変数の付け替えの簡約 (optVarBind)
-- * 値コンストラクタを含む、関数のインライン展開 (optPackInline, optCallInline)
-- * 不要なletの削除 (removeUnusedLet)
-- * 無意味なcastの削除（optIdCast）
optimizeProgram ::
  MonadIO m =>
  UniqSupply ->
  -- | インライン展開する関数のサイズ
  Int ->
  Program (Id Type) ->
  m (Program (Id Type))
optimizeProgram us level prog@Program {..} = runReaderT ?? OptimizeEnv {_optimizeUniqSupply = us, _inlineLevel = level} $ do
  state <- execStateT (traverse (checkInlineable . uncurry LocalDef . over _2 (uncurry Fun)) _topFuncs) mempty
  appProgram (optimizeExpr state) prog

optimizeExpr :: (MonadReader OptimizeEnv f, MonadIO f) => CallInlineMap -> Exp (Id Type) -> f (Exp (Id Type))
optimizeExpr state = 10 `times` opt
  where
    opt =
      pure
        >=> optVarBind
        >=> (flip runReaderT mempty . optPackInline)
        >=> removeUnusedLet
        >=> (flip evalStateT state . optCallInline)
        >=> optIdCast
        >=> optTrivialCall
        -- >=> optCast
        >=> pure
          . flat

-- (let ((f (fun () body))) (f)) = body
optTrivialCall :: (Eq a, Applicative f) => Exp a -> f (Exp a)
optTrivialCall (Let [LocalDef f (Fun [] body)] (Call (Var f') [])) | f == f' = optTrivialCall body
optTrivialCall (Let ds e) = Let <$> traverseOf (traversed . localDefObj . appObj) optTrivialCall ds <*> optTrivialCall e
optTrivialCall (Match v cs) = Match <$> optTrivialCall v <*> traverseOf (traversed . appCase) optTrivialCall cs
optTrivialCall e = pure e

type CallInlineMap = HashMap (Id Type) ([Id Type], Exp (Id Type))

optCallInline ::
  (MonadState CallInlineMap f, MonadReader OptimizeEnv f, MonadIO f) =>
  Exp (Id Type) ->
  f (Exp (Id Type))
optCallInline (Call (Var f) xs) = lookupCallInline (Call . Var) f xs
optCallInline (CallDirect f xs) = lookupCallInline CallDirect f xs
optCallInline (Match v cs) =
  Match <$> optCallInline v <*> traverseOf (traversed . appCase) optCallInline cs
optCallInline (Let ds e) = do
  ds' <- traverseOf (traversed . localDefObj . appObj) optCallInline ds
  traverse_ checkInlineable ds'
  Let ds' <$> optCallInline e
optCallInline e = pure e

checkInlineable ::
  (MonadState CallInlineMap m, MonadReader OptimizeEnv m) =>
  LocalDef (Id Type) ->
  m ()
checkInlineable (LocalDef f (Fun ps v)) = do
  level <- view inlineLevel
  -- ノードの数がlevel以下ならインライン展開する
  when (lengthOf cosmos v <= level {-  || f `notElem` freevars v -}) $ at f ?= (ps, v)
checkInlineable _ = pure ()
{'fraction_non_alphanumeric': 0.11274030628869339, 'fraction_numerical': 0.0019550342130987292, 'mean_word_length': 3.619262603461249, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2344647', 'n_tokens_mistral': 2245, 'n_tokens_neox': 2239, 'n_words': 659}