
   0 | module Data.HashMap.Internal
   1 | 
   2 | import public Data.Hashable
   3 | import Data.HashMap.Bits
   4 | import Data.HashMap.Array
   5 | import Data.HashMap.SparseArray
   6 | import Decidable.Equality
   7 | 
   8 | %default total
   9 | 
  10 | chunkSize : Bits64
  11 | chunkSize = 6
  12 | 
  13 | -- 10 * 6 + 4
  14 | -- 10 groups of 6, then one group of 4
  15 | maxDepth : Bits64
  16 | maxDepth = 10
  17 | 
  18 | {-
  19 | Take 6 bits at a time
  20 | Starting from least significant bits
  21 | -}
  22 | 
  23 | ||| A non-empty dependently-typed hash-array mapped trie
  24 | export
  25 | data HAMT : (key : Type) -> (val : key -> Type) -> Type where
  26 |     Leaf : (hash : Bits64) -> (k : key) -> (v : val k) -> HAMT key val
  27 |     Node : SparseArray (HAMT key val) -> HAMT key val
  28 |     Collision : (hash : Bits64) -> Array (k ** val k) -> HAMT key val
  29 | 
  30 | %name HAMT hamt
  31 | 
  32 | getMask : (depth : Bits64) -> Bits64
  33 | getMask depth = baseMask `shiftL` (depth * chunkSize)
  34 |   where
  35 |     baseMask : Bits64
  36 |     baseMask = 0b111111
  37 | 
  38 | export
  39 | getIndex : (depth : Bits64) -> (hash : Bits64) -> Bits32
  40 | getIndex depth hash = unsafeCast $ (getMask depth .&. hash) `shiftR` (depth * chunkSize)
  41 | 
  42 | export
  43 | singletonWithHash : (hash : Bits64) -> (k : key) -> val k -> HAMT key val
  44 | singletonWithHash = Leaf
  45 | 
  46 | export
  47 | singleton : Hashable key => (k : key) -> val k -> HAMT key val
  48 | singleton k x = singletonWithHash (hash k) k x
  49 | 
  50 | parameters
  51 |     {0 key : Type}
  52 |     {0 val : key -> Type}
  53 |     (keyEq : (x : key) -> (y : key) -> Bool)
  54 | 
  55 |     lookupEntry : (k : key) -> (idx : Bits32) -> List (k ** val k) -> Maybe (Bits32, (k ** val k))
  56 |     lookupEntry k idx [] = Nothing
  57 |     lookupEntry k idx (entry :: xs) = if keyEq k entry.fst
  58 |         then Just (idx, entry)
  59 |         else lookupEntry k (idx + 1) xs
  60 | 
  61 |     export
  62 |     lookupWithHash :
  63 |         (k : key) ->
  64 |         (hash : Bits64) ->
  65 |         (depth : Bits64) ->
  66 |         HAMT key val ->
  67 |         Maybe (k ** val k)
  68 |     lookupWithHash k0 hash0 depth (Leaf hash1 k1 val) = if hash0 == hash1
  69 |         then if keyEq k0 k1
  70 |             then Just (k1 ** val)
  71 |             else Nothing
  72 |         else Nothing
  73 |     lookupWithHash k0 hash0 depth (Node arr) =
  74 |         let idx = getIndex depth hash0
  75 |          in index idx arr >>=
  76 |             lookupWithHash k0 hash0 (assert_smaller depth $ depth + 1)
  77 |     lookupWithHash k0 hash0 depth (Collision hash1 arr) = if hash0 == hash1
  78 |         then
  79 |             let arrL = toList arr
  80 |              in snd <$> lookupEntry k0 0 arrL
  81 |         else Nothing
  82 | 
  83 |     export
  84 |     lookup :
  85 |         Hashable key =>
  86 |         (k : key) ->
  87 |         HAMT key val ->
  88 |         Maybe (k ** val k)
  89 |     lookup k hamt = lookupWithHash k (hash k) 0 hamt
  90 | 
  91 |     -- Pre: hash0 /= hash1
  92 |     export
  93 |     node2 :
  94 |         (tree0 : HAMT key val) ->
  95 |         (hash0 : Bits64) ->
  96 |         (tree1 : HAMT key val) ->
  97 |         (hash1 : Bits64) ->
  98 |         (depth : Bits64) ->
  99 |         HAMT key val
 100 |     node2 hamt0 hash0 hamt1 hash1 depth =
 101 |         let idx0 = getIndex depth hash0
 102 |             idx1 = getIndex depth hash1
 103 |          in if idx0 == idx1
 104 |             then Node $ singleton
 105 |                 (idx0, (node2 hamt0 hash0 hamt1 hash1 (assert_smaller depth $ depth + 1)))
 106 |             else Node $ doubleton (idx0, hamt0) (idx1, hamt1)
 107 | 
 108 |     export
 109 |     insertWithHash :
 110 |         (k : key) ->
 111 |         val k ->
 112 |         (hash : Bits64) ->
 113 |         (depth : Bits64) ->
 114 |         HAMT key val ->
 115 |         HAMT key val
 116 |     insertWithHash k0 val0 hash0 depth hamt@(Leaf hash1 k1 val1) =
 117 |         if hash0 /= hash1
 118 |             then node2 (singletonWithHash hash0 k0 val0) hash0 hamt hash1 depth
 119 |         else if keyEq k0 k1
 120 |             then Leaf hash0 k0 val0
 121 |         else Collision hash0 (fromList [(k0 ** val0), (k1 ** val1)]) 
 122 |     insertWithHash k val hash0 depth (Node arr) =
 123 |         let idx = getIndex depth hash0
 124 |          in case index idx arr of
 125 |             Just hamt => Node $ set idx
 126 |                 (insertWithHash k val hash0 (assert_smaller depth $ depth + 1) hamt)
 127 |                 arr
 128 |             Nothing => Node $ set idx (singletonWithHash hash0 k val) arr 
 129 |     insertWithHash k val hash0 depth hamt@(Collision hash1 arr) =
 130 |         if hash0 == hash1
 131 |             then case lookupEntry k 0 (toList arr) of
 132 |                 Just (idx, _) => Collision hash1 (update arr [(idx, (k ** val))])
 133 |                 Nothing => Collision hash1 (append (k ** val) arr)
 134 |             else node2 (singletonWithHash hash0 k val) hash0 hamt hash1 depth
 135 | 
 136 |     export
 137 |     insert :
 138 |         Hashable key =>
 139 |         (k : key) ->
 140 |         val k ->
 141 |         HAMT key val ->
 142 |         HAMT key val
 143 |     insert k x hamt = insertWithHash k x (hash k) 0 hamt
 144 | 
 145 |     export
 146 |     deleteWithHash :
 147 |         Hashable key =>
 148 |         (k : key) ->
 149 |         (hash : Bits64) ->
 150 |         (depth : Bits64) ->
 151 |         HAMT key val ->
 152 |         Maybe (HAMT key val)
 153 |     deleteWithHash k0 h0 depth hamt@(Leaf h1 k1 _) =
 154 |         if h0 == h1 && keyEq k0 k1
 155 |             then Nothing
 156 |             else Just hamt
 157 |     deleteWithHash k hash depth hamt0@(Node arr) =
 158 |         let idx = getIndex depth hash
 159 |          in case index idx arr of
 160 |             Just hamt1 => case deleteWithHash k hash (depth + 1) (assert_smaller hamt0 hamt1) of
 161 |                 Just hamt2 => Just $ Node $ set idx hamt2 arr
 162 |                 Nothing =>
 163 |                     let arr' = delete idx arr
 164 |                      in case length arr' of
 165 |                         0 => Nothing
 166 |                         1 => case index arr'.array 0 of
 167 |                             Just (Node _) => Just $ Node arr'
 168 |                             hamt2 => hamt2
 169 |                         _ => Just $ Node arr'
 170 |             Nothing => Just hamt0
 171 |     deleteWithHash k h0 depth hamt@(Collision h1 arr) =
 172 |         if h0 == h1
 173 |             then case findIndex (keyEq k . fst) arr of
 174 |                 Nothing => Just hamt
 175 |                 Just idx =>
 176 |                     let arr' = delete idx arr
 177 |                      in case length arr' of
 178 |                         0 => Nothing
 179 |                         1 => map (\(key ** val) => Leaf h1 key val) $ index arr' 0
 180 |                         _ => Just $ Collision h1 arr'
 181 |             else Just hamt
 182 | 
 183 |     export
 184 |     delete :
 185 |         Hashable key =>
 186 |         (k : key) ->
 187 |         HAMT key val ->
 188 |         Maybe (HAMT key val)
 189 |     delete k hamt = deleteWithHash k (hash k) 0 hamt
 190 | 
 191 | export
 192 | trieMap : ({k : _} -> val0 k -> val1 k) -> HAMT key val0 -> HAMT key val1
 193 | trieMap f (Leaf hash k v) = Leaf hash k (f v)
 194 | trieMap f (Node arr) = Node $ assert_total $ map (trieMap f) arr
 195 | trieMap f (Collision hash arr) = Collision hash $ map ({ snd $= f }) arr
 196 | 
 197 | export
 198 | foldWithKey : ((k : _) -> val k -> acc -> acc) -> acc -> HAMT key val -> acc
 199 | foldWithKey f z (Leaf hash k v) = f k v z
 200 | foldWithKey f z (Node arr) = assert_total $ foldr (\trie, acc => foldWithKey f acc trie) z arr
 201 | foldWithKey f z (Collision hash arr) = foldr (\(k ** v), acc => f k v z) z arr
 202 |