1 files changed, 119 insertions, 0 deletions
diff --git a/src/Data/VarMap.hs b/src/Data/VarMap.hs
new file mode 100644
index 0000000..2712b08
--- /dev/null
+++ b/src/Data/VarMap.hs
@@ -0,0 +1,119 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+module Data.VarMap (
+  VarMap,
+  empty,
+  insert,
+  delete,
+  TypedIdx(..),
+  lookup,
+  disjointUnion,
+  sink1,
+  unsink1,
+  subMap,
+  superMap,
+) where
+
+import Prelude hiding (lookup)
+
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map)
+import Data.Maybe (mapMaybe)
+import Data.Some
+import qualified Data.Vector.Storable as VS
+import Unsafe.Coerce
+
+import AST.Env
+import AST.Types
+import AST.Weaken
+
+
+type role VarMap _ nominal  -- ensure that 'env' is not phantom
+data VarMap k (env :: [Ty]) =
+  VarMap Int  -- ^ Global offset; must be added to any value in the map in order to get the proper index
+         Int  -- ^ Time since last cleanup
+         (Map k (Some STy, Int))
+deriving instance Show k => Show (VarMap k env)
+
+empty :: VarMap k env
+empty = VarMap 0 0 Map.empty
+
+insert :: Ord k => k -> STy t -> Idx env t -> VarMap k env -> VarMap k env
+insert k ty idx (VarMap off interval mp) =
+  maybeCleanup $ VarMap off (interval + 1) (Map.insert k (Some ty, idx2int idx - off) mp)
+
+delete :: Ord k => k -> VarMap k env -> VarMap k env
+delete k (VarMap off interval mp) =
+  maybeCleanup $ VarMap off (interval + 1) (Map.delete k mp)
+
+data TypedIdx env t = TypedIdx (STy t) (Idx env t)
+  deriving (Show)
+
+lookup :: Ord k => k -> VarMap k env -> Maybe (Some (TypedIdx env))
+lookup k (VarMap off _ mp) = do
+  (Some ty, i) <- Map.lookup k mp
+  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
+
+unsink1 :: VarMap k (t : env) -> VarMap k env
+unsink1 (VarMap off interval mp) = VarMap (off - 1) interval mp
+
+subMap :: Eq k => Subenv env env' -> VarMap k env -> VarMap k env'
+subMap subenv =
+  let bools = let loop :: Subenv env env' -> [Bool]
+                  loop SETop = []
+                  loop (SEYesR sub) = True : loop sub
+                  loop (SENo sub) = False : loop sub
+              in VS.fromList $ loop subenv
+      newIndices = VS.init $ VS.scanl' (\n b -> if b then n + 1 else n) (0 :: Int) bools
+      modify off (k, (ty, i))
+        | i + off < 0 = Nothing
+        | i + off >= VS.length bools = error "VarMap.subMap: found negative indices in map"
+        | bools VS.! (i + off) = Just (k, (ty, newIndices VS.! (i + off)))
+        | 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 (SEYesR 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 vm@(VarMap _ interval mp)
+  | let sz = Map.size mp
+  , sz > 0, 2 * interval >= 3 * sz
+  = 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
+    go :: Int -> Idx env t
+    go 0 = unsafeCoerce IZ
+    go n = unsafeCoerce (IS (go (n-1)))