module UHC.Light.Compiler.LamInfo
( StackTraceInfo (..)
, LamInfoBindAsp (..)
, LamInfo (..), emptyLamInfo, emptyLamInfo'
, LamMp, emptyLamMp
, lamMpUnionBindAspMp, lamMpUnionsBindAspMp
, lamMpMergeInto
, lamMpLookupAsp, lamMpLookupAsp2, lamMpLookupLam, lamMpLookupCaf
, lamMpFilterLam, lamMpFilterCaf
, lamMpMergeFrom
, initLamMp
, laminfo1stArgIsStackTrace
, FusionRole (..) )
where
import UHC.Light.Compiler.Base.Common
import UHC.Light.Compiler.AbstractCore
import UHC.Light.Compiler.Ty
import UHC.Light.Compiler.Core
import UHC.Util.Utils
import UHC.Util.Pretty
import UHC.Light.Compiler.Ty.Pretty
import qualified Data.Map as Map
import qualified Data.Set as Set
import Control.Monad
import UHC.Util.Serialize



{-# LINE 68 "src/ehc/LamInfo.chs" #-}
data StackTraceInfo
  = StackTraceInfo_None
  | StackTraceInfo_HasStackTraceEquiv	HsName		-- has a stack traced equivalent
  | StackTraceInfo_IsStackTraceEquiv	HsName		-- is a stack traced equivalent
  deriving ( Show
           , Typeable, Generic
           )

{-# LINE 80 "src/ehc/LamInfo.chs" #-}
-- | The role a value takes in fusion
data FusionRole
  = FusionRole_Fuse			-- fuse this, i.e. inline, turned on by 'fuse f' for f
  | FusionRole_BuildLeft	-- role of g in 'convert g,h'
  | FusionRole_BuildRight	-- role of h in 'convert g,h'
  deriving ( Enum, Show
           , Typeable
           , Generic
           )

{-# LINE 92 "src/ehc/LamInfo.chs" #-}
instance PP FusionRole where
  pp r = pp $ drop l $ show r
       where l = length "FusionRole_"

{-# LINE 98 "src/ehc/LamInfo.chs" #-}
-- | per aspect info
data LamInfoBindAsp
  = LamInfoBindAsp_Ty								-- plain good old type
      { libindaspTy 			:: !Ty
      }
  | LamInfoBindAsp_Core								-- actual Core, should go paired with Ty (?? maybe pair them directly)
      { libindaspMetaLev		:: !MetaLev
      , libindaspCore			:: !CExpr
      }
  | LamInfoBindAsp_FusionRole						-- role in fusion
      { libindaspFusionRole 	:: !FusionRole
      }
  deriving ( Show
           , Typeable, Generic
           )

type LamInfoBindAspMp = Map.Map ACoreBindAspectKeyS LamInfoBindAsp

{-# LINE 132 "src/ehc/LamInfo.chs" #-}
instance PP LamInfoBindAsp where
  pp (LamInfoBindAsp_Ty      	t) = "Ty"   >#< ppTy t
  pp (LamInfoBindAsp_Core    ml	c) = pp "Core" -- >#< pp c -- Core.Pretty uses LamInfo, so module cycle...
  pp (LamInfoBindAsp_FusionRole	r) = "Fuse" >#< pp r

{-# LINE 144 "src/ehc/LamInfo.chs" #-}
-- | per lambda implementation info
data LamInfo
  = LamInfo
      { laminfoArity				:: !Int							-- arity of function
      , laminfoStackTrace  			:: !StackTraceInfo				-- stacktrace
      , laminfoBindAspMp			:: !LamInfoBindAspMp			-- info organized per/keyed on aspect
      }
  deriving ( Show
           , Typeable, Generic
           )

emptyLamInfo' :: LamInfo
emptyLamInfo'
  = LamInfo 0 StackTraceInfo_None
            Map.empty

emptyLamInfo :: LamInfo
emptyLamInfo
  = LamInfo 0 StackTraceInfo_None
            Map.empty

{-# LINE 178 "src/ehc/LamInfo.chs" #-}
instance PP LamInfo where
  pp (LamInfo {laminfoBindAspMp=m}) = ppAssocL $ assocLMapKey ppACBaspKeyS $ Map.toList m

{-# LINE 183 "src/ehc/LamInfo.chs" #-}
laminfo1stArgIsStackTrace :: LamInfo -> Bool
laminfo1stArgIsStackTrace (LamInfo {laminfoStackTrace=StackTraceInfo_IsStackTraceEquiv _}) = True
laminfo1stArgIsStackTrace _                                                                = False

{-# LINE 195 "src/ehc/LamInfo.chs" #-}
type LamMp    = Map.Map HsName LamInfo

emptyLamMp :: LamMp
emptyLamMp = Map.empty

{-# LINE 202 "src/ehc/LamInfo.chs" #-}
-- union, including the aspect map, but arbitrary for the info itself
lamMpUnionBindAspMp :: LamMp -> LamMp -> LamMp
lamMpUnionBindAspMp = Map.unionWith (\i1 i2 -> i1 {laminfoBindAspMp = laminfoBindAspMp i1 `Map.union` laminfoBindAspMp i2})

lamMpUnionsBindAspMp :: [LamMp] -> LamMp
lamMpUnionsBindAspMp = foldr lamMpUnionBindAspMp Map.empty

{-# LINE 211 "src/ehc/LamInfo.chs" #-}
-- propagate from new (left) to prev (right), adding new entries if necessary, combining with mergeL2RInfo, finally combining/choosing maps with mergeL2RMp
lamMpMergeInto
  :: (LamInfo -> LamInfo -> LamInfo)		-- ^ 'mergeL2RInfo', merge info
  -> (LamMp -> LamMp -> LamMp)				-- ^ 'mergeL2RMp', merge map
  -> LamMp									-- ^ new map
  -> LamMp									-- ^ prev map
  -> LamMp
lamMpMergeInto mergeL2RInfo mergeL2RMp newMp prevMp
  = mergeL2RMp newMpMerge prevMp
  where newMpMerge
          = Map.mapWithKey
              (\n i -> maybe i (mergeL2RInfo i) $ Map.lookup n prevMp
              ) newMp

{-# LINE 238 "src/ehc/LamInfo.chs" #-}
lamMpLookupAsp :: HsName -> ACoreBindAspectKeyS -> LamMp -> Maybe LamInfoBindAsp
lamMpLookupAsp n a m
  = fmap snd $ mapLookup2' laminfoBindAspMp n a m

lamMpLookupAsp2 :: ACoreBindRef -> LamMp -> Maybe LamInfoBindAsp
lamMpLookupAsp2 (ACoreBindRef n (Just a)) m = lamMpLookupAsp n a m

lamMpLookupLam :: HsName -> LamMp -> Maybe Int
lamMpLookupLam n m
  = case Map.lookup n m of
      j@(Just (LamInfo {laminfoArity=a})) | a > 0
        -> Just a
      _ -> Nothing

lamMpLookupCaf :: HsName -> LamMp -> Maybe Int
lamMpLookupCaf n m
  = case Map.lookup n m of
      j@(Just (LamInfo {laminfoArity=a})) | a == 0
        -> Just a
      _ -> Nothing

{-# LINE 261 "src/ehc/LamInfo.chs" #-}
lamMpFilterLam :: LamMp -> LamMp
lamMpFilterLam = Map.filter ((>0) . laminfoArity)

lamMpFilterCaf :: LamMp -> LamMp
lamMpFilterCaf = Map.filter ((==0) . laminfoArity)

{-# LINE 269 "src/ehc/LamInfo.chs" #-}
-- | merge info from arbitrary map m into LamMp holding LamInfo's
lamMpMergeFrom
  :: (LamInfo -> Maybe x)					-- extract relevant info from a LamInfo
     -> (Maybe x -> LamInfo -> LamInfo)		-- set the info
     -> (z -> x -> x)						-- merge info from new map and old info
     -> LamInfo								-- default, empty LamInfo
     -> Map.Map HsName z					-- arbitrary map holding info to merge
     -> LamMp -> LamMp
lamMpMergeFrom get set merge empty m lm
  = Map.foldrWithKey (\n z lm -> Map.alter (Just . upd z) n lm)
                    lm m
  where upd z (Just i) = set (Just (merge z $ maybe emptyExtra id $ get i)) i
        upd z Nothing  = set (Just (merge z         emptyExtra           )) empty
        emptyExtra = panicJust "lamMpMergeFrom" $ get $ empty

{-# LINE 326 "src/ehc/LamInfo.chs" #-}
initLamMp :: LamMp
initLamMp = emptyLamMp

{-# LINE 355 "src/ehc/LamInfo.chs" #-}
instance Serialize FusionRole where
  sput = sputEnum8
  sget = sgetEnum8

{-# LINE 361 "src/ehc/LamInfo.chs" #-}
instance Serialize LamInfoBindAsp
instance Serialize LamInfo
instance Serialize StackTraceInfo