From 2836070488d10c89ceeb747688f4497b6f5f3bab Mon Sep 17 00:00:00 2001
From: "Iavor S. Diatchki" <diatchki@galois.com>
Date: Tue, 16 Dec 2014 17:47:22 -0800
Subject: [PATCH] Fixes to applying substitutin to TypeMap keys.

---
 src/Cryptol/TypeCheck/Subst.hs   | 23 ++++++++++++-----------
 src/Cryptol/TypeCheck/TypeMap.hs | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 11 deletions(-)

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