LamInfo.hs

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
           )

{-# 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
           )

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

{-# LINE 97 "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
           )

type LamInfoBindAspMp = Map.Map ACoreBindAspectKeyS LamInfoBindAsp

{-# LINE 131 "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 143 "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
           )

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

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

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

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

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

emptyLamMp :: LamMp
emptyLamMp = Map.empty

{-# LINE 201 "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 210 "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) -> (LamMp -> LamMp -> LamMp) -> LamMp -> LamMp -> 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 232 "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 255 "src/ehc/LamInfo.chs" #-}
lamMpFilterLam :: LamMp -> LamMp
lamMpFilterLam = Map.filter ((>0) . laminfoArity)

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

{-# LINE 263 "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 320 "src/ehc/LamInfo.chs" #-}
initLamMp :: LamMp
initLamMp = emptyLamMp

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

{-# LINE 355 "src/ehc/LamInfo.chs" #-}
instance Serialize LamInfoBindAsp where
  sput (LamInfoBindAsp_Ty 			a) = sputWord8 1 >> sput a
  sput (LamInfoBindAsp_Core 	  a b) = sputWord8 2 >> sput a >> sput b
  sput (LamInfoBindAsp_FusionRole 	a) = sputWord8 3 >> sput a
  sget = do
    t <- sgetWord8
    case t of
      1 -> liftM  LamInfoBindAsp_Ty 		sget
      2 -> liftM2 LamInfoBindAsp_Core 		sget sget
      3 -> liftM  LamInfoBindAsp_FusionRole sget

instance Serialize LamInfo where
  sput (LamInfo a b c) = sput a >> sput b >> sput c
  sget = liftM3 LamInfo  sget sget sget

instance Serialize StackTraceInfo where
  sput (StackTraceInfo_None                ) = sputWord8 0
  sput (StackTraceInfo_HasStackTraceEquiv a) = sputWord8 1 >> sput a
  sput (StackTraceInfo_IsStackTraceEquiv  a) = sputWord8 2 >> sput a
  sget
    = do t <- sgetWord8
         case t of
           0 -> return StackTraceInfo_None
           1 -> liftM  StackTraceInfo_HasStackTraceEquiv sget
           2 -> liftM  StackTraceInfo_IsStackTraceEquiv  sget