-- Copyright (C) 2018 Jun Zhang <zhangjunphy[at]gmail[dot]com>
--
-- This file is a part of decafc.
--
-- decafc is free software: you can redistribute it and/or modify it under the
-- terms of the MIT (X11) License as described in the LICENSE file.
--
-- decafc is distributed in the hope that it will be useful, but WITHOUT ANY
-- WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE.  See the X11 license for more details.

module CFG.Optimizations.Optimizer where

import Data.Generics.Labels
import GHC.Generics (Generic)
import CFG.Types
import Control.Monad.Except
import Control.Monad.State
import Types
import Util.Graph qualified as G
import Control.Lens (use, uses, view, (%=), (%~), (&), (+=), (.=), (.~), (^.), _1, _2, _3)

data CFGOptimizerState = CFGOptimizerState
  { CFGOptimizerState -> CFG
cfg :: CFG
  } deriving ((forall x. CFGOptimizerState -> Rep CFGOptimizerState x)
-> (forall x. Rep CFGOptimizerState x -> CFGOptimizerState)
-> Generic CFGOptimizerState
forall x. Rep CFGOptimizerState x -> CFGOptimizerState
forall x. CFGOptimizerState -> Rep CFGOptimizerState x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. CFGOptimizerState -> Rep CFGOptimizerState x
from :: forall x. CFGOptimizerState -> Rep CFGOptimizerState x
$cto :: forall x. Rep CFGOptimizerState x -> CFGOptimizerState
to :: forall x. Rep CFGOptimizerState x -> CFGOptimizerState
Generic)

newtype CFGOptimizer a = CFGOptmizer
  { forall a.
CFGOptimizer a -> StateT CFGOptimizerState (Except CompileError) a
runOptimizer :: StateT CFGOptimizerState (Except CompileError) a
  } deriving ((forall a b. (a -> b) -> CFGOptimizer a -> CFGOptimizer b)
-> (forall a b. a -> CFGOptimizer b -> CFGOptimizer a)
-> Functor CFGOptimizer
forall a b. a -> CFGOptimizer b -> CFGOptimizer a
forall a b. (a -> b) -> CFGOptimizer a -> CFGOptimizer b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> CFGOptimizer a -> CFGOptimizer b
fmap :: forall a b. (a -> b) -> CFGOptimizer a -> CFGOptimizer b
$c<$ :: forall a b. a -> CFGOptimizer b -> CFGOptimizer a
<$ :: forall a b. a -> CFGOptimizer b -> CFGOptimizer a
Functor, Functor CFGOptimizer
Functor CFGOptimizer
-> (forall a. a -> CFGOptimizer a)
-> (forall a b.
    CFGOptimizer (a -> b) -> CFGOptimizer a -> CFGOptimizer b)
-> (forall a b c.
    (a -> b -> c)
    -> CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer c)
-> (forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b)
-> (forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer a)
-> Applicative CFGOptimizer
forall a. a -> CFGOptimizer a
forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer a
forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
forall a b.
CFGOptimizer (a -> b) -> CFGOptimizer a -> CFGOptimizer b
forall a b c.
(a -> b -> c) -> CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> CFGOptimizer a
pure :: forall a. a -> CFGOptimizer a
$c<*> :: forall a b.
CFGOptimizer (a -> b) -> CFGOptimizer a -> CFGOptimizer b
<*> :: forall a b.
CFGOptimizer (a -> b) -> CFGOptimizer a -> CFGOptimizer b
$cliftA2 :: forall a b c.
(a -> b -> c) -> CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer c
liftA2 :: forall a b c.
(a -> b -> c) -> CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer c
$c*> :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
*> :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
$c<* :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer a
<* :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer a
Applicative, Applicative CFGOptimizer
Applicative CFGOptimizer
-> (forall a b.
    CFGOptimizer a -> (a -> CFGOptimizer b) -> CFGOptimizer b)
-> (forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b)
-> (forall a. a -> CFGOptimizer a)
-> Monad CFGOptimizer
forall a. a -> CFGOptimizer a
forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
forall a b.
CFGOptimizer a -> (a -> CFGOptimizer b) -> CFGOptimizer b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b.
CFGOptimizer a -> (a -> CFGOptimizer b) -> CFGOptimizer b
>>= :: forall a b.
CFGOptimizer a -> (a -> CFGOptimizer b) -> CFGOptimizer b
$c>> :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
>> :: forall a b. CFGOptimizer a -> CFGOptimizer b -> CFGOptimizer b
$creturn :: forall a. a -> CFGOptimizer a
return :: forall a. a -> CFGOptimizer a
Monad, MonadError CompileError, MonadState CFGOptimizerState)

