1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
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 (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)
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 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)))
|