0 | ||| Linear RRB Vector Internals
  1 | module Data.RRBVector1.Unsized.Internal
  2 |
  3 | import Data.Array.Core
  4 | import Data.Array.Indexed
  5 | import Data.Array.Mutable
  6 | import Data.Bits
  7 | import Data.List
  8 | import Data.Linear.Ref1
  9 | import Data.Linear.Token
 10 | import Data.Nat
 11 | import Data.String
 12 | import Derive.Prelude
 13 |
 14 | %default total
 15 | %hide Data.Vect.Quantifiers.All.get
 16 | %language ElabReflection
 17 |
 18 | --------------------------------------------------------------------------------
 19 | --          Internal Utility
 20 | --------------------------------------------------------------------------------
 21 |
 22 | ||| Convenience interface for bitSize that doesn't use an implicit parameter.
 23 | private
 24 | bitSizeOf :  (ty : Type)
 25 |           -> FiniteBits ty
 26 |           => Nat
 27 | bitSizeOf ty = bitSize {a = ty}
 28 |
 29 | --------------------------------------------------------------------------------
 30 | --          Internals
 31 | --------------------------------------------------------------------------------
 32 |
 33 | public export
 34 | Shift : Type
 35 | Shift = Nat
 36 |
 37 | ||| The number of bits used per level.
 38 | export
 39 | blockshift : Shift
 40 | blockshift = 4
 41 |
 42 | ||| The maximum size of a block.
 43 | export
 44 | blocksize : Nat
 45 | blocksize = integerToNat $ 1 `shiftL` blockshift
 46 |
 47 | ||| The mask used to extract the index into the array.
 48 | export
 49 | blockmask : Nat
 50 | blockmask = minus blocksize 1
 51 |
 52 | export
 53 | up :  Shift
 54 |    -> Shift
 55 | up sh = plus sh blockshift
 56 |
 57 | export
 58 | down :  Shift
 59 |      -> Shift
 60 | down sh = minus sh blockshift
 61 |
 62 | export
 63 | radixIndex :  Nat
 64 |            -> Shift
 65 |            -> Nat
 66 | radixIndex i sh = integerToNat ((natToInteger i) `shiftR` sh .&. (natToInteger blockmask))
 67 |
 68 | export
 69 | relaxedRadixIndex :  {n : Nat}
 70 |                   -> IArray n Nat
 71 |                   -> Nat
 72 |                   -> Shift
 73 |                   -> (Nat, Nat)
 74 | relaxedRadixIndex sizes i sh =
 75 |   let guess  = radixIndex i sh -- guess <= idx
 76 |       idx    = loop sizes guess
 77 |       subIdx = case idx == 0 of
 78 |                  True  =>
 79 |                    i
 80 |                  False =>
 81 |                    let idx' = case tryNatToFin $ minus idx 1 of
 82 |                                 Nothing    =>
 83 |                                   assert_total $ idris_crash "Data.RRBVector1.Internal.relaxedRadixIndex: index out of bounds"
 84 |                                 Just idx'' =>
 85 |                                   idx''
 86 |                      in minus i (at sizes idx')
 87 |     in (idx, subIdx)
 88 |   where
 89 |     loop :  IArray n Nat
 90 |          -> Nat
 91 |          -> Nat
 92 |     loop sizes idx =
 93 |       let current = case tryNatToFin idx of
 94 |                       Nothing       =>
 95 |                         assert_total $ idris_crash "Data.RRBVector1.Internal.relaxedRadixIndex.loop: index out of bounds"
 96 |                       Just idx' =>
 97 |                         at sizes idx' -- idx will always be in range for a well-formed tree
 98 |         in case i < current of
 99 |              True  =>
100 |                idx
101 |              False =>
102 |                assert_total $ loop sizes (plus idx 1)
103 |
104 | --------------------------------------------------------------------------------
105 | --          Internal Tree Representation
106 | --------------------------------------------------------------------------------
107 |
108 | ||| A linear internal tree representation.
109 | public export
110 | data Tree1 : (s : Type) -> Type -> Type where
111 |   Balanced   : {bsize : Nat} -> (bsize ** MArray s bsize (Tree1 s a)) -> Tree1 s a
112 |   Unbalanced : {usize : Nat} -> (usize ** MArray s usize (Tree1 s a)) -> IArray usize Nat -> Tree1 s a
113 |   Leaf       : {lsize : Nat} -> (lsize ** MArray s lsize a) -> Tree1 s a
114 |
115 | --------------------------------------------------------------------------------
116 | --          Tree Utilities
117 | --------------------------------------------------------------------------------
118 |
119 | export
120 | singleton :  a
121 |           -> F1 s (MArray s 1 a)
122 | singleton x t =
123 |   marray1 1 x t
124 |
125 | export
126 | singleton' :  Tree1 s a
127 |            -> F1 s (MArray s 1 (Tree1 s a))
128 | singleton' tree t =
129 |   marray1 1 tree t
130 |
131 | export
132 | treeToArray :  Tree1 s a
133 |             -> Either (bsize ** MArray s bsize (Tree1 s a))
134 |                       (usize ** MArray s usize (Tree1 s a))
135 | treeToArray (Balanced   arr)   =
136 |   Left arr
137 | treeToArray (Unbalanced arr _) =
138 |   Right arr
139 | treeToArray (Leaf       _)     =
140 |   assert_total $ idris_crash "Data.RRBVector1.Internal.treeToArray: leaf"
141 |
142 | export
143 | treeToArray' :  Tree1 s a
144 |              -> (lsize ** MArray s lsize a)
145 | treeToArray' (Balanced _)     =
146 |   assert_total $ idris_crash "Data.RRBVector1.Internal.treeToArray': balanced"
147 | treeToArray' (Unbalanced _ _) =
148 |   assert_total $ idris_crash "Data.RRBVector1.Internal.treeToArray': unbalanced"
149 | treeToArray' (Leaf arr)       =
150 |   arr
151 |
152 | export
153 | treeBalanced :  Tree1 s a
154 |              -> Bool
155 | treeBalanced (Balanced   _)   =
156 |   True
157 | treeBalanced (Unbalanced _ _) =
158 |   False
159 | treeBalanced (Leaf       _ )  =
160 |   True
161 |
162 | ||| Computes the size of a tree with shift.
163 | export
164 | treeSize :  Shift
165 |          -> Tree1 s a
166 |          -> F1 s Nat
167 | treeSize t =
168 |   go 0 t
169 |   where
170 |     go :  Shift
171 |        -> Shift
172 |        -> Tree1 s a
173 |        -> F1 s Nat
174 |     go acc _  (Leaf       (l ** arr))       t =
175 |       plus acc l # t
176 |     go acc _  (Unbalanced (u ** arr) sizes) t =
177 |       let i := tryNatToFin $ minus u 1
178 |         in case i of
179 |              Nothing =>
180 |                (assert_total $ idris_crash "Data.RRBVector1.Internal.treeSize: index out of bounds") # t
181 |              Just i' =>
182 |                (plus acc (at sizes i')) # t
183 |     go acc sh (Balanced  (b ** arr))        t =
184 |       let i := minus b 1
185 |         in case tryNatToFin i of
186 |              Nothing =>
187 |                (assert_total $ idris_crash "Data.RRBVector1.Internal.treeSize: index out of bounds") # t
188 |              Just i' =>
189 |                let i'' # t := get arr i' t
190 |                  in go (plus acc (mult i (integerToNat (1 `shiftL` sh))))
191 |                        (down sh)
192 |                        (assert_smaller arr i'')
193 |                        t
194 |
195 | ||| Turns an array into a tree node by computing the sizes of its subtrees.
196 | ||| sh is the shift of the resulting tree.
197 | export
198 | computeSizes :  {n : Nat}
199 |              -> Shift
200 |              -> MArray s n (Tree1 s a)
201 |              -> F1 s (Tree1 s a)
202 | computeSizes sh arr t =
203 |   let isbalanced # t := isBalanced arr t
204 |     in case isbalanced of
205 |          True  =>
206 |            (Balanced {bsize=n} (n ** arr)) # t
207 |          False =>
208 |            let arrnat # t := createArrNat arr t
209 |              in (Unbalanced {usize=n} (n ** arr) arrnat) # t
210 |   where
211 |     createArrNat :  MArray s n (Tree1 s a)
212 |                  -> F1 s (IArray n Nat)
213 |     createArrNat arr t =
214 |       let tant # t := unsafeMArray1 n t
215 |         in go 0 n 0 tant t
216 |       where
217 |         go :  (m, x, acc : Nat)
218 |            -> (an : MArray s n Nat)
219 |            -> {auto v : Ix x n}
220 |            -> F1 s (IArray n Nat)
221 |         go m Z     acc an t =
222 |           freeze an t
223 |         go m (S j) acc an t =
224 |           let j'       # t := getIx arr j t
225 |               treesize # t := treeSize (down sh) j' t
226 |               acc'         := plus acc treesize
227 |             in case tryNatToFin m of
228 |                  Nothing =>
229 |                    (assert_total $ idris_crash "Data.RRBVector.Internal.computeSizes.createArrNat.go: can't convert Nat to Fin") # t
230 |                  Just m' =>
231 |                    let () # t := casswap an m' acc' t
232 |                      in go (S m) j acc' an t
233 |     maxsize : Integer
234 |     maxsize = 1 `shiftL` sh -- the maximum size of a subtree
235 |     lenM1 : Nat
236 |     lenM1 = minus n 1
237 |     isBalanced :  {n : Nat}
238 |                -> MArray s n (Tree1 s a)
239 |                -> F1 s Bool
240 |     isBalanced arr t =
241 |       let bool' # t := go arr 0 t
242 |         in bool' # t
243 |       where
244 |         go :  {n : Nat}
245 |            -> MArray s n (Tree1 s a)
246 |            -> Nat
247 |            -> F1 s Bool
248 |         go arr i t =
249 |           case tryNatToFin i of
250 |             Nothing =>
251 |               (assert_total $ idris_crash "Data.RRBVector.Internal.computeSizes.isBalanced: can't convert Nat to Fin") # t
252 |             Just i' =>
253 |               let subtree # t := get arr i' t
254 |                 in case i < lenM1 of
255 |                      True  =>
256 |                        let go'      # t := assert_total $ go arr (plus i 1) t
257 |                            treesize # t := treeSize (down sh) subtree t
258 |                          in ((natToInteger treesize == maxsize) && go') # t
259 |                      False =>
260 |                        treeBalanced subtree # t
261 |
262 | export
263 | countTrailingZeros :  Nat
264 |                    -> Nat
265 | countTrailingZeros x =
266 |   go 0
267 |   where
268 |     w : Nat
269 |     w = bitSizeOf Int
270 |     go : Nat -> Nat
271 |     go i =
272 |       case i >= w of
273 |         True  =>
274 |           i
275 |         False =>
276 |           case tryNatToFin i of
277 |             Nothing =>
278 |               assert_total $ idris_crash "Data.RRBVector1.Internal.countTrailingZeros: can't convert Nat to Fin"
279 |             Just i' =>
280 |               case testBit (the Int (cast x)) i' of
281 |                 True  =>
282 |                   i
283 |                 False =>
284 |                   assert_total $ go (plus i 1)
285 |
286 | ||| Nat log base 2.
287 | export
288 | log2 :  Nat
289 |      -> Nat
290 | log2 x =
291 |   let bitSizeMinus1 = minus (bitSizeOf Int) 1
292 |     in minus bitSizeMinus1 (countLeadingZeros x)
293 |   where
294 |     countLeadingZeros : Nat -> Nat
295 |     countLeadingZeros x =
296 |       minus (minus w 1) (go (minus w 1))
297 |       where
298 |         w : Nat
299 |         w = bitSizeOf Int
300 |         go : Nat -> Nat
301 |         go i =
302 |           case i < 0 of
303 |             True  =>
304 |               i
305 |             False =>
306 |               case tryNatToFin i of
307 |                 Nothing =>
308 |                   assert_total $ idris_crash "Data.RRBVector1.Internal.log2: can't convert Nat to Fin"
309 |                 Just i' =>
310 |                   case testBit (the Int (cast x)) i' of
311 |                     True  =>
312 |                       i
313 |                     False =>
314 |                       assert_total $ go (minus i 1)
315 |
316 | --------------------------------------------------------------------------------
317 | --          Linear RRB Vectors
318 | --------------------------------------------------------------------------------
319 |
320 | ||| A linear relaxed radix balanced vector (RRBVector1).
321 | ||| It supports fast indexing, iteration, concatenation and splitting.
322 | public export
323 | data RRBVector1 : (s : Type) -> Type -> Type where
324 |   Root  :  Nat   -- size
325 |         -> Shift -- shift (blockshift * height)
326 |         -> Tree1 s a
327 |         -> RRBVector1 s a
328 |   Empty : RRBVector1 s a
329 |