{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE CPP #-}

-- {-# LANGUAGE DataKinds #-}
-- {-# LANGUAGE TypeOperators #-}

module Ivory.Opts.CFG
  ( callGraphDot
  -- ** Generate a dot file of the control flow for a program.
  , SizeMap(..)
  -- ** Implementation-defined size map for stack elements.
  , defaultSizeMap
  , hasLoop
  -- ** Does the control-flow graph contain a loop?
  , WithTop
  , maxStack
  -- ** What is the maximum stack size of the program?
  )
    where

import Prelude ()
import Prelude.Compat hiding (lookup)

import qualified Ivory.Language.Array       as I
import qualified Ivory.Language.Syntax.AST  as I
import qualified Ivory.Language.Syntax.Type as I
import qualified Data.Graph.Inductive as G

import Control.Applicative (liftA2)
import System.FilePath
import Data.Maybe
import Data.List (find,(\\))
import qualified Data.IntMap as M
import MonadLib (StateT, get, set, Id, StateM, runM)
import MonadLib.Derive (derive_get, derive_set, Iso(..))

-- import Ivory.Language hiding (Top)

--------------------------------------------------------------------------------
-- Types

-- | Add Top to a type.
data WithTop a = Top | Val a
  deriving (Eq, Functor)

instance Show a => Show (WithTop a) where
  show Top     = "Top"
  show (Val a) = show a

instance Ord a => Ord (WithTop a) where
  compare Top     (Val _) = GT
  compare (Val _) Top     = LT
  compare Top      Top    = EQ
  compare (Val a) (Val b) = compare a b

instance Applicative WithTop where
  pure a = Val a
  Val _ <*> Top     = Top
  Val f <*> (Val a) = Val (f a)
  _     <*> _       = Top

type Size = Integer
type CallNm = String

-- | We track those objects for which the size might be
-- implementation-dependent.
data StackType = TyStruct String      -- ^ Name of struct
               | TyArr StackType Size -- ^ Array and size
               | Ptr                  -- ^ Pointer or ref type
               | TyVoid
               | TyInt IntSize
               | TyWord WordSize
               | TyBool
               | TyChar
               | TyFloat
               | TyDouble
  deriving (Show, Eq)

data IntSize = Int8
             | Int16
             | Int32
             | Int64
  deriving (Show,Eq)

data WordSize = Word8
              | Word16
              | Word32
              | Word64
  deriving (Show,Eq)

data Block a
  = Stmt a
  | Branch [Block a] [Block a]
  | Loop (Maybe Integer) [Block a] -- If we know how many loops we make, we
                                   -- store it.  Otherwise, `Nothing`.
  deriving (Show, Eq)

type Control = Block CallNm    -- ^ Name of function being called
type Alloc   = Block StackType -- ^ Memory being allocated

-- | Describes the CFG and memory usage for a single procedure.
data ProcInfo = ProcInfo
  { procSym :: CallNm
  , params  :: [StackType]
    -- ^ Parameters pushed onto the stack.
  , alloc   :: [Alloc]
    -- ^ Allocated elements.
  , calls   :: [Control]
  } deriving (Show, Eq)

data ModuleInfo = ModuleInfo
  { modName :: String
  , procs   :: [ProcInfo]
  } deriving (Show, Eq)

--------------------------------------------------------------------------------
-- Procedure Analysis

toStackTyped :: I.Typed a -> StackType
toStackTyped ty = toStackType (I.tType ty)

toStackType :: I.Type -> StackType
toStackType ty =
  case ty of
    I.TyStruct nm  -> TyStruct nm
    I.TyArr i t    -> TyArr (toStackType t) (fromIntegral i)
    I.TyRef{}      -> Ptr
    I.TyConstRef{} -> Ptr
    I.TyPtr{}      -> Ptr
    I.TyVoid       -> TyVoid
    I.TyInt i      -> TyInt (toIntType i)
    I.TyWord w     -> TyWord (toWordType w)
    I.TyIndex _n   -> toStackType I.ixRep
    I.TyBool       -> TyBool
    I.TyChar       -> TyChar
    I.TyFloat      -> TyFloat
    I.TyDouble     -> TyDouble
    t              -> error $ "Unhandled stack type: " ++ show t

toIntType :: I.IntSize -> IntSize
toIntType i =
  case i of
    I.Int8  -> Int8
    I.Int16 -> Int16
    I.Int32 -> Int32
    I.Int64 -> Int64

toWordType :: I.WordSize -> WordSize
toWordType i =
  case i of
    I.Word8  -> Word8
    I.Word16 -> Word16
    I.Word32 -> Word32
    I.Word64 -> Word64

cfgProc :: I.Proc -> ProcInfo
cfgProc proc = ProcInfo
  { procSym = I.procSym proc
  , params  = map toStackTyped  (I.procArgs proc)
  , alloc   = concatMap toAlloc (I.procBody proc)
  , calls   = concatMap toCall  (I.procBody proc)
  }

toAlloc :: I.Stmt -> [Alloc]
toAlloc stmt =
  case stmt of
    I.Assign ty _ _                     -> [Stmt $ toStackType ty]
    I.AllocRef ty _ _                   -> [Stmt $ toStackType ty]
    I.Deref ty _ _                      -> [Stmt $ toStackType ty]
    -- Descend into blocks
    I.IfTE _ blk0 blk1                  -> [ Branch (concatMap toAlloc blk0)
                                                    (concatMap toAlloc blk1) ]
                                           -- For the loop variable.
    I.Loop m _ e _ blk                  ->
      let ty = I.ixRep in
      [Stmt (toStackType ty), Loop (Just (loopIdx m e)) (concatMap toAlloc blk)]
    I.Forever blk                       ->
      [Loop Nothing (concatMap toAlloc blk)]
    _                                   -> []

toCall :: I.Stmt -> [Control]
toCall stmt =
  case stmt of
    I.IfTE _ blk0 blk1 -> [ Branch (concatMap toCall blk0)
                                   (concatMap toCall blk1) ]
    I.Call _ _ nm _  -> case nm of
                          I.NameSym sym -> [Stmt sym]
                          I.NameVar _   -> error $ "XXX need to implement function pointers."
    I.Loop m _ e _ blk -> [Loop (Just (loopIdx m e)) (concatMap toCall blk)]
    _                  -> []

loopIdx :: Integer -> I.Expr -> Integer
loopIdx _ (I.ExpLit (I.LitInteger i)) = i
loopIdx m _                           = m

_getIdx :: I.Expr -> Maybe Integer
_getIdx e = case e of
             I.ExpLit (I.LitInteger i) -> Just i
             _                         -> Nothing

--------------------------------------------------------------------------------
-- Call-graph construction from a module.

-- type Node = G.LNode ProcInfo

-- | A call graph is a graph in which nodes are labeled by their procedure info
-- and edges are unlabeled.
type CFG = G.Gr ProcInfo ()

flattenControl :: Control -> [CallNm]
flattenControl ctrl =
  case ctrl of
    Stmt str           -> [str]
    Branch ctrl0 ctrl1 ->
      concatMap flattenControl ctrl0 ++ concatMap flattenControl ctrl1
    Loop _ ctrl0       -> concatMap flattenControl ctrl0

cfgModule :: I.Module -> ModuleInfo
cfgModule m = ModuleInfo
  { modName = I.modName m
  , procs   = map cfgProc ps
  }
    where ps = I.public (I.modProcs m) ++ I.private (I.modProcs m)

-- | Construct a control-flow graph from an Ivory module.
cfg :: I.Module -> CFG
cfg m = G.insEdges (concatMap go nodes) $ G.insNodes nodes G.empty
  where
  nodes :: [G.LNode ProcInfo]
  nodes = zip [0,1..] (procs $ cfgModule m)

  go :: (Int, ProcInfo) -> [G.LEdge ()]
  go (i,p) =
    let outCalls = concatMap flattenControl (calls p) in
    let outIdxs  = catMaybes (map (lookup nodes) outCalls) in
    zip3 (repeat i) outIdxs (repeat ()) -- outboud edges

  lookup ls sym | [] <- ls         = Nothing
                | ((i,p):_) <- ls
                , procSym p == sym = Just i
                | (_:ls') <- ls    = lookup ls' sym
                | otherwise        = error "Unreachable in cfg" -- To make GHC happy.

-- | Just label the nodes with the function names.
procSymGraph :: CFG -> G.Gr CallNm ()
procSymGraph = G.nmap procSym

-- | Does the program have a loop in it?
hasLoop :: CFG -> Bool
hasLoop = G.hasLoop

--------------------------------------------------------------------------------
-- Stack usage analysis.

data SizeMap = SizeMap
  { stackElemMap :: StackType -> Size -- ^ Mapping from `StackType` to their
                                      --   implementation-dependent size.
  , retSize      :: Size              -- ^ Size of a return address.
  }


-- | Maps everything to being size 1.  Useful for testing.
defaultSizeMap :: SizeMap
defaultSizeMap = SizeMap
  { stackElemMap = const 1
  , retSize      = 1
  }

-- A mapping from `Node` (Ints) to the maximum stack usage for that function.
type MaxMap = M.IntMap Size

newtype MaxState a = MaxState
  { unSt :: StateT MaxMap Id a
  } deriving (Functor, Monad, Applicative)

instance StateM MaxState MaxMap where
  get = derive_get (Iso MaxState unSt)
  set = derive_set (Iso MaxState unSt)

emptyMaxSt :: MaxMap
emptyMaxSt = M.empty

getMaxMap :: MaxState MaxMap
getMaxMap = return =<< get

-- | Takes a procedure name, a control-flow graph, and a `SizeMap` and produces
-- the maximum stack usage starting at the procedure give.  Returns `Top` if
-- there is an unanalyzable loop in the program (ie., non-constant loop) and
-- @Val max@ otherwise.
maxStack :: CallNm -> CFG -> SizeMap -> WithTop Size
maxStack proc cf szmap = go proc
  where
  go p = fst $ runM (unSt (maxStack' cf szmap [] (findNode cf p))) emptyMaxSt

-- Get the node from it's name.
findNode :: CFG -> CallNm -> G.Node
findNode cf proc = fst $
  fromMaybe (error $ "Proc " ++ proc ++ " is not in the graph!")
            (find ((== proc) . procSym . snd) (G.labNodes cf))

maxStack' :: CFG -> SizeMap -> [G.Node] -> G.Node -> MaxState (WithTop Size)
maxStack' cf szmap visited curr
  | curr `elem` visited = return Top -- A loop is detected.
  | otherwise           = maxStackNode
  where
  -- Process the max stack for a single node.
  maxStackNode :: MaxState (WithTop Size)
  maxStackNode = do
    blkMax <- goBlks cf szmap visited curr alloc' calls'
    let sz = (topAllocSz + paramsSz + retSize szmap +) <$> blkMax
    return sz

    where
    cxt = G.context cf curr
    procInfo = G.lab' cxt

    alloc' = alloc procInfo
    calls' = calls procInfo

    -- Top-level allocation
    topAllocSz :: Size
    topAllocSz = getSize szmap (getBlock alloc')

    paramsSz :: Size
    paramsSz = getSize szmap (params procInfo)

goBlks :: CFG -> SizeMap -> [G.Node] -> G.Node
       -> [Alloc] -> [Control] -> MaxState (WithTop Size)
goBlks cf szmap visited curr acs cns =
  case (acs, cns) of
    -- Done with all blocks/statements in this function.
    ([] ,[])                               -> return (Val 0)
    (Branch a0 a1:acs', Branch c0 c1:cns') -> do
      sz0 <- goBlks' a0 c0
      sz1 <- goBlks' a1 c1
      sz2 <- goBlks' acs' cns'
      return (liftA2 max sz0 sz1 <+> sz2)
    -- There's no new loop allocation, just assignments.  So we assume that on
    -- each iteration, the stack usage doesn't change.
    (Loop _ a:acs', Loop _ c:cns')         -> do
      sz0 <- goBlks' a c
      sz1 <- goBlks' acs' cns'
      return (sz0 <+> sz1)

      -- case idx of
      --   -- Unknown loop bound.
      --   Nothing -> if sz0 == Val 0 then return sz1

      --                else return Top -- Did some allocation in the loop, so
      --                                -- can't compute.
      --   Just _  ->
  -- There is either a straight-line call or assignment.
    _                                       -> do
      sz0 <- goBlk cf szmap visited curr (getBlock acs) (getBlock cns)
      sz1 <- goBlks' (nxtBlock acs) (nxtBlock cns)
      return (sz0 <+> sz1)
  where goBlks' = goBlks cf szmap visited curr

goBlk :: CFG -> SizeMap -> [G.Node] -> G.Node -> [StackType]
      -> [CallNm] -> MaxState (WithTop Size)
goBlk cf szmap visited curr acs cns = do
  maxMp <- getMaxMap
  let localAlloc = getSize szmap acs
  let callNodes  = map (findNode cf) cns
  let allCalls   = zip callNodes (map (flip M.lookup $ maxMp) callNodes)
  newMaxs <- mapM cachedCalls allCalls
  return $ Val localAlloc <+> if null newMaxs then Val 0
                                     -- Depends on Top being >= all vals.
                                else maximum newMaxs

  where
  cachedCalls :: (G.Node, Maybe Size) -> MaxState (WithTop Size)
  cachedCalls (n, msz) | Just sz <- msz = return (Val sz)
                       | otherwise      =
    maxStack' cf szmap (curr:visited) n

(<+>) :: Num a => WithTop a -> WithTop a -> WithTop a
(<+>) = liftA2 (+)

-- Get the cumulative sizes of allocated things.
getSize :: SizeMap -> [StackType] -> Size
getSize szmap = sum . map (stackElemMap szmap)

-- Get the prefix of a list of @[Block a] in the current block.
getBlock :: (Show a, Eq a) => [Block a] -> [a]
getBlock bls | (Stmt b:bs) <- bls = b : getBlock bs
             | otherwise          = []

nxtBlock :: (Show a, Eq a) => [Block a] -> [Block a]
nxtBlock bls | (Stmt _:bs) <- bls = nxtBlock bs
             | otherwise           = bls

-- | Call-graph output.  Takes a start function foo, a filepath, and emits the
-- graph named "foo.dot" in the filepath.
callGraphDot :: CallNm -> FilePath -> [I.Module] -> IO CFG
callGraphDot proc path mods =
  writeFile (path </> proc `addExtension` "dot") grOut >> return graph
  where
  m = mconcat mods
  grOut = graphviz filterG (I.modName m)
  filterG :: G.Gr CallNm ()
  filterG = let closure = G.reachable (findNode graph proc) graph in
            G.delNodes (G.nodes graph \\ closure) (procSymGraph graph)
  graph = cfg m

--------------------------------------------------------------------------------
-- Adapted from fgl (BSD3) to remove sizing info.
graphviz :: (G.Graph g, Show a, Show b)
  => g a b   -- ^ The graph to format
  -> String  -- ^ The title of the graph
  -> String
graphviz g t =
    let n = G.labNodes g
        e = G.labEdges g
        ns = concatMap sn n
        es = concatMap se e
    in "digraph "++t++" {\n"
            ++ ns
            ++ es
        ++"}"
    where sn (n, a) | sa == ""  = ""
                    | otherwise = '\t':(show n ++ sa ++ "\n")
            where sa = sl a
          se (n1, n2, b) = '\t':(show n1 ++ " -> " ++ show n2 ++ sl b ++ "\n")

sl :: (Show a) => a -> String
sl a = let l = sq (show a)
       in if (l /= "()") then (" [label = \""++l++"\"]") else ""

sq :: String -> String
sq s@[_]                     = s
sq ('"':s)  | last s == '"'  = init s
            | otherwise      = s
sq ('\'':s) | last s == '\'' = init s
            | otherwise      = s
sq s                         = s

--------------------------------------------------------------------------------
-- Testing

{-

fibCFG :: CFG
fibCFG = cfg fibMod

outGraph = writeFile "fib.dot" gr
  where
  gr = G.graphviz (procSymGraph fibCFG) "fib" (4,4) (4,4) G.Portrait
----------------------

fibMod :: Module
fibMod = package "fib" $ do
  incl fib
  incl fib_aux

fib :: Def ('[Uint32] :-> Uint64)
fib  = proc "fib" (\n -> body (ret =<< call fib_aux 0 1 n))

fib_aux :: Def ('[Uint32,Uint32,Uint32] :-> Uint64)
fib_aux  = proc "fib_aux" $ \ a b n -> body $ do
  ifte_ (n ==? 0)
    (ret (safeCast a))
    (ret . safeCast =<< call fib_aux b (a + b) (n - 1))

fibSz = SizeMap
  { stackElemMap = \_ -> 1
  , retSize = 1 }

-------------------------------------

x :: G.Gr Int ()
x = G.insEdges edges gr
  where
  edges = [(0,1,()), (0,2,()), (2,1,()), (2,3,())]
  gr = G.insNodes nodes G.empty
  nodes = zip [0,1 ..] [2,3,4,1]

y :: G.Gr Int ()
y = G.insEdges edges gr
  where
  edges = [(0,1,()), (0,2,()), (2,1,()), (1,3,())]
  gr = G.insNodes nodes G.empty
  nodes = zip [0,1 ..] [1,4,5,8]

maxf :: G.Gr Int () -> G.Node -> Maybe Int
maxf g n = maxf' [] n
  where
  maxf' :: [G.Node] -> G.Node -> Maybe Int
  maxf' visited curr
    | curr `elem` visited = Nothing -- We made a loop on this path!  That's Top.
    | otherwise           = liftA2 (+) (Just $ G.lab' cxt) go
        where
        cxt = G.context g curr
        res = map (maxf' (G.node' cxt : visited)) (G.suc' cxt)
        go | null (G.suc' cxt)  = Just 0 -- No more successors to try
           | Nothing `elem` res = Nothing
           | otherwise          = Just $ maximum (catMaybes res)

-}