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.Light.Compiler.AnaDomain
import UHC.Util.Utils
import UHC.Util.Pretty
import UHC.Light.Compiler.AnaDomain.Pretty
import qualified Data.Map as Map
import qualified Data.Set as Set
import Control.Monad
import UHC.Util.Serialize
import Data.Typeable
import Data.Generics (Data)



{-# 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
           , Data, 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
           , Data,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_RelevTy							-- relevance typing
      { libindaspRelevTy 		:: !RelevTy
      }
  | 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
           , Data, Typeable
           )

type LamInfoBindAspMp = Map.Map ACoreBindAspectKeyS LamInfoBindAsp

{-# LINE 129 "src/ehc/LamInfo.chs" #-}
instance PP LamInfoBindAsp where
  pp (LamInfoBindAsp_RelevTy 	t) = "RTy"  >#< pp t
  pp (LamInfoBindAsp_Ty      	t) = "Ty"   >#< pp 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 139 "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
           , Data, Typeable
           )

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

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

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

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

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

emptyLamMp :: LamMp
emptyLamMp = Map.empty

{-# LINE 197 "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 206 "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 228 "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 251 "src/ehc/LamInfo.chs" #-}
lamMpFilterLam :: LamMp -> LamMp
lamMpFilterLam = Map.filter ((>0) . laminfoArity)

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

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

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

{-# LINE 351 "src/ehc/LamInfo.chs" #-}
instance Serialize LamInfoBindAsp where
  sput (LamInfoBindAsp_RelevTy  	a) = sputWord8 0 >> sput a
  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
      0 -> liftM  LamInfoBindAsp_RelevTy  	sget
      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