module ACleanup(aCleanup) where
import ASyntax
import ASyntaxUtil
import Prim
import DisjointTest(DisjointTestState, initDisjointTestState,
                    addADefToDisjointTestState, checkDisjointExprWithCtx)
import Data.Maybe
import Flags(Flags)
import Control.Monad.State
import FStringCompat(mkFString)
import Position(noPosition)
import Id
import Util
import IOUtil(progArgs)
import PPrint(ppReadable, ppString)
import Error(ErrorHandle)

trace_disjoint_tests :: Bool
trace_disjoint_tests = "-trace-disjoint-tests" `elem` progArgs

-- =====
-- Naming conventions

acleanupPref :: String
acleanupPref = "_dfoo"

-- =====

-- A state monad for generating identifiers and capturing definitions during the cleanup pass
data CState = CState Integer [ADef] DisjointTestState

-- Identifiers are _prefix# and the first is _prefix1
-- prefix is acleanupPref
initCState :: DisjointTestState -> CState
initCState dts = CState 1 [] dts

-- the state monad itself
type CMonad   = StateT CState IO

addDef :: ADef -> CMonad ()
addDef d = do
    (CState i ds dts) <- get
    dts' <- liftIO $ addADefToDisjointTestState dts [d]
    put (CState i (d:ds) dts')

getDefs :: CMonad [ADef]
getDefs = do
            (CState i ds dts) <- get
            return (reverse ds)

getDisjointTestState :: CMonad DisjointTestState
getDisjointTestState = do
    (CState i ds dts) <- get
    return dts

updateDisjointTestState :: DisjointTestState -> CMonad ()
updateDisjointTestState dts = do
    (CState i ds _) <- get
    put (CState i ds dts)

newName :: CMonad AId
newName = do
    (CState i ds dts) <- get
    put (CState (i+1) ds dts)
    return $ mkId noPosition (mkFString (acleanupPref ++ itos i))


aCleanup :: ErrorHandle -> Flags -> APackage -> IO APackage
aCleanup errh flags apkg = do
  let userDefs = apkg_local_defs apkg
      state = apkg_state_instances apkg
      userARules = apkg_rules apkg
      ifs = apkg_interface apkg
      -- definitions of the value methods in the interface
      ifcDefs = [d | (AIDef { aif_value = d }) <- ifs] ++
                [d | (AIActionValue { aif_value = d }) <- ifs]
  let str = "cleanup_" ++ ppString (apkg_name apkg)
  dts <- initDisjointTestState str errh flags (userDefs ++ ifcDefs) state []
  let initstate = initCState dts
  evalStateT (do
                userARules' <- mapM (cleanupRule flags Nothing) userARules
                ifs' <- mapM (cleanupIfc flags) ifs
                newDefs <- getDefs
--              traceM (show newDefs)
                let defs' = userDefs ++ newDefs
                return (apkg { apkg_local_defs = defs', apkg_rules = userARules', apkg_interface = ifs' }))
            initstate


cleanupRule :: Flags -> Maybe AExpr -> ARule -> CMonad ARule
cleanupRule flags mrdy (ARule id rps descr wp pred actions asmps splitorig) =
  do -- if the rule is the Action part of a method (possible a split method)
     -- then we need to include the method's ready condiion
     let pred' = case mrdy of
                   Nothing -> pred
                   Just rdy -> aAnd rdy pred
     actions' <- cleanupActions flags pred' actions
     -- don't cleanup assumption actions because they include no methods
     return (ARule id rps descr wp pred actions' asmps splitorig)

cleanupIfc :: Flags -> AIFace -> CMonad AIFace
cleanupIfc flags a@(AIAction { aif_pred = rdy, aif_body = rs }) =
  do rs' <- mapM (cleanupRule flags (Just rdy)) rs
     return a { aif_body = rs' }

cleanupIfc flags a@(AIActionValue { aif_pred = rdy, aif_body = rs }) =
  do rs' <- mapM (cleanupRule flags (Just rdy)) rs
     return a { aif_body = rs' }

cleanupIfc _ ifc = return ifc

-- merge mutually exclusive calls to the same action method
-- really a later pass lifting (so I suppose no-lift no longer breaks anything)
cleanupActions :: Flags -> AExpr -> [AAction] -> CMonad [AAction]
cleanupActions flags pred as =
  let
      loop :: [AAction] -> [AAction] -> CMonad [AAction]
      loop merged [] = return (reverse merged)

      -- Found an action method
      loop merged (first@(ACall id methodid (cond:args)):rest) =
         -- Internal loop to scan for matching actions that might be ME
         let loopR :: [AAction] -> [AAction] -> CMonad [AAction]
             loopR scanned [] =
                return ((reverse merged) ++ [first] ++ (reverse scanned))

         -- not necessary - rest has been cleaned already
         -- loopR scanned [] = loop (first:merged) (reverse scanned)
             loopR scanned (firstR@(ACall id' methodid' (cond':args')):restR)
                                | (id == id') &&
                                    (methodid == methodid') &&
                                    ((length args) == (length args')) =
                do
                   dtState <- getDisjointTestState
                   (isDisjoint,newstate) <-
                        liftIO $ checkDisjointCond dtState pred cond cond'
                   updateDisjointTestState newstate

                   if (isDisjoint) then
                     do
                        newid <- newName
                        addDef (ADef newid aTBool
                                (APrim newid aTBool PrimBOr [cond, cond']) [])
                        newargs <-
                            (mapM (\ (arg, arg') ->
                                do
                                    argid <- newName
                                    let argtyp = (aType arg)
                                    addDef (ADef argid argtyp
                                        (APrim argid argtyp PrimIf [cond, arg, arg']) [])
                                    return (ASDef argtyp argid))
                                         (zip args args'))
                        let newcall = (ACall id methodid
                                ((ASDef aTBool newid):newargs))

          -- restR is guaranteed merged amongst itself (see below)
          -- so no more work need be done...
                        return ((reverse merged)
                                ++ [newcall]
                                ++ (reverse scanned)
                                ++ restR)
                     else loopR (firstR:scanned) restR
             loopR scanned (firstR:restR) = loopR (firstR:scanned) restR

             -- aggressively merge the rest
             -- which allows the shortcuts in loopR above
         in do rest' <- (loop [] rest)
               (loopR [] rest')

      -- don't try to merge foreign function calls since
      -- those shouldn't get muxed anyway
      loop merged (first:rest) = loop (first:merged) rest
  in (loop [] as)


checkDisjointCond :: DisjointTestState -> AExpr -> AExpr -> AExpr ->
                     IO (Bool, DisjointTestState)
checkDisjointCond dtState pred cond1 cond2 =
    do
      (mres, dtState') <- checkDisjointExprWithCtx dtState pred pred cond1 cond2
      let res  = if isNothing mres then False else fromJust mres
      when(trace_disjoint_tests) $
          putStrLn("checkDisjoint(ACleanup)\n" ++
                   -- "pred: " ++ (ppReadable pred) ++
                   "cond1: " ++ (ppReadable cond1) ++
                   "cond2: " ++ (ppReadable cond2) ++
                   "Result: " ++ (show mres))
      return (res, dtState')