-- Copyright (C) 2018-2024 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.

-- CFG -- Control Flow Graph with SSA nodes
module CFG (plot, CFGContext (..), Condition (..), BasicBlock (..), CFGEdge (..), CFG (..), SingleFileCFG (..), buildAndOptimize) where

import AST qualified
import CFG.Build (CFGContext (..), buildCFGs)
import CFG.Optimizations.Optimizer (CFGOptimizer, runOptimizerOnCFG)
import CFG.Optimizations.RemoveDeadBlock (removeDeadBlock)
import CFG.Optimizations.RemoveNoOp (removeNoOp)
import CFG.Plot (fileCFGsToDot)
import CFG.Types
import Control.Lens (views, (%~), (&), (.~), (^.))
import Control.Monad (mapM_)
import Data.Functor ((<&>))
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Text qualified as Text
import SSA
import Semantic qualified as SE
import Types

{-
Refactor and clean up.
TODO:
1. Find better ways to add phi nodes. [DROP]
2. Refactor control start/exit related code. [DONE]
3. Produce dot plot with some proper library. [DONE]
4. Add unit tests.
5. Other chores.
-}

optimizations :: [CFGOptimizer ()]
optimizations :: [CFGOptimizer ()]
optimizations = [CFGOptimizer ()
removeDeadBlock, CFGOptimizer ()
removeNoOp]

buildAndOptimize :: AST.ASTRoot -> SE.SemanticInfo -> Either [CompileError] SingleFileCFG
buildAndOptimize :: ASTRoot -> SemanticInfo -> Either [CompileError] SingleFileCFG
buildAndOptimize ASTRoot
root SemanticInfo
si = do
  let context :: CFGContext
context = SemanticInfo -> CFGContext
CFGContext SemanticInfo
si
  SingleFileCFG
fileCFG <- ASTRoot -> CFGContext -> Either [CompileError] SingleFileCFG
buildCFGs ASTRoot
root CFGContext
context
  let runOpts :: SingleFileCFG -> Either [CompileError] (Map Name CFG)
runOpts = LensLike'
  (Const (Either [CompileError] (Map Name CFG)))
  SingleFileCFG
  (Map Name CFG)
-> (Map Name CFG -> Either [CompileError] (Map Name CFG))
-> SingleFileCFG
-> Either [CompileError] (Map Name CFG)
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike'
  (Const (Either [CompileError] (Map Name CFG)))
  SingleFileCFG
  (Map Name CFG)
#cfgs ((Map Name CFG -> Either [CompileError] (Map Name CFG))
 -> SingleFileCFG -> Either [CompileError] (Map Name CFG))
-> (Map Name CFG -> Either [CompileError] (Map Name CFG))
-> SingleFileCFG
-> Either [CompileError] (Map Name CFG)
forall a b. (a -> b) -> a -> b
$ (CFG -> Either [CompileError] CFG)
-> Map Name CFG -> Either [CompileError] (Map Name CFG)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Map Name a -> m (Map Name b)
mapM (CFGOptimizer () -> CFG -> Either [CompileError] CFG
runOptimizerOnCFG ([CFGOptimizer ()] -> CFGOptimizer ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ [CFGOptimizer ()]
optimizations))
  Map Name CFG
cfgs <- SingleFileCFG -> Either [CompileError] (Map Name CFG)
runOpts SingleFileCFG
fileCFG
  SingleFileCFG -> Either [CompileError] SingleFileCFG
forall a. a -> Either [CompileError] a
forall (m :: * -> *) a. Monad m => a -> m a
return (SingleFileCFG -> Either [CompileError] SingleFileCFG)
-> SingleFileCFG -> Either [CompileError] SingleFileCFG
forall a b. (a -> b) -> a -> b
$ SingleFileCFG
fileCFG {$sel:cfgs:SingleFileCFG :: Map Name CFG
cfgs = Map Name CFG
cfgs}

plot :: AST.ASTRoot -> SE.SemanticInfo -> Either [CompileError] String
plot :: ASTRoot -> SemanticInfo -> Either [CompileError] String
plot ASTRoot
root SemanticInfo
si = do
  SingleFileCFG
fileCFG <- ASTRoot -> SemanticInfo -> Either [CompileError] SingleFileCFG
buildAndOptimize ASTRoot
root SemanticInfo
si
  String -> Either [CompileError] String
forall a. a -> Either [CompileError] a
forall (m :: * -> *) a. Monad m => a -> m a
return (String -> Either [CompileError] String)
-> String -> Either [CompileError] String
forall a b. (a -> b) -> a -> b
$ Name -> String
Text.unpack (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ SingleFileCFG -> Name
fileCFGsToDot SingleFileCFG
fileCFG

-- linearize :: CFG -> [SSA]
-- linearize cfg = _