{-# LANGUAGE RecordWildCards, LambdaCase, TupleSections, OverloadedStrings #-}
module Stg.DeadFunctionElimination.Facts where

import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.ByteString.Char8 as BS8
import Control.Monad

import System.IO
import System.FilePath

import Stg.Syntax

writeDfeFacts :: FilePath -> Module -> IO ()
writeDfeFacts prefixPath Module{..} = do
  let
      openFact :: String -> IO Handle
      openFact name = do
        let factFile = prefixPath -<.> name
        h <- openFile factFile WriteMode
        putStrLn factFile
        pure h

      closeFact :: Handle -> IO ()
      closeFact h = do
        hFlush h
        hClose h

      getTopBinders :: TopBinding -> [Binder]
      getTopBinders = \case
        StgTopLifted (StgNonRec b _)  -> [b]
        StgTopLifted (StgRec bs)      ->  map fst bs
        StgTopStringLit b _           -> [b]

      topBinders :: [Binder]
      topBinders = concatMap getTopBinders moduleTopBindings

      extBinders :: [Binder]
      extBinders = concatMap snd . concatMap snd $ moduleExternalTopIds

  -- from module

  -- export LiveSource
  factLiveSource <- openFact "LiveSource.facts"
  case moduleForeignStubs of
    NoStubs -> pure ()
    ForeignStubs{..} -> do
      forM_ [b | StubDeclExport _ b _ <- fsDecls] $ \b -> do
        BS8.hPutStrLn factLiveSource $ binderUniqueName b
  closeFact factLiveSource

  -- export TyCon with DataCons
  factTyCon <- openFact "TyCon.facts"
  forM_ (concatMap snd . concatMap snd $ moduleTyCons) $ \tc -> do
    let tcName = tcUniqueName tc
    forM_ (tcDataCons tc) $ \dc -> do
      BS8.hPutStrLn factTyCon $ tcName <> "\t" <> dcUniqueName dc
  closeFact factTyCon

  -- from AST
  factTyConReference    <- openFact "TyConReference.facts"
  factDataConReference  <- openFact "DataConReference.facts"
  factFunReference      <- openFact "FunReference.facts"

  let
      funNameMap :: Map BinderId Name
      funNameMap = Map.fromList [(binderId b, binderUniqueName b) | b <- topBinders ++ extBinders]

      addFunRef :: Name -> Binder -> IO ()
      addFunRef fun b = case Map.lookup (binderId b) funNameMap of
        Nothing -> pure ()
        Just n  -> BS8.hPutStrLn factFunReference $ fun <> "\t" <> n

      addDataConRef :: Name -> DataCon -> IO ()
      addDataConRef fun dc = do
        BS8.hPutStrLn factDataConReference $ fun <> "\t" <> dcUniqueName dc

      addTyConRef :: Name -> TyCon -> IO ()
      addTyConRef fun tc = do
        BS8.hPutStrLn factTyConReference $ fun <> "\t" <> tcUniqueName tc

  let
      visitTopBinding :: TopBinding -> IO ()
      visitTopBinding = \case
        StgTopLifted (StgNonRec b e)  -> visitRhs (binderUniqueName b) e
        StgTopLifted (StgRec bs)      -> sequence_ [visitRhs (binderUniqueName b) e | (b, e) <- bs]
        StgTopStringLit{}             -> pure ()

      visitBinding :: Name -> Binding -> IO ()
      visitBinding fun = \case
        StgNonRec _ rhs -> visitRhs fun rhs
        StgRec    l     -> mapM_ (visitRhs fun) $ map snd l

      visitRhs :: Name -> Rhs -> IO ()
      visitRhs fun = \case
        StgRhsClosure _ _ args expr -> mapM_ (addFunRef fun) args >> visitExpr fun expr
        StgRhsCon con args          -> addDataConRef fun con >> mapM_ (visitArg fun) args

      visitArg :: Name -> Arg -> IO ()
      visitArg fun = \case
        StgVarArg v               -> addFunRef fun v
        StgLitArg{}               -> pure ()

      visitExpr :: Name -> Expr -> IO ()
      visitExpr fun = \case
        StgApp f args             -> addFunRef fun f >> mapM_ (visitArg fun) args
        StgLit{}                  -> pure ()
        StgConApp con args _      -> addDataConRef fun con >> mapM_ (visitArg fun) args

        -- NOTE: _op could only refer to function that are implemented in the low level IRs but not in STG
        StgOpApp _op args _ mtc   -> mapM_ (visitArg fun) args >> case mtc of
                                      Nothing -> pure ()
                                      Just tc -> addTyConRef fun tc

        StgCase expr _ aty alts   -> visitExpr fun expr >> visitAltType fun aty >> mapM_ (visitAlt fun) alts
        StgLet b  e               -> visitBinding fun b >> visitExpr fun e
        StgLetNoEscape b  e       -> visitBinding fun b >> visitExpr fun e
        StgTick _t e              -> visitExpr fun e

      visitAltType :: Name -> AltType -> IO ()
      visitAltType fun = \case
        PolyAlt       -> pure ()
        MultiValAlt{} -> pure ()
        PrimAlt{}     -> pure ()
        AlgAlt tc     -> addTyConRef fun tc

      visitAlt :: Name -> Alt -> IO ()
      visitAlt fun (Alt ac _ expr) = visitAltCon fun ac >> visitExpr fun expr

      visitAltCon :: Name -> AltCon -> IO ()
      visitAltCon fun = \case
        AltDataCon dc -> addDataConRef fun dc
        AltLit{}      -> pure ()
        AltDefault    -> pure ()

  -- traverse AST and collect facts
  mapM_ visitTopBinding moduleTopBindings

  -- close facts
  closeFact factTyConReference
  closeFact factDataConReference
  closeFact factFunReference
