1 files changed, 28 insertions, 2 deletions
diff --git a/src/Data/VarMap.hs b/src/Data/VarMap.hs
index 16c2d27..9c10421 100644
--- a/src/Data/VarMap.hs
+++ b/src/Data/VarMap.hs
@@ -11,9 +11,11 @@ module Data.VarMap (
   delete,
   TypedIdx(..),
   lookup,
+  disjointUnion,
   sink1,
   unsink1,
   subMap,
+  superMap,
 ) where
 
 import Prelude hiding (lookup)
@@ -57,6 +59,11 @@ lookup k (VarMap off _ mp) = do
   idx <- unsafeInt2idx (i + off)
   return (Some (TypedIdx ty idx))
 
+disjointUnion :: Ord k => VarMap k env -> VarMap k env -> VarMap k env
+disjointUnion (VarMap off1 cl1 m1) (VarMap off2 cl2 m2) | off1 == off2 =
+  VarMap off1 (min cl1 cl2) (Map.unionWith (error "VarMap.disjointUnion: overlapping keys") m1 m2)
+disjointUnion vm1 vm2 = disjointUnion (cleanup vm1) (cleanup vm2)
+
 sink1 :: VarMap k env -> VarMap k (t : env)
 sink1 (VarMap off interval mp) = VarMap (off + 1) interval mp
 
@@ -78,13 +85,32 @@ subMap subenv =
         | otherwise = Nothing
   in \(VarMap off _ mp) -> VarMap 0 0 (Map.fromAscList . mapMaybe (modify off) . Map.toAscList $ mp)
 
+superMap :: Eq k => Subenv env env' -> VarMap k env' -> VarMap k env
+superMap subenv =
+  let loop :: Subenv env env' -> Int -> [Int]
+      loop SETop _ = []
+      loop (SEYes sub) i = i : loop sub (i+1)
+      loop (SENo sub) i = loop sub (i+1)
+
+      newIndices = VS.fromList $ loop subenv 0
+      modify off (k, (ty, i))
+        | i + off < 0 = Nothing
+        | i + off >= VS.length newIndices = error "VarMap.superMap: found negative indices in map"
+        | otherwise = let j = newIndices VS.! (i + off)
+                      in if j == -1 then Nothing else Just (k, (ty, j))
+
+  in \(VarMap off _ mp) -> VarMap 0 0 (Map.fromAscList . mapMaybe (modify off) . Map.toAscList $ mp)
+
 maybeCleanup :: VarMap k env -> VarMap k env
-maybeCleanup (VarMap off interval mp)
+maybeCleanup vm@(VarMap _ interval mp)
   | let sz = Map.size mp
   , sz > 0, 2 * interval >= 3 * sz
-  = VarMap off 0 (Map.filter (\(_, i) -> i + off >= 0) mp)
+  = cleanup vm
 maybeCleanup vm = vm
 
+cleanup :: VarMap k env -> VarMap k env
+cleanup (VarMap off _ mp) = VarMap 0 0 (Map.mapMaybe (\(t, i) -> if i + off >= 0 then Just (t, i + off) else Nothing) mp)
+
 unsafeInt2idx :: Int -> Maybe (Idx env t)
 unsafeInt2idx = \n -> if n < 0 then Nothing else Just (go n)
   where