content: module SimPackageOpt ( simPackageOpt ) where
  
  import Flags
  import SimPackage
  import ASyntax
  import ASyntaxUtil
  import AOpt(aExpandDynSel, aInsertCaseDef)
  import Prim
  import IntLit
  import IntegerUtil(aaaa)
  import Util(mapSndM, map_mapM)
  import Error(internalError, ErrMsg(..), ErrorHandle, bsError)
  import Id(isFromRHSId, dummy_id, mkIdCanFire, mkIdWillFire)
  import Position(noPosition)
  import Data.Maybe
  import PPrint
  
  import qualified Data.Map as M
  import qualified Data.Set as S
  import Util(map_insertManyWith)
  
  -- import Debug.Trace
  
  -- -------------------------
  
  simPackageOpt :: ErrorHandle -> Flags -> SimSystem -> IO SimSystem
  simPackageOpt errh flags ss = do
    let pkgs0 = M.toList (ssys_packages ss)
    pkgs1 <- mapSndM (pkgOpt errh flags) pkgs0
    return $ ss { ssys_packages = M.fromList pkgs1 }
  
  pkgOpt :: ErrorHandle -> Flags -> SimPackage -> IO SimPackage
  pkgOpt errh flags pkg0 = do
    -- Replace dynamic array selection with case expressions.
    -- This needs to be done before aInsertCase, which can merge adjacent
    -- case/if expressions.
    let pkg1A = expandDynSel pkg0
  
    -- Before any inlining, use aInsertCase to infer case expression.
    -- It only infers in defs, so it needs to be applied before any inlining.
    -- Note: This can leave unused defs, so we rely on the later opts
    -- to remove unused defs.
    let pkg1B = insertCase pkg1A
  
    -- remove occurrences of ASAny
    --
    -- replace occurrences of ASAny with a user-specified value;
    -- this is monadic because it can error if the user specified X or Z
    pkg2 <- convertASAny errh flags pkg1B
  
    -- inline defs (that were created in anticipation of CSE that didn't happen)
    -- and remove unused defs
    let pkg3 = inlineDefs pkg2
  
    -- combine nested concats and constants exposed by converting ASAny,
    -- and optimize wide concats for Bluesim codegen
    let pkg4 = optimizeConcats pkg3
  
    return pkg4
  
  -- -------------------------
  
  data UseInfo = None | Only AId | Multiple
    deriving (Eq,Show)
  
  instance PPrint UseInfo where
    pPrint _ _ None     = text "None"
    pPrint d _ (Only x) = (text "Only") <+> pPrint d 0 x
    pPrint _ _ Multiple = text "Multiple"
  
  combine_uses :: UseInfo -> UseInfo -> UseInfo
  combine_uses None     x        = x
  combine_uses x        None     = x
  combine_uses (Only _) (Only _) = Multiple
  combine_uses Multiple  _       = Multiple
  combine_uses _        Multiple = Multiple
  
  collectUses :: DefMap -> [AId] -> M.Map AId UseInfo
  collectUses defmap used_vars = collectUses' use_map0 S.empty used_vars
    where -- consider the vars we start with as "multiple"
          use_map0 = M.fromList [ (v, Multiple) | v <- used_vars ]
          collectUses' m s [] = m
          collectUses' m s (v:vs) | (S.member v s) =
              -- we've already visited this def
              collectUses' m s vs
          collectUses' m s (v:vs) =
              -- some vars are not defs
              case (M.lookup v defmap) of
                Nothing -> collectUses' m s vs
                Just d@(ADef i _ _ _) ->
                    let new_vs = aVars d
                        new_uses = [ (new_v, Only i) | new_v <- new_vs ]
                        m' = map_insertManyWith combine_uses new_uses m
                        s' = S.insert v s
                    in  collectUses' m' s' (new_vs ++ vs)
  
  -- This is inlining defs that were created during AConv.
  --   In AConv, expressions are added to a map in an attempt to CSE.
  --   If the expr doesn't already exist in the map, then a new name is
  --   created, with the FromRHS Id property.  If that expression is
  --   seen again, then the same name is used; but if the expression
  --   doesn't appear again, then the new name was not necessary.
  --   This code checks for FromRHS property on defs that are used
  --   only once, and inlines them.
  -- This code also eliminates unused defs (which might have resulted
  -- from the "insertCase" step).
  --
  inlineDefs :: SimPackage -> SimPackage
  inlineDefs pkg =
    let
        def_map0 = sp_local_defs pkg
        defs0 = M.toList def_map0
  
        used_in_rule = concatMap aVars (sp_rules pkg)
        used_in_method = concatMap aVars (sp_interface pkg)
        -- We don't extract the uses in state instantiations
        -- because we assume that those exprs have been inlined
        --used_in_state = concatMap aVars (sp_instances pkg)
        -- we want to make sure to preserve the CF/WF signals
        used_in_sched =
            let sched_ids = map arule_id (sp_rules pkg) ++
                            concatMap aIfaceSchedNames (sp_interface pkg)
            in  concat [ [mkIdCanFire i, mkIdWillFire i] | i <- sched_ids ]
        used_vars = used_in_rule ++ used_in_method ++ used_in_sched
  
        use_map = collectUses def_map0 used_vars
  
        isTaskOrForeignFunc (ADef _ _ (AFunCall {}) _)   = True
        isTaskOrForeignFunc (ADef _ _ (ATaskValue {}) _) = True
        isTaskOrForeignFunc _                          = False
        -- case-expr becomes case-stmt, which needs a def name to assign to
        isCase (ADef _ _ (APrim _ _ PrimCase _) _) = True
        isCase _ = False
        -- is an unsized string or an array of unsized strings
        isUnsized (ATString Nothing) = True
        isUnsized (ATArray _ t) = isUnsized t
        isUnsized _ = False
        isNotOk d
          | isUnsized (adef_type d)         = True
          | (aSize d > 64)                  = True
          | isTaskOrForeignFunc d           = True
          | isCase d                        = True
          | otherwise                       = False
        dangerous = [ i | (i,d) <- defs0, isNotOk d ]
  
        -- don't inline these
        keep = S.fromList (used_vars ++ dangerous)
  
        -- construct a substitution
        subst = M.fromList [ (i, adef_expr (fromJust md))
                           | (i,(Only _)) <- M.toList use_map
                           , isFromRHSId i
                           , not (i `S.member` keep)
                           , let md = M.lookup i (sp_local_defs pkg)
                           , isJust md
                           ]
  
        defs1 = mapMaybe (processDef use_map subst) defs0
    in pkg { sp_local_defs = M.fromList defs1 }
  
  processDef :: M.Map AId UseInfo -> M.Map AId AExpr ->
                (AId, ADef) -> Maybe (AId, ADef)
  processDef use_map subst (name,def) =
    let has_subst = name `M.member` subst
        has_use = name `M.member` use_map
        doSubst (ASDef _ i) = (\e -> exprMap doSubst e) <$> M.lookup i subst
        doSubst _           = Nothing
        opt d@(ADef _ _ e _) = d { adef_expr = exprMap doSubst e }
    in if has_subst
       then Nothing
       else if has_use
       then Just (name, opt def)
       else Nothing
  
  -- -------------------------
  
  optimizeConcats :: SimPackage -> SimPackage
  optimizeConcats pkg =
     let optConcat p@(APrim pid pty PrimConcat args) =
            let -- flatten sub-concats into this one
                flatten ((APrim _ _ PrimConcat es):rest) = (flatten es) ++ (flatten rest)
                flatten (x:xs)                           = x:(flatten xs)
                flatten []                               = []
                -- combine adjacent constants (this can happen because ASAny
                -- was only replaced with a value at the start of the Bluesim
                -- backend
                combine_constants (0,_) [] = []
                combine_constants (w,n) [] = [ ASInt dummy_id (ATBit w) (ilSizedBin w n) ]
                combine_constants (w,n) ((ASInt _ (ATBit w') (IntLit _ _ n')):es) =
                  combine_constants (w+w', n*(2^w') + n') es
                combine_constants (0,_) (e:es) =
                  [e] ++ (combine_constants (0,0) es)
                combine_constants (w,n) (e:es) =
                  [ASInt dummy_id (ATBit w) (ilSizedBin w n),e] ++ (combine_constants (0,0) es)
                -- this is the optimized concat arguments regardless of width
                args' = map optConcat $ combine_constants (0,0) $ flatten args
            in if ((aSize pty) <= 64)
               then -- for non-wide concats, we have no optimizations
                    -- beyond the basic ones
                    case args' of
                      []  -> internalError "SimPackageOpt: optimizeConcats has empty arg list"
                      [e] -> e
                      _   -> APrim pid pty PrimConcat args'
               else -- restructure wide concats into chunks aligned on
                    -- 32-bit boundaries
                    let extract (APrim aid (ATBit sz) PrimExtract [e,_,ASInt _ _ (IntLit _ _ base)]) nbits lo =
                          APrim aid (ATBit nbits) PrimExtract [e, aNat (base + lo + nbits - 1), aNat (base + lo) ]
                        extract e nbits lo =
                          APrim dummy_id (ATBit nbits) PrimExtract [e, aNat (lo + nbits - 1), aNat lo]
                        restructure grp sz []     = [grp]
                        restructure grp sz (e:es) =
                          let bits_left_in_dword = if ((sz /= 0) && ((sz `rem` 32) == 0)) then 32 else (sz `rem` 32);
                              e_sz = aSize e
                          in if e_sz <= bits_left_in_dword
                             then restructure (grp ++ [e]) (sz - e_sz) es
                             else -- need to break e across word boundary
                                  let ehi = extract e bits_left_in_dword (e_sz - bits_left_in_dword)
                                      elo = extract e (e_sz - bits_left_in_dword) 0
                                  in (grp ++ [ehi]):(restructure [] (sz - bits_left_in_dword) (elo:es))
                        args'' = [ case grp of
                                      []  -> internalError "SimPackageOpt: optimizeConcats has empty arg group"
                                      [e] -> e
                                      _   -> APrim dummy_id (ATBit (sum (map aSize grp))) PrimConcat grp
                                 | grp <- restructure [] (aSize pty) args'
                                 ]
                    in case args'' of
                         []  -> internalError "SimPackageOpt: optimizeConcats has empty arg list"
                         [e] -> e
                         _   -> APrim pid pty PrimConcat args''
         -- recurse for other
         optConcat (APrim i t o as) = APrim i t o (map optConcat as)
         optConcat (AMethCall t o m as) = AMethCall t o m (map optConcat as)
         optConcat (AFunCall t i f isC as) = AFunCall t i f isC (map optConcat as)
         optConcat e = e
     in mapAExprs optConcat pkg
  
  -- -------------------------
  
  -- Replace dynamic array selection with case-expressions.
  -- Note: This can inline defs, causing there to be unused defs;
  --       we rely on "inlineUses" to remove them.
  --
  expandDynSel :: SimPackage -> SimPackage
  expandDynSel pkg =
    let
        defmap0 = sp_local_defs pkg
        findFn i = adef_expr (findDef defmap0 i)
  
        defmap1 = M.map (aExpandDynSel True findFn) defmap0
  
        -- because method return values are not lifted to local defs,
        -- we need to expand dynamic array selection there, too
        ifc1 = aExpandDynSel True findFn (sp_interface pkg)
  
        -- and they have been inlined into instantiation arguments
        -- XXX Until there is a way to codegen case statements for submodule
        -- XXX instantiation arguments (see mkStateInitList?)
        ss1 = aExpandDynSel False findFn (sp_state_instances pkg)
    in
        pkg { sp_local_defs = defmap1,
              sp_interface = ifc1,
              sp_state_instances = ss1 }
  
  -- -------------------------
  
  -- Convert nested if-expressions into case-expressions in some situations
  -- (using the same "aInsertCaseDef" function from the Verilog backend)
  -- Note: This can inline defs, causing there to be unused defs;
  --       we rely on "inlineUses" to remove them.
  --
  insertCase :: SimPackage -> SimPackage
  insertCase pkg =
    let
        defmap0 = sp_local_defs pkg
        findFn i = adef_expr (findDef defmap0 i)
  
        defmap1 = M.map (aInsertCaseDef True findFn) defmap0
  
        -- because method return values are not lifted to local defs,
        -- we need to infer case expressions there, too
        aInsertCaseIfc ai@(AIDef { aif_value = d }) =
            ai { aif_value = aInsertCaseDef True findFn d }
        aInsertCaseIfc ai@(AIActionValue { aif_value = d }) =
            ai { aif_value = aInsertCaseDef True findFn d }
        aInsertCaseIfc ai = ai
  
        ifc0 = sp_interface pkg
        ifc1 = map aInsertCaseIfc ifc0
    in
        pkg { sp_local_defs = defmap1,
              sp_interface = ifc1 }
  
  -- -------------------------
  
  -- This is monadic because it can report an error
  --
  convertASAny :: ErrorHandle -> Flags -> SimPackage -> IO SimPackage
  convertASAny errh flags apkg = do
    let
        tgt :: String
        tgt = unSpecTo flags
  
        mkVal :: Integer -> IO AExpr
        mkVal n = do
          v <- case tgt of
                 "0" -> return $ 0
                 "1" -> return $ (2^n) - 1
                 "A" -> return $ aaaa n
                 _   -> -- this situation should have been rejected when
                        -- we checked command-line flags
                        bsError errh [(noPosition, EBluesimNoXZ tgt)]
          return $ ASInt dummy_id (ATBit n) (ilSizedHex n v)
  
        cvtASAnyExpr :: AExpr -> IO AExpr
        cvtASAnyExpr (ASAny ty _) = mkVal (aSize ty)
        cvtASAnyExpr (APrim aid ty op args) =
          do args' <- mapM cvtASAnyExpr args
             return $ APrim aid ty op args'
        cvtASAnyExpr (AMethCall ty aid  mid args) =
          do args' <- mapM cvtASAnyExpr args
             return $ AMethCall ty aid mid args'
        cvtASAnyExpr (ANoInlineFunCall ty aid fun args) =
          do args' <- mapM cvtASAnyExpr args
             return $ ANoInlineFunCall ty aid fun args'
        cvtASAnyExpr (AFunCall ty aid fun isC args) =
          do args' <- mapM cvtASAnyExpr args
             return $ AFunCall ty aid fun isC args'
        cvtASAnyExpr expr = return expr
  
    -- convert in the state, rules, ifc, and defs
    -- XXX consider an instance of mapMAExprs for SimPackage?
    --
    let ss = sp_state_instances apkg
        rs = sp_rules apkg
        ifc = sp_interface apkg
        dmap = sp_local_defs apkg
  
    ss' <- mapMAExprs cvtASAnyExpr ss
    rs' <- mapMAExprs cvtASAnyExpr rs
    ifc' <- mapMAExprs cvtASAnyExpr ifc
    dmap' <- map_mapM (mapMAExprs cvtASAnyExpr) dmap
  
    return $ apkg { sp_state_instances = ss'
                  , sp_rules = rs'
                  , sp_interface = ifc'
                  , sp_local_defs = dmap' }
  
  -- -------------------------
  