runOptimizerOnCFG :: CFGOptimizer () -> CFG -> Either [CompileError] CFG
runOptimizerOnCFG :: CFGOptimizer () -> CFG -> Either [CompileError] CFG
runOptimizerOnCFG CFGOptimizer ()
opt CFG
cfg =
  let initState :: CFGOptimizerState
initState = CFG -> CFGOptimizerState
CFGOptimizerState CFG
cfg
      result :: Either CompileError ((), CFGOptimizerState)
result = Except CompileError ((), CFGOptimizerState)
-> Either CompileError ((), CFGOptimizerState)
forall e a. Except e a -> Either e a
runExcept (Except CompileError ((), CFGOptimizerState)
 -> Either CompileError ((), CFGOptimizerState))
-> Except CompileError ((), CFGOptimizerState)
-> Either CompileError ((), CFGOptimizerState)
forall a b. (a -> b) -> a -> b
$ StateT CFGOptimizerState (Except CompileError) ()
-> CFGOptimizerState -> Except CompileError ((), CFGOptimizerState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (CFGOptimizer ()
-> StateT CFGOptimizerState (Except CompileError) ()
forall a.
CFGOptimizer a -> StateT CFGOptimizerState (Except CompileError) a
runOptimizer CFGOptimizer ()
opt) CFGOptimizerState
initState
  in case Either CompileError ((), CFGOptimizerState)
result of
    Left CompileError
err -> [CompileError] -> Either [CompileError] CFG
forall a b. a -> Either a b
Left [CompileError
err]
    Right (()
_, CFGOptimizerState CFG
cfg) -> CFG -> Either [CompileError] CFG
forall a b. b -> Either a b
Right CFG
cfg

getCFG :: CFGOptimizer CFG
getCFG :: CFGOptimizer CFG
getCFG = (CFGOptimizerState -> CFG) -> CFGOptimizer CFG
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CFGOptimizerState -> CFG
cfg

updateCFG :: G.GraphBuilder BBID BasicBlock CFGEdge a -> CFGOptimizer ()
updateCFG :: forall a. GraphBuilder BBID BasicBlock CFGEdge a -> CFGOptimizer ()
updateCFG GraphBuilder BBID BasicBlock CFGEdge a
update = do
  (CFG Graph BBID BasicBlock CFGEdge
g BBID
_ BBID
_ [Var]
_ MethodSig
_) <- CFGOptimizer CFG
getCFG
  let g' :: Either Text (Graph BBID BasicBlock CFGEdge)
g' = GraphBuilder BBID BasicBlock CFGEdge a
-> Graph BBID BasicBlock CFGEdge
-> Either Text (Graph BBID BasicBlock CFGEdge)
forall ni nd ed a.
(Eq ni, Ord ni) =>
GraphBuilder ni nd ed a
-> Graph ni nd ed -> Either Text (Graph ni nd ed)
G.update GraphBuilder BBID BasicBlock CFGEdge a
update Graph BBID BasicBlock CFGEdge
g
  case Either Text (Graph BBID BasicBlock CFGEdge)
g' of
    Left Text
m -> CompileError -> CFGOptimizer ()
forall a. CompileError -> CFGOptimizer a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (CompileError -> CFGOptimizer ())
-> CompileError -> CFGOptimizer ()
forall a b. (a -> b) -> a -> b
$ Maybe Range -> Text -> CompileError
CompileError Maybe Range
forall a. Maybe a
Nothing Text
m
    Right Graph BBID BasicBlock CFGEdge
g -> (CFG -> Identity CFG)
-> CFGOptimizerState -> Identity CFGOptimizerState
#cfg ((CFG -> Identity CFG)
 -> CFGOptimizerState -> Identity CFGOptimizerState)
-> ((Graph BBID BasicBlock CFGEdge
     -> Identity (Graph BBID BasicBlock CFGEdge))
    -> CFG -> Identity CFG)
-> (Graph BBID BasicBlock CFGEdge
    -> Identity (Graph BBID BasicBlock CFGEdge))
-> CFGOptimizerState
-> Identity CFGOptimizerState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Graph BBID BasicBlock CFGEdge
 -> Identity (Graph BBID BasicBlock CFGEdge))
-> CFG -> Identity CFG
#graph ((Graph BBID BasicBlock CFGEdge
  -> Identity (Graph BBID BasicBlock CFGEdge))
 -> CFGOptimizerState -> Identity CFGOptimizerState)
-> Graph BBID BasicBlock CFGEdge -> CFGOptimizer ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> b -> m ()
.= Graph BBID BasicBlock CFGEdge
g