diff options
Diffstat (limited to 'compiler/coreSyn/TrieMap.lhs')
| -rw-r--r-- | compiler/coreSyn/TrieMap.lhs | 36 |
1 files changed, 34 insertions, 2 deletions
diff --git a/compiler/coreSyn/TrieMap.lhs b/compiler/coreSyn/TrieMap.lhs index d8a134ed87..e481886460 100644 --- a/compiler/coreSyn/TrieMap.lhs +++ b/compiler/coreSyn/TrieMap.lhs @@ -486,7 +486,10 @@ data TypeMap a , tm_app :: TypeMap (TypeMap a) , tm_fun :: TypeMap (TypeMap a) , tm_tc_app :: NameEnv (ListMap TypeMap a) - , tm_forall :: TypeMap (BndrMap a) } + , tm_forall :: TypeMap (BndrMap a) + , tm_tylit :: TyLitMap a + } + instance Outputable a => Outputable (TypeMap a) where ppr m = text "TypeMap elts" <+> ppr (foldTypeMap (:) [] m) @@ -499,7 +502,8 @@ wrapEmptyTypeMap = TM { tm_var = emptyTM , tm_app = EmptyTM , tm_fun = EmptyTM , tm_tc_app = emptyNameEnv - , tm_forall = EmptyTM } + , tm_forall = EmptyTM + , tm_tylit = emptyTyLitMap } instance TrieMap TypeMap where type Key TypeMap = Type @@ -519,6 +523,7 @@ lkT env ty m go (AppTy t1 t2) = tm_app >.> lkT env t1 >=> lkT env t2 go (FunTy t1 t2) = tm_fun >.> lkT env t1 >=> lkT env t2 go (TyConApp tc tys) = tm_tc_app >.> lkNamed tc >=> lkList (lkT env) tys + go (LiteralTy l) = tm_tylit >.> lkTyLit l go (ForAllTy tv ty) = tm_forall >.> lkT (extendCME env tv) ty >=> lkBndr env tv ----------------- @@ -535,6 +540,8 @@ xtT env (ForAllTy tv ty) f m = m { tm_forall = tm_forall m |> xtT (extendCME e xtT env (TyConApp tc tys) f m = m { tm_tc_app = tm_tc_app m |> xtNamed tc |>> xtList (xtT env) tys f } +xtT _ (LiteralTy l) f m = m { tm_tylit = tm_tylit m |> xtTyLit l f } + fdT :: (a -> b -> b) -> TypeMap a -> b -> b fdT _ EmptyTM = \z -> z fdT k m = foldTM k (tm_var m) @@ -542,6 +549,31 @@ fdT k m = foldTM k (tm_var m) . foldTM (foldTM k) (tm_fun m) . foldTM (foldTM k) (tm_tc_app m) . foldTM (foldTM k) (tm_forall m) + . foldTyLit k (tm_tylit m) + + + +------------------------ +data TyLitMap a + = EmptyTLM + | TLM { tlm_number :: Map.Map Integer a } + +emptyTyLitMap :: TyLitMap a +emptyTyLitMap = EmptyTLM + +lkTyLit :: TyLit -> TyLitMap a -> Maybe a +lkTyLit l = + case l of + NumberTyLit n -> tlm_number >.> Map.lookup n + +xtTyLit :: TyLit -> XT a -> TyLitMap a -> TyLitMap a +xtTyLit l f m = + case l of + NumberTyLit n -> m { tlm_number = tlm_number m |> Map.alter f n } + +foldTyLit :: (a -> b -> b) -> TyLitMap a -> b -> b +foldTyLit l m x = Map.fold l x (tlm_number m) + \end{code} |