diff --git a/src/Cryptol/TypeCheck/Subst.hs b/src/Cryptol/TypeCheck/Subst.hs index ed603e5d..e4edc309 100644 --- a/src/Cryptol/TypeCheck/Subst.hs +++ b/src/Cryptol/TypeCheck/Subst.hs @@ -140,18 +140,17 @@ instance TVars a => TVars (TypeMap a) where -- | Apply the substitution to the keys of a type map. apSubstTypeMapKeys :: Subst -> TypeMap a -> TypeMap a -apSubstTypeMapKeys su = go id +apSubstTypeMapKeys su = go (\_ x -> x) id where - go :: (a -> a) -> TypeMap a -> TypeMap a - go atNode TM { .. } = foldl addKey tm' tys + go :: (a -> a -> a) -> (a -> a) -> TypeMap a -> TypeMap a + go merge atNode TM { .. } = foldl addKey tm' tys where - - addKey tm (ty,a) = insertTM ty a tm + addKey tm (ty,a) = insertWithTM merge ty a tm tm' = TM { tvar = Map.fromList vars - , tcon = fmap (lgo atNode) tcon - , trec = fmap (lgo atNode) trec + , tcon = fmap (lgo merge atNode) tcon + , trec = fmap (lgo merge atNode) trec } -- partition out variables that have been replaced with more specific types @@ -164,10 +163,12 @@ apSubstTypeMapKeys su = go id , let a' = atNode a ] - lgo :: (a -> a) -> List TypeMap a -> List TypeMap a - lgo atNode k = k { nil = fmap atNode (nil k) - , cons = go (lgo atNode) (cons k) - } + lgo :: (a -> a -> a) -> (a -> a) -> List TypeMap a -> List TypeMap a + lgo merge atNode k = k { nil = fmap atNode (nil k) + , cons = go (unionTM merge) + (lgo merge atNode) + (cons k) + } {- | WARNING: This instance assumes that the quantified variables in the diff --git a/src/Cryptol/TypeCheck/TypeMap.hs b/src/Cryptol/TypeCheck/TypeMap.hs index 7d497a77..b09a459b 100644 --- a/src/Cryptol/TypeCheck/TypeMap.hs +++ b/src/Cryptol/TypeCheck/TypeMap.hs @@ -33,6 +33,7 @@ class TrieMap m k | m -> k where nullTM :: m a -> Bool lookupTM :: k -> m a -> Maybe a alterTM :: k -> (Maybe a -> Maybe a) -> m a -> m a + unionTM :: (a -> a -> a) -> m a -> m a -> m a toListTM :: m a -> [(k,a)] membersTM :: TrieMap m k => m a -> [a] @@ -69,6 +70,14 @@ instance TrieMap m a => TrieMap (List m) [a] where [ ([], v) | v <- maybeToList (nil m) ] ++ [ (x:xs,v) | (x,m1) <- toListTM (cons m), (xs,v) <- toListTM m1 ] + unionTM f m1 m2 = L { nil = case (nil m1, nil m2) of + (Just x, Just y) -> Just (f x y) + (Just x, _) -> Just x + (_, Just y) -> Just y + _ -> Nothing + , cons = unionTM (unionTM f) (cons m1) (cons m2) + } + instance Ord a => TrieMap (Map a) a where emptyTM = Map.empty @@ -76,6 +85,7 @@ instance Ord a => TrieMap (Map a) a where lookupTM = Map.lookup alterTM = flip Map.alter toListTM = Map.toList + unionTM = Map.unionWith type TypesMap = List TypeMap @@ -115,6 +125,12 @@ instance TrieMap TypeMap Type where [ (TRec (zip fs ts), v) | (fs,m1) <- toListTM (trec m) , (ts,v) <- toListTM m1 ] + unionTM f m1 m2 = TM { tvar = unionTM f (tvar m1) (tvar m2) + , tcon = unionTM (unionTM f) (tcon m1) (tcon m2) + , trec = unionTM (unionTM f) (trec m1) (trec m2) + } + + updSub :: TrieMap m k => k -> (Maybe a -> Maybe a) -> Maybe (m a) -> Maybe (m a) updSub k f = Just . alterTM k f . fromMaybe emptyTM