1 files changed, 93 insertions, 0 deletions
diff --git a/src/Data/VarMap.hs b/src/Data/VarMap.hs
new file mode 100644
index 0000000..16c2d27
--- /dev/null
+++ b/src/Data/VarMap.hs
@@ -0,0 +1,93 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeOperators #-}
+module Data.VarMap (
+  VarMap,
+  empty,
+  insert,
+  delete,
+  TypedIdx(..),
+  lookup,
+  sink1,
+  unsink1,
+  subMap,
+) 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))
+
+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 (SEYes 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)
+
+maybeCleanup :: VarMap k env -> VarMap k env
+maybeCleanup (VarMap off interval mp)
+  | let sz = Map.size mp
+  , sz > 0, 2 * interval >= 3 * sz
+  = VarMap off 0 (Map.filter (\(_, i) -> i + off >= 0) mp)
+maybeCleanup vm = vm
+
+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)))