0 | ||| Relaxed Radix Balanced Vectors (RRBVector)
   1 | module Data.RRBVector
   2 |
   3 | import public Data.RRBVector.Internal
   4 |
   5 | import Data.Array
   6 | import Data.Array.Core
   7 | import Data.Array.Index
   8 | import Data.Array.Indexed
   9 | import Data.Bits
  10 | import Data.Linear.Ref1
  11 | import Data.Linear.Traverse1
  12 | import Data.List
  13 | import Data.List1
  14 | import Data.Maybe
  15 | import Data.SnocList
  16 | import Data.Vect
  17 | import Data.Zippable
  18 | import Syntax.T1 as T1
  19 |
  20 | %hide Prelude.null
  21 | %hide Prelude.Ops.infixr.(<|)
  22 | %hide Prelude.Ops.infixl.(|>)
  23 |
  24 | %default total
  25 |
  26 | --------------------------------------------------------------------------------
  27 | --          Fixity
  28 | --------------------------------------------------------------------------------
  29 |
  30 | export
  31 | infixr 5 ><
  32 |
  33 | export
  34 | infixr 5 <|
  35 |
  36 | export
  37 | infixl 5 |>
  38 |
  39 | --------------------------------------------------------------------------------
  40 | --          Creating RRB-Vectors
  41 | --------------------------------------------------------------------------------
  42 |
  43 | ||| The empty vector. O(1)
  44 | |||
  45 | export
  46 | empty : RRBVector a
  47 | empty = Empty
  48 |
  49 | ||| A vector with a single element. O(1)
  50 | |||
  51 | export
  52 | singleton :  a
  53 |           -> RRBVector a
  54 | singleton x = Root 1 0 (Leaf $ A 1 $ fill 1 x)
  55 |
  56 | ||| Create a new vector from a list. O(n)
  57 | |||
  58 | export
  59 | fromList :  List a
  60 |          -> RRBVector a
  61 | fromList []  = Empty
  62 | fromList [x] = singleton x
  63 | fromList xs  =
  64 |   case nodes Leaf xs of
  65 |     [tree] =>
  66 |       Root (treeSize 0 tree) 0 tree -- tree is a single leaf
  67 |     xs'    =>
  68 |       assert_smaller xs (iterateNodes blockshift xs')
  69 |   where
  70 |     nodes :  (Array a -> Tree a)
  71 |           -> List a
  72 |           -> List (Tree a)
  73 |     nodes f trees =
  74 |       let (trees', rest) = unsafeAlloc blocksize (go 0 blocksize f trees)
  75 |         in case rest of
  76 |              []    =>
  77 |                [trees']
  78 |              rest' =>
  79 |                (trees' :: nodes f (assert_smaller trees rest'))
  80 |       where
  81 |         go :  (cur,n : Nat)
  82 |            -> (Array a -> Tree a)
  83 |            -> List a
  84 |            -> WithMArray n a (Tree a,List a)
  85 |         go cur n f []        r = T1.do
  86 |           res <- unsafeFreeze r
  87 |           pure $ (f $ force $ take cur $ A n res,[])
  88 |         go cur n f (x :: xs) r =
  89 |           case cur == n of
  90 |             True  => T1.do
  91 |               res <- unsafeFreeze r
  92 |               pure $ (f $ A n res, x :: xs)
  93 |             False =>
  94 |               case tryNatToFin cur of
  95 |                 Nothing   =>
  96 |                   assert_total $ idris_crash "Data.RRBVector.fromList.node: can't convert Nat to Fin"
  97 |                 Just cur' => T1.do
  98 |                   set r cur' x
  99 |                   go (S cur) n f xs r
 100 |     nodes' :  (Array (Tree a) -> Tree a)
 101 |            -> List (Tree a)
 102 |            -> List (Tree a)
 103 |     nodes' f trees =
 104 |       let (trees', rest) = unsafeAlloc blocksize (go 0 blocksize f trees)
 105 |         in case rest of
 106 |              []    =>
 107 |                [trees']
 108 |              rest' =>
 109 |                (trees' :: nodes' f (assert_smaller trees rest'))
 110 |       where
 111 |         go :  (cur,n : Nat)
 112 |            -> (Array (Tree a) -> Tree a)
 113 |            -> List (Tree a)
 114 |            -> WithMArray n (Tree a) (Tree a,List (Tree a))
 115 |         go cur n f []        r = T1.do
 116 |           res <- unsafeFreeze r
 117 |           pure $ (f $ force $ take cur $ A n res,[])
 118 |         go cur n f (x :: xs) r =
 119 |           case cur == n of
 120 |             True  => T1.do
 121 |               res <- unsafeFreeze r
 122 |               pure $ (f $ A n res, x :: xs)
 123 |             False =>
 124 |               case tryNatToFin cur of
 125 |                 Nothing   =>
 126 |                   assert_total $ idris_crash "Data.RRBVector.fromList.node': can't convert Nat to Fin"
 127 |                 Just cur' => T1.do
 128 |                   set r cur' x
 129 |                   go (S cur) n f xs r
 130 |     iterateNodes :  Nat
 131 |                  -> List (Tree a)
 132 |                  -> RRBVector a
 133 |     iterateNodes sh trees =
 134 |       case nodes' Balanced trees of
 135 |         [tree] =>
 136 |           Root (treeSize sh tree) sh tree
 137 |         trees' =>
 138 |           iterateNodes (up sh) (assert_smaller trees trees')
 139 |
 140 | ||| Creates a vector of length n with every element set to x. O(log n)
 141 | |||
 142 | export
 143 | replicate :  Nat
 144 |           -> a
 145 |           -> RRBVector a
 146 | replicate n x =
 147 |   case compare n 0 of
 148 |     LT =>
 149 |       Empty
 150 |     EQ =>
 151 |       Empty
 152 |     GT =>
 153 |       case compare n blocksize of
 154 |         LT =>
 155 |           Root n 0 (Leaf $ A n $ fill n x)
 156 |         EQ =>
 157 |           Root n 0 (Leaf $ A n $ fill n x)
 158 |         GT =>
 159 |           let size' = integerToNat ((natToInteger $ minus n 1) .&. (natToInteger $ plus blockmask 1))
 160 |             in iterateNodes blockshift
 161 |                             (Leaf $ A blocksize $ fill blocksize x)
 162 |                             (Leaf $ A size' $ fill size' x)
 163 |   where
 164 |     iterateNodes :  Shift
 165 |                  -> Tree a
 166 |                  -> Tree a
 167 |                  -> RRBVector a
 168 |     iterateNodes sh full rest =
 169 |       let subtreesm1  = (natToInteger $ minus n 1) `shiftR` sh
 170 |           restsize    = integerToNat (subtreesm1 .&. (natToInteger blockmask))
 171 |           rest'       = Balanced $ A (plus restsize 1) $ append (fill restsize full) (fill 1 rest)
 172 |         in case compare subtreesm1 (natToInteger blocksize) of
 173 |              LT =>
 174 |                Root n sh rest'
 175 |              EQ =>
 176 |                let full' = Balanced (A blocksize $ fill blocksize full)
 177 |                  in iterateNodes (up sh) (assert_smaller full full') (assert_smaller rest rest')
 178 |              GT =>
 179 |                let full' = Balanced (A blocksize $ fill blocksize full)
 180 |                  in iterateNodes (up sh) (assert_smaller full full') (assert_smaller rest rest')
 181 |
 182 | --------------------------------------------------------------------------------
 183 | --          Creating Lists from RRB-Vectors
 184 | --------------------------------------------------------------------------------
 185 |
 186 | ||| Convert a vector to a list. O(n)
 187 | |||
 188 | export
 189 | toList :  RRBVector a
 190 |        -> List a
 191 | toList Empty           = []
 192 | toList (Root _ _ tree) = treeToList tree
 193 |   where
 194 |     treeToList :  Tree a
 195 |                -> List a
 196 |     treeToList (Balanced trees)     = assert_total $ concat (map treeToList (toList trees))
 197 |     treeToList (Unbalanced trees _) = assert_total $ concat (map treeToList (toList trees))
 198 |     treeToList (Leaf arr)           = toList arr
 199 |
 200 | --------------------------------------------------------------------------------
 201 | --          Folds
 202 | --------------------------------------------------------------------------------
 203 |
 204 | export
 205 | foldl :  (b -> a -> b)
 206 |       -> b
 207 |       -> RRBVector a
 208 |       -> b
 209 | foldl f acc = go
 210 |   where
 211 |     foldlTree :  b
 212 |               -> Tree a
 213 |               -> b
 214 |     foldlTree acc' (Balanced arr)     = assert_total $ foldl foldlTree acc' arr
 215 |     foldlTree acc' (Unbalanced arr _) = assert_total $ foldl foldlTree acc' arr
 216 |     foldlTree acc' (Leaf arr)         = assert_total $ foldl f acc' arr
 217 |     go :  RRBVector a
 218 |        -> b
 219 |     go Empty           = acc
 220 |     go (Root _ _ tree) = assert_total $ foldlTree acc tree
 221 |
 222 | export
 223 | foldr :  (a -> b -> b)
 224 |       -> b
 225 |       -> RRBVector a
 226 |       -> b
 227 | foldr f acc = go
 228 |   where
 229 |     foldrTree :  Tree a
 230 |               -> b
 231 |               -> b
 232 |     foldrTree (Balanced arr) acc'     = assert_total $ foldr foldrTree acc' arr
 233 |     foldrTree (Unbalanced arr _) acc' = assert_total $ foldr foldrTree acc' arr
 234 |     foldrTree (Leaf arr) acc'         = assert_total $ foldr f acc' arr
 235 |     go :  RRBVector a
 236 |        -> b
 237 |     go Empty           = acc
 238 |     go (Root _ _ tree) = assert_total $ foldrTree tree acc
 239 |
 240 | --------------------------------------------------------------------------------
 241 | --          Query
 242 | --------------------------------------------------------------------------------
 243 |
 244 | ||| Is the vector empty? O(1)
 245 | |||
 246 | export
 247 | null :  RRBVector a
 248 |      -> Bool
 249 | null Empty = True
 250 | null _     = False
 251 |
 252 | ||| Return the size of a vector. O(1)
 253 | |||
 254 | export
 255 | length :  RRBVector a
 256 |        -> Nat
 257 | length Empty        = 0
 258 | length (Root s _ _) = s
 259 |
 260 | --------------------------------------------------------------------------------
 261 | --          Indexing
 262 | --------------------------------------------------------------------------------
 263 |
 264 | ||| The element at the index or Nothing if the index is out of range. O(log n)
 265 | |||
 266 | export
 267 | lookup :  Nat
 268 |        -> RRBVector a
 269 |        -> Maybe a
 270 | lookup _ Empty               = Nothing
 271 | lookup i (Root size sh tree) =
 272 |   case compare i 0 of
 273 |     LT =>
 274 |       Nothing -- index out of range
 275 |     GT =>
 276 |       case compare i size of
 277 |         EQ =>
 278 |           Nothing -- index out of range
 279 |         GT =>
 280 |           Nothing -- index out of range
 281 |         LT =>
 282 |           Just $ lookupTree i sh tree
 283 |     EQ =>
 284 |       case compare i size of
 285 |         EQ =>
 286 |           Nothing -- index out of range
 287 |         GT =>
 288 |           Nothing -- index out of range
 289 |         LT =>
 290 |           Just $ lookupTree i sh tree
 291 |   where
 292 |     lookupTree :  Nat
 293 |                -> Nat
 294 |                -> Tree a
 295 |                -> a
 296 |     lookupTree i sh (Balanced arr)         =
 297 |       case tryNatToFin (radixIndex i sh) of
 298 |         Nothing =>
 299 |           assert_total $ idris_crash "Data.RRBVector.lookup: can't convert Nat to Fin"
 300 |         Just i' =>
 301 |           assert_total $ lookupTree i (down sh) (at arr.arr i')
 302 |     lookupTree i sh (Unbalanced arr sizes) =
 303 |       let (idx, subidx) = relaxedRadixIndex sizes i sh
 304 |         in case tryNatToFin idx of
 305 |              Nothing   =>
 306 |                assert_total $ idris_crash "Data.RRBVector.lookup: can't convert Nat to Fin"
 307 |              Just idx' =>
 308 |                assert_total $ lookupTree subidx (down sh) (at arr.arr idx')
 309 |     lookupTree i _ (Leaf arr)              =
 310 |       let i' = integerToNat ((natToInteger i) .&. (natToInteger blockmask))
 311 |         in case tryNatToFin i' of
 312 |              Nothing =>
 313 |                assert_total $ idris_crash "Data.RRBVector.lookup: can't convert Nat to Fin"
 314 |              Just i'' =>
 315 |                at arr.arr i''
 316 |
 317 | ||| The element at the index.
 318 | ||| Calls 'idris_crash' if the index is out of range. O(log n)
 319 | |||
 320 | export
 321 | index :  Nat
 322 |       -> RRBVector a
 323 |       -> a
 324 | index i = fromMaybe (assert_total $ idris_crash "Data.RRBVector.index: index out of range") . lookup i
 325 |
 326 | ||| A flipped version of lookup. O(log n)
 327 | |||
 328 | export
 329 | (!?) :  RRBVector a
 330 |      -> Nat
 331 |      -> Maybe a
 332 | (!?) = flip lookup
 333 |
 334 | ||| A flipped version of index. O(log n)
 335 | |||
 336 | export
 337 | (!!) :  RRBVector a
 338 |      -> Nat
 339 |      -> a
 340 | (!!) = flip index
 341 |
 342 | ||| Update the element at the index with a new element.
 343 | ||| If the index is out of range, the original vector is returned. O (log n)
 344 | |||
 345 | export
 346 | update :  Nat
 347 |        -> a
 348 |        -> RRBVector a
 349 |        -> RRBVector a
 350 | update _ _ Empty                 = Empty
 351 | update i x v@(Root size sh tree) =
 352 |   case compare i 0 of
 353 |     LT =>
 354 |       v -- index out of range
 355 |     GT =>
 356 |       case compare i size of
 357 |         EQ =>
 358 |           v -- index out of range
 359 |         GT =>
 360 |           v -- index out of range
 361 |         LT =>
 362 |           Root size sh (updateTree i sh tree)
 363 |     EQ =>
 364 |       case compare i size of
 365 |         EQ =>
 366 |           v -- index out of range
 367 |         GT =>
 368 |           v -- index out of range
 369 |         LT =>
 370 |           Root size sh (updateTree i sh tree)
 371 |   where
 372 |     updateTree :  Nat
 373 |                -> Nat
 374 |                -> Tree a
 375 |                -> Tree a
 376 |     updateTree i sh (Balanced arr)         =
 377 |       case tryNatToFin (radixIndex i sh) of
 378 |         Nothing =>
 379 |           assert_total $ idris_crash "Data.RRBVector.update: can't convert Nat to Fin"
 380 |         Just i' =>
 381 |           assert_total $ Balanced (A arr.size (updateAt i' (updateTree i (down sh)) arr.arr))
 382 |     updateTree i sh (Unbalanced arr sizes) =
 383 |       let (idx, subidx) = relaxedRadixIndex sizes i sh
 384 |         in case tryNatToFin idx of
 385 |              Nothing   =>
 386 |                assert_total $ idris_crash "Data.RRBVector.update: can't convert Nat to Fin"
 387 |              Just idx' =>
 388 |                assert_total $ Unbalanced (A arr.size (updateAt idx' (updateTree subidx (down sh)) arr.arr)) sizes
 389 |     updateTree i _ (Leaf arr)              =
 390 |       let i' = integerToNat ((natToInteger i) .&. (natToInteger blockmask))
 391 |         in case tryNatToFin i' of
 392 |              Nothing =>
 393 |                assert_total $ idris_crash "Data.RRBVector.update: can't convert Nat to Fin"
 394 |              Just i'' =>
 395 |                Leaf (A arr.size (setAt i'' x arr.arr))
 396 |
 397 | ||| Adjust the element at the index by applying the function to it.
 398 | ||| If the index is out of range, the original vector is returned. O(log n)
 399 | |||
 400 | export
 401 | adjust :  Nat
 402 |        -> (a -> a)
 403 |        -> RRBVector a
 404 |        -> RRBVector a
 405 | adjust _ _ Empty                 = Empty
 406 | adjust i f v@(Root size sh tree) =
 407 |   case compare i 0 of
 408 |     LT =>
 409 |       v -- index out of range
 410 |     GT =>
 411 |       case compare i size of
 412 |         EQ =>
 413 |           v -- index out of range
 414 |         GT =>
 415 |           v -- index out of range
 416 |         LT =>
 417 |           Root size sh (adjustTree i sh tree)
 418 |     EQ =>
 419 |       case compare i size of
 420 |         EQ =>
 421 |           v -- index out of range
 422 |         GT =>
 423 |           v -- index out of range
 424 |         LT =>
 425 |           Root size sh (adjustTree i sh tree)
 426 |   where
 427 |     adjustTree :  Nat
 428 |                -> Nat
 429 |                -> Tree a
 430 |                -> Tree a
 431 |     adjustTree i sh (Balanced arr)         =
 432 |       case tryNatToFin (radixIndex i sh) of
 433 |         Nothing =>
 434 |           assert_total $ idris_crash "Data.RRBVector.adjust: can't convert Nat to Fin"
 435 |         Just i' =>
 436 |           assert_total $ Balanced (A arr.size (updateAt i' (adjustTree i (down sh)) arr.arr))
 437 |     adjustTree i sh (Unbalanced arr sizes) =
 438 |       let (idx, subidx) = relaxedRadixIndex sizes i sh
 439 |         in case tryNatToFin idx of
 440 |              Nothing   =>
 441 |                assert_total $ idris_crash "Data.RRBVector.adjust: can't convert Nat to Fin"
 442 |              Just idx' =>
 443 |                assert_total $ Unbalanced (A arr.size (updateAt idx' (adjustTree subidx (down sh)) arr.arr)) sizes
 444 |     adjustTree i _ (Leaf arr)              =
 445 |       let i' = integerToNat ((natToInteger i) .&. (natToInteger blockmask))
 446 |         in case tryNatToFin i' of
 447 |              Nothing =>
 448 |                assert_total $ idris_crash "Data.RRBVector.adjust: can't convert Nat to Fin"
 449 |              Just i'' =>
 450 |                Leaf (A arr.size (updateAt i'' f arr.arr))
 451 |
 452 | private
 453 | normalize :  RRBVector a
 454 |           -> RRBVector a
 455 | normalize v@(Root size sh (Balanced arr))     =
 456 |   case compare arr.size 1 of
 457 |     LT =>
 458 |       v
 459 |     EQ =>
 460 |       case tryNatToFin 0 of
 461 |         Nothing =>
 462 |           assert_total $ idris_crash "Data.RRBVector.normalize: can't convert Nat to Fin"
 463 |         Just i  =>
 464 |           assert_total $ normalize $ Root size (down sh) (at arr.arr i)
 465 |     GT =>
 466 |       v
 467 | normalize v@(Root size sh (Unbalanced arr _)) =
 468 |   case compare arr.size 1 of
 469 |     LT =>
 470 |       v
 471 |     EQ =>
 472 |       case tryNatToFin 0 of
 473 |         Nothing =>
 474 |           assert_total $ idris_crash "Data.RRBVector.normalize: can't convert Nat to Fin"
 475 |         Just i  =>
 476 |           assert_total $ normalize $ Root size (down sh) (at arr.arr i)
 477 |     GT =>
 478 |       v
 479 | normalize v                                   =
 480 |   v
 481 |
 482 | ||| The initial i is n - 1 (the index of the last element in the new tree).
 483 | |||
 484 | private
 485 | takeTree :  Nat
 486 |          -> Shift
 487 |          -> Tree a
 488 |          -> Tree a
 489 | takeTree i sh (Balanced arr) with (radixIndex i sh) | ((plus (radixIndex i sh) 1) <= arr.size) proof eq
 490 |   _ | i' | True  =
 491 |     case tryNatToFin i' of
 492 |       Nothing =>
 493 |         assert_total $ idris_crash "Data.RRBVector.takeTree: can't convert Nat to Fin"
 494 |       Just i'' =>
 495 |         let newarr = force $ take (plus (radixIndex i sh) 1) arr.arr @{lteOpReflectsLTE _ _ eq}
 496 |           in assert_total $ Balanced (A (plus (radixIndex i sh) 1) (updateAt i'' (takeTree i (down sh)) newarr))
 497 |   _ | _  | False =
 498 |     assert_total $ idris_crash "Data.RRBVector.takeTree: index out of bounds"
 499 | takeTree i sh (Unbalanced arr sizes) with (relaxedRadixIndex sizes i sh) | ((plus (fst (relaxedRadixIndex sizes i sh)) 1) <= arr.size) proof eq
 500 |   _ | (idx, subidx) | True  =
 501 |     case tryNatToFin idx of
 502 |       Nothing   =>
 503 |         assert_total $ idris_crash "Data.RRBVector.takeTree: can't convert Nat to Fin"
 504 |       Just idx' =>
 505 |         let newarr = force $ take (plus (fst (relaxedRadixIndex sizes i sh)) 1) arr.arr @{lteOpReflectsLTE _ _ eq}
 506 |           in assert_total $ computeSizes sh (A (plus (fst (relaxedRadixIndex sizes i sh)) 1) (updateAt idx' (takeTree subidx (down sh)) newarr))
 507 |   _ | _             | False =
 508 |     assert_total $ idris_crash "Data.RRBVector.takeTree: index out of bounds"
 509 | takeTree i _ (Leaf arr) with (integerToNat (((natToInteger i) .&. (natToInteger blockmask)) + 1) <= arr.size) proof eq
 510 |   _ | True  =
 511 |     let newarr = force $ take (integerToNat (((natToInteger i) .&. (natToInteger blockmask)) + 1)) arr.arr @{lteOpReflectsLTE _ _ eq}
 512 |       in Leaf (A (integerToNat (((natToInteger i) .&. (natToInteger blockmask)) + 1)) newarr)
 513 |   _ | False =
 514 |     assert_total $ idris_crash "Data.RRBVector.takeTree: index out of bounds"
 515 |
 516 | private
 517 | dropTree :  Nat
 518 |          -> Shift
 519 |          -> Tree a
 520 |          -> Tree a
 521 | dropTree n sh (Balanced arr) =
 522 |   case tryNatToFin 0 of
 523 |     Nothing   =>
 524 |       assert_total $ idris_crash "Data.RRBVector.dropTree: can't convert Nat to Fin"
 525 |     Just zero =>
 526 |       let newarr = force $ drop (radixIndex n sh) arr.arr
 527 |         in assert_total $ computeSizes sh (A (minus arr.size (radixIndex n sh)) (updateAt zero (dropTree n (down sh)) newarr))
 528 | dropTree n sh (Unbalanced arr sizes) =
 529 |   case tryNatToFin 0 of
 530 |     Nothing   =>
 531 |       assert_total $ idris_crash "Data.RRBVector.dropTree: can't convert Nat to Fin"
 532 |     Just zero =>
 533 |       let newarr = force $ drop (fst $ relaxedRadixIndex sizes n sh) arr.arr
 534 |         in assert_total $ computeSizes sh (A (minus arr.size (fst $ relaxedRadixIndex sizes n sh)) (updateAt zero (dropTree (snd $ relaxedRadixIndex sizes n sh) (down sh)) newarr))
 535 | dropTree n _  (Leaf arr) =
 536 |   let n      = integerToNat ((natToInteger n) .&. (natToInteger blockmask))
 537 |       newarr = force $ drop n arr.arr
 538 |     in Leaf (A (minus arr.size n) newarr)
 539 |
 540 | ||| The first i elements of the vector.
 541 | ||| If the vector contains less than or equal to i elements, the whole vector is returned. O(log n)
 542 | |||
 543 | export
 544 | take :  Nat
 545 |      -> RRBVector a
 546 |      -> RRBVector a
 547 | take _ Empty                 = Empty
 548 | take n v@(Root size sh tree) =
 549 |   case compare n 0 of
 550 |     LT =>
 551 |       empty
 552 |     EQ =>
 553 |       empty
 554 |     GT =>
 555 |       case compare n size of
 556 |         LT =>
 557 |           normalize $ Root n sh (takeTree (minus n 1) sh tree)
 558 |         EQ =>
 559 |           v
 560 |         GT =>
 561 |           v
 562 |
 563 | ||| The vector without the first i elements.
 564 | ||| If the vector contains less than or equal to i elements, the empty vector is returned. O(log n)
 565 | |||
 566 | export
 567 | drop :  Nat
 568 |      -> RRBVector a
 569 |      -> RRBVector a
 570 | drop _ Empty                 = Empty
 571 | drop n v@(Root size sh tree) =
 572 |   case compare n 0 of
 573 |     LT =>
 574 |       v
 575 |     EQ =>
 576 |       v
 577 |     GT =>
 578 |       case compare n size of
 579 |         LT =>
 580 |           normalize $ Root (minus size n) sh (dropTree n sh tree)
 581 |         EQ =>
 582 |           empty
 583 |         GT =>
 584 |           empty
 585 |
 586 | ||| Split the vector at the given index. O(log n)
 587 | |||
 588 | export
 589 | splitAt :  Nat
 590 |         -> RRBVector a
 591 |         -> (RRBVector a, RRBVector a)
 592 | splitAt _ Empty                 = (Empty, Empty)
 593 | splitAt n v@(Root size sh tree) =
 594 |   case compare n 0 of
 595 |     LT =>
 596 |       (empty, v)
 597 |     EQ =>
 598 |       (empty, v)
 599 |     GT =>
 600 |       case compare n size of
 601 |         LT =>
 602 |           let left  = normalize $ Root n sh (takeTree (minus n 1) sh tree)
 603 |               right = normalize $ Root (minus size n) sh (dropTree n sh tree)
 604 |             in (left, right)
 605 |         EQ =>
 606 |           (v, empty)
 607 |         GT =>
 608 |           (v, empty)
 609 |
 610 | --------------------------------------------------------------------------------
 611 | --          Deconstruction
 612 | --------------------------------------------------------------------------------
 613 |
 614 | ||| The first element and the vector without the first element, or 'Nothing' if the vector is empty. O(log n)
 615 | |||
 616 | export
 617 | viewl :  RRBVector a
 618 |       -> Maybe (a, RRBVector a)
 619 | viewl Empty             = Nothing
 620 | viewl v@(Root _ _ tree) =
 621 |   let tail = drop 1 v
 622 |     in Just (headTree tree, tail)
 623 |   where
 624 |     headTree : Tree a -> a
 625 |     headTree (Balanced arr)     =
 626 |       case tryNatToFin 0 of
 627 |         Nothing   =>
 628 |           assert_total $ idris_crash "Data.RRBVector.viewl: can't convert Nat to Fin"
 629 |         Just zero =>
 630 |           assert_total $ headTree (at arr.arr zero)
 631 |     headTree (Unbalanced arr _) =
 632 |       case tryNatToFin 0 of
 633 |         Nothing   =>
 634 |           assert_total $ idris_crash "Data.RRBVector.viewl: can't convert Nat to Fin"
 635 |         Just zero =>
 636 |           assert_total $ headTree (at arr.arr zero)
 637 |     headTree (Leaf arr)         =
 638 |       case tryNatToFin 0 of
 639 |         Nothing   =>
 640 |           assert_total $ idris_crash "Data.RRBVector.viewl: can't convert Nat to Fin"
 641 |         Just zero =>
 642 |           at arr.arr zero
 643 |
 644 | ||| The vector without the last element and the last element, or 'Nothing' if the vector is empty. O(log n)
 645 | |||
 646 | export
 647 | viewr :  RRBVector a
 648 |       -> Maybe (RRBVector a, a)
 649 | viewr Empty                = Nothing
 650 | viewr v@(Root size _ tree) =
 651 |   let init = take (minus size 1) v
 652 |     in Just (init, lastTree tree)
 653 |   where
 654 |     lastTree : Tree a -> a
 655 |     lastTree (Balanced arr)     =
 656 |       case tryNatToFin (minus size 1) of
 657 |         Nothing   =>
 658 |           assert_total $ idris_crash "Data.RRBVector.viewr: can't convert Nat to Fin"
 659 |         Just last =>
 660 |           assert_total $ lastTree (at arr.arr last)
 661 |     lastTree (Unbalanced arr _) =
 662 |       case tryNatToFin (minus size 1) of
 663 |         Nothing   =>
 664 |           assert_total $ idris_crash "Data.RRBVector.viewr: can't convert Nat to Fin"
 665 |         Just last =>
 666 |           assert_total $ lastTree (at arr.arr last)
 667 |     lastTree (Leaf arr)         =
 668 |       case tryNatToFin (minus size 1) of
 669 |         Nothing   =>
 670 |           assert_total $ idris_crash "Data.RRBVector.viewr: can't convert Nat to Fin"
 671 |         Just last =>
 672 |           at arr.arr last
 673 |
 674 | --------------------------------------------------------------------------------
 675 | --          Transformation
 676 | --------------------------------------------------------------------------------
 677 |
 678 | ||| Apply the function to every element. O(n)
 679 | |||
 680 | export
 681 | map :  (a -> b)
 682 |     -> RRBVector a
 683 |     -> RRBVector b
 684 | map _ Empty               = Empty
 685 | map f (Root size sh tree) = Root size sh (mapTree tree)
 686 |   where
 687 |     mapTree : Tree a -> Tree b
 688 |     mapTree (Balanced arr)         =
 689 |       assert_total $ Balanced (map mapTree arr)
 690 |     mapTree (Unbalanced arr sizes) =
 691 |       assert_total $ Unbalanced (map mapTree arr) sizes
 692 |     mapTree (Leaf arr)             =
 693 |       Leaf (map f arr)
 694 |
 695 | ||| Reverse the vector. O(n)
 696 | |||
 697 | export
 698 | reverse :  RRBVector a
 699 |         -> RRBVector a
 700 | reverse v =
 701 |   case compare (length v) 1 of
 702 |     LT =>
 703 |       v
 704 |     EQ =>
 705 |       v
 706 |     GT =>
 707 |       case fromList $ toList v of
 708 |         Nothing =>
 709 |           assert_total $ idris_crash "Data.RRBVector.reverse: can't convert to List1"
 710 |         Just v' =>
 711 |           fromList $ forget $ reverse v'
 712 |
 713 | ||| Take two vectors and return a vector of corresponding pairs.
 714 | ||| If one input is longer, excess elements are discarded from the right end. O(min(n1,n2))
 715 | |||
 716 | export
 717 | zip :  RRBVector a
 718 |     -> RRBVector b
 719 |     -> RRBVector (a, b)
 720 | zip v1 v2 =
 721 |   case fromList $ toList v1 of
 722 |     Nothing  =>
 723 |       assert_total $ idris_crash "Data.RRBVector.zip: can't convert to List1"
 724 |     Just v1' =>
 725 |       case fromList $ toList v2 of
 726 |         Nothing  =>
 727 |           assert_total $ idris_crash "Data.RRBVector.zip: can't convert to List1"
 728 |         Just v2' =>
 729 |           fromList $ forget $ zip v1' v2'
 730 |
 731 | --------------------------------------------------------------------------------
 732 | --          Concatenation
 733 | --------------------------------------------------------------------------------
 734 |
 735 | ||| Create a new tree with shift sh.
 736 | |||
 737 | private
 738 | newBranch :  a
 739 |           -> Shift
 740 |           -> Tree a
 741 | newBranch x 0  = Leaf (singleton x)
 742 | newBranch x sh = assert_total $ Balanced (singleton $ newBranch x (down sh))
 743 |
 744 | ||| Add an element to the left end of the vector. O(log n)
 745 | |||
 746 | export
 747 | (<|) :  a
 748 |      -> RRBVector a
 749 |      -> RRBVector a
 750 | x <| Empty             = singleton x
 751 | x <| Root size sh tree =
 752 |   case compare insertshift sh of
 753 |     LT =>
 754 |       Root (plus size 1) sh (consTree sh tree)
 755 |     EQ =>
 756 |       Root (plus size 1) sh (consTree sh tree)
 757 |     GT =>
 758 |       let new = A 2 $ array $ fromList [(newBranch x sh), tree]
 759 |         in Root (plus size 1) insertshift (computeSizes insertshift new)
 760 |   where
 761 |     -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
 762 |     -- the size is computed for efficient calculation of the shift in a balanced subtree
 763 |     computeShift :  Nat
 764 |                  -> Nat
 765 |                  -> Nat
 766 |                  -> Tree a
 767 |                  -> Nat
 768 |     computeShift sz sh min (Balanced _)          =
 769 |       -- @sz - 1@ is the index of the last element
 770 |       let hishift  = let comp = mult (log2 (minus sz 1) `div` blockshift) blockshift  -- the shift of the root when normalizing
 771 |                        in case compare comp 0 of
 772 |                             LT =>
 773 |                               0
 774 |                             EQ =>
 775 |                               0
 776 |                             GT =>
 777 |                               comp
 778 |           hi       = (natToInteger $ minus sz 1) `shiftR` hishift -- the length of the root node when normalizing minus 1
 779 |           newshift = case compare hi (natToInteger blockmask) of
 780 |                        LT =>
 781 |                          hishift
 782 |                        EQ =>
 783 |                          plus hishift blockshift
 784 |                        GT =>
 785 |                          plus hishift blockshift
 786 |         in case compare newshift sh of
 787 |              LT =>
 788 |                newshift
 789 |              EQ =>
 790 |                newshift
 791 |              GT =>
 792 |                min
 793 |     computeShift _ sh min (Unbalanced arr sizes) =
 794 |       let sz'     = case tryNatToFin 0 of
 795 |                       Nothing   =>
 796 |                         assert_total $ idris_crash "Data.RRBVector.(<|).computeShift.Unbalanced: can't convert Nat to Fin"
 797 |                       Just zero =>
 798 |                         at sizes.arr zero -- the size of the first subtree
 799 |           newtree = case tryNatToFin 0 of
 800 |                       Nothing   =>
 801 |                         assert_total $ idris_crash "Data.RRBVector.(<|).computeShift.Unbalanced: can't convert Nat to Fin"
 802 |                       Just zero =>
 803 |                         at arr.arr zero
 804 |           newmin  = case compare arr.size blocksize of
 805 |                       LT =>
 806 |                         sh
 807 |                       EQ =>
 808 |                         min
 809 |                       GT =>
 810 |                         min
 811 |         in assert_total $ computeShift sz' (down sh) newmin newtree
 812 |     computeShift _ _ min (Leaf arr)              =
 813 |       case compare arr.size blocksize of
 814 |         LT =>
 815 |           0
 816 |         EQ =>
 817 |           min
 818 |         GT =>
 819 |           min
 820 |     insertshift : Nat
 821 |     insertshift = computeShift size sh (up sh) tree
 822 |     consTree :  Nat
 823 |              -> Tree a
 824 |              -> Tree a
 825 |     consTree sh (Balanced arr)     =
 826 |       case compare sh insertshift of
 827 |         LT =>
 828 |           case tryNatToFin 0 of
 829 |             Nothing   =>
 830 |               assert_total $ idris_crash "Data.RRBVector.(<|).consTree.Balanced: can't convert Nat to Fin"
 831 |             Just zero =>
 832 |               assert_total $ computeSizes sh (A arr.size $ updateAt zero (consTree (down sh)) arr.arr)
 833 |         EQ =>
 834 |           computeSizes sh (A (S arr.size) (append (fill 1 (newBranch x (down sh))) arr.arr))
 835 |         GT =>
 836 |           case tryNatToFin 0 of
 837 |             Nothing   =>
 838 |               assert_total $ idris_crash "Data.RRBVector.(<|).consTree.Balanced: can't convert Nat to Fin"
 839 |             Just zero =>
 840 |               assert_total $ computeSizes sh (A arr.size $ updateAt zero (consTree (down sh)) arr.arr)
 841 |     consTree sh (Unbalanced arr _) =
 842 |       case compare sh insertshift of
 843 |         LT =>
 844 |           case tryNatToFin 0 of
 845 |             Nothing   =>
 846 |               assert_total $ idris_crash "Data.RRBVector.(<|).consTree.Unbalanced: can't convert Nat to Fin"
 847 |             Just zero =>
 848 |               assert_total $ computeSizes sh (A arr.size $ updateAt zero (consTree (down sh)) arr.arr)
 849 |         EQ =>
 850 |           computeSizes sh (A (S arr.size) (append (fill 1 (newBranch x (down sh))) arr.arr))
 851 |         GT =>
 852 |           case tryNatToFin 0 of
 853 |             Nothing   =>
 854 |               assert_total $ idris_crash "Data.RRBVector.(<|).consTree.Unbalanced: can't convert Nat to Fin"
 855 |             Just zero =>
 856 |               assert_total $ computeSizes sh (A arr.size $ updateAt zero (consTree (down sh)) arr.arr)
 857 |     consTree _ (Leaf arr)          =
 858 |       Leaf (A (S arr.size) (append (fill 1 x) arr.arr))
 859 |
 860 | ||| Add an element to the right end of the vector. O(log n)
 861 | |||
 862 | export
 863 | (|>) :  RRBVector a
 864 |      -> a
 865 |      -> RRBVector a
 866 | Empty             |> x = singleton x
 867 | Root size sh tree |> x =
 868 |   case compare insertshift sh of
 869 |     LT =>
 870 |       Root (plus size 1) sh (snocTree sh tree)
 871 |     EQ =>
 872 |       Root (plus size 1) sh (snocTree sh tree)
 873 |     GT =>
 874 |       let new = A 2 $ array $ fromList [tree,(newBranch x sh)]
 875 |         in Root (plus size 1) insertshift (computeSizes insertshift new)
 876 |   where
 877 |     -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
 878 |     -- the size is computed for efficient calculation of the shift in a balanced subtree
 879 |     computeShift :  Nat
 880 |                  -> Nat
 881 |                  -> Nat
 882 |                  -> Tree a
 883 |                  -> Nat
 884 |     computeShift sz sh min (Balanced _)          =
 885 |       -- @sz - 1@ is the index of the last element
 886 |       let newshift = mult (countTrailingZeros sz `div` blockshift) blockshift
 887 |         in case compare newshift sh of
 888 |              LT =>
 889 |                newshift
 890 |              EQ =>
 891 |                newshift
 892 |              GT =>
 893 |                min
 894 |     computeShift _ sh min (Unbalanced arr sizes) =
 895 |       let lastidx = minus arr.size 1
 896 |           sz'     = case tryNatToFin lastidx of
 897 |                       Nothing       =>
 898 |                         assert_total $ idris_crash "Data.RRBVector.(|>).computeShift.Unbalanced: can't convert Nat to Fin"
 899 |                       Just lastidx' =>
 900 |                         case tryNatToFin $ minus lastidx 1 of
 901 |                           Nothing        =>
 902 |                             assert_total $ idris_crash "Data.RRBVector.(|>).computeShift.Unbalanced: can't convert Nat to Fin"
 903 |                           Just lastidx'' =>
 904 |                             minus (at sizes.arr lastidx') (at sizes.arr lastidx'')
 905 |           newtree = case tryNatToFin lastidx of
 906 |                       Nothing       =>
 907 |                         assert_total $ idris_crash "Data.RRBVector.(|>).computeShift.Unbalanced: can't convert Nat to Fin"
 908 |                       Just lastidx' =>
 909 |                         at arr.arr lastidx'
 910 |           newmin  = case compare arr.size blocksize of
 911 |                       LT =>
 912 |                         sh
 913 |                       EQ =>
 914 |                         min
 915 |                       GT =>
 916 |                         min
 917 |         in assert_total $ computeShift sz' (down sh) newmin newtree
 918 |     computeShift _ _ min (Leaf arr)              =
 919 |       case compare arr.size blocksize of
 920 |         LT =>
 921 |           0
 922 |         EQ =>
 923 |           min
 924 |         GT =>
 925 |           min
 926 |     insertshift : Nat
 927 |     insertshift = computeShift size sh (up sh) tree
 928 |     snocTree :  Nat
 929 |              -> Tree a
 930 |              -> Tree a
 931 |     snocTree sh (Balanced arr) =
 932 |       case compare sh insertshift of
 933 |         LT =>
 934 |           case tryNatToFin $ minus arr.size 1 of
 935 |             Nothing   =>
 936 |               assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Balanced: can't convert Nat to Fin"
 937 |             Just lastidx =>
 938 |               assert_total $ Balanced (A arr.size $ updateAt lastidx (snocTree (down sh)) arr.arr)
 939 |         EQ =>
 940 |           Balanced (A (plus arr.size 1) (append arr.arr (fill 1 (newBranch x (down sh))))) -- the current subtree is fully balanced
 941 |         GT =>
 942 |           case tryNatToFin $ minus arr.size 1 of
 943 |             Nothing   =>
 944 |               assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Balanced: can't convert Nat to Fin"
 945 |             Just lastidx =>
 946 |               assert_total $ Balanced (A arr.size $ updateAt lastidx (snocTree (down sh)) arr.arr)
 947 |     snocTree sh (Unbalanced arr sizes) =
 948 |       case compare sh insertshift of
 949 |         LT =>
 950 |           case tryNatToFin $ minus arr.size 1 of
 951 |             Nothing       =>
 952 |               assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Unbalanced: can't convert Nat to Fin"
 953 |             Just lastidxa =>
 954 |               case tryNatToFin $ minus sizes.size 1 of
 955 |                 Nothing       =>
 956 |                   assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Unbalanced: can't convert Nat to Fin"
 957 |                 Just lastidxs =>
 958 |                   let lastsize = plus (at sizes.arr lastidxs) 1
 959 |                     in assert_total $ Unbalanced (A arr.size (updateAt lastidxa (snocTree (down sh)) arr.arr))
 960 |                                                  (A sizes.size (setAt lastidxs lastsize sizes.arr))
 961 |         EQ =>
 962 |           case tryNatToFin $ minus sizes.size 1 of
 963 |             Nothing      =>
 964 |               assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Unbalanced: can't convert Nat to Fin"
 965 |             Just lastidx =>
 966 |               let lastsize = plus (at sizes.arr lastidx) 1
 967 |                 in assert_total $ Unbalanced (A (plus arr.size 1) (append arr.arr (fill 1 (newBranch x (down sh)))))
 968 |                                              (A (plus sizes.size 1) (append sizes.arr (fill 1 lastsize)))
 969 |         GT =>
 970 |           case tryNatToFin $ minus arr.size 1 of
 971 |             Nothing       =>
 972 |               assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Unbalanced: can't convert Nat to Fin"
 973 |             Just lastidxa =>
 974 |               case tryNatToFin $ minus sizes.size 1 of
 975 |                 Nothing       =>
 976 |                   assert_total $ idris_crash "Data.RRBVector.(|>).snocTree.Unbalanced: can't convert Nat to Fin"
 977 |                 Just lastidxs =>
 978 |                   let lastsize = plus (at sizes.arr lastidxs) 1
 979 |                     in assert_total $ Unbalanced (A arr.size (updateAt lastidxa (snocTree (down sh)) arr.arr))
 980 |                                                  (A sizes.size (setAt lastidxs lastsize sizes.arr))
 981 |     snocTree _ (Leaf arr) = Leaf (A (plus arr.size 1) (append arr.arr (fill 1 x)))
 982 |
 983 | ||| Concatenates two vectors. O(log(max(n1,n2)))
 984 | |||
 985 | export
 986 | (><) :  RRBVector a
 987 |      -> RRBVector a
 988 |      -> RRBVector a
 989 | Empty                >< v                    = v
 990 | v                    >< Empty                = v
 991 | Root size1 sh1 tree1 >< Root size2 sh2 tree2 =
 992 |   let upmaxshift = case compare sh1 sh2 of
 993 |                      LT =>
 994 |                        up sh2
 995 |                      EQ =>
 996 |                        up sh1
 997 |                      GT =>
 998 |                        up sh1
 999 |       newarr     = mergeTrees tree1 sh1 tree2 sh2
1000 |     in normalize $ Root (plus size1 size2) upmaxshift (computeSizes upmaxshift newarr)
1001 |   where
1002 |     viewlArr : Array (Tree a) -> (Tree a, Array (Tree a))
1003 |     viewlArr arr =
1004 |       case tryNatToFin 0 of
1005 |         Nothing   =>
1006 |           assert_total $ idris_crash "Data.RRBVector.(><).viewlArr: can't convert Nat to Fin"
1007 |         Just zero =>
1008 |           (at arr.arr zero, drop 1 arr)
1009 |     viewrArr : Array (Tree b) -> (Array (Tree b), Tree b)
1010 |     viewrArr arr =
1011 |       case tryNatToFin $ minus arr.size 1 of
1012 |         Nothing   =>
1013 |           assert_total $ idris_crash "Data.RRBVector.(><).viewrArr: can't convert Nat to Fin"
1014 |         Just last =>
1015 |           (take (minus arr.size 1) arr, at arr.arr last)
1016 |     mergeRebalance' :  Shift
1017 |                     -> Array (Tree a)
1018 |                     -> Array (Tree a)
1019 |                     -> Array (Tree a)
1020 |                     -> (Tree a -> Array (Tree a))
1021 |                     -> (Array (Tree a) -> Tree a)
1022 |                     -> Array (Tree a)
1023 |     mergeRebalance' sh left center right extract construct =
1024 |       run1 $ \t =>
1025 |         let nodecounter    # t := ref1 Z t
1026 |             subtreecounter # t := ref1 Z t
1027 |             newnode        # t := ref1 Lin t
1028 |             newsubtree     # t := ref1 Lin t
1029 |             newroot        # t := ref1 Lin t
1030 |             ()             # t := mergeRebalanceSubtree' sh nodecounter subtreecounter newnode newsubtree newroot extract construct (toList left ++ toList center ++ toList right) t
1031 |             newnode'       # t := read1 newnode t
1032 |             ()             # t := casmod1 newsubtree (\y => y :< (construct $ A (SnocSize newnode')
1033 |                                                                                 (snocConcat newnode'))
1034 |                                                      ) t                
1035 |             newsubtree'    # t := read1 newsubtree t
1036 |             ()             # t := casmod1 newroot (\y => y :< (computeSizes sh (fromList (cast {to=List (Tree a)} newsubtree')))
1037 |                                                   ) t
1038 |             newroot'       # t := read1 newroot t
1039 |           in fromList (cast {to=List (Tree a)} newroot') # t
1040 |       where
1041 |         mergeRebalanceSubtreeNodeCounter :  Ref s Nat
1042 |                                          -> Ref s Nat
1043 |                                          -> Ref s (SnocList (Array (Tree a)))
1044 |                                          -> Ref s (SnocList (Tree a))
1045 |                                          -> (Array (Tree a) -> Tree a)
1046 |                                          -> F1' s
1047 |         mergeRebalanceSubtreeNodeCounter nodecounter subtreecounter newnode newsubtree construct t =
1048 |           let newnode' # t := read1 newnode t
1049 |               ()       # t := casmod1 newsubtree (\y => y :< (construct $ A (SnocSize newnode')
1050 |                                                                             (snocConcat newnode'))
1051 |                                                  ) t
1052 |               ()       # t := write1 newnode Lin t
1053 |               ()       # t := write1 nodecounter Z t
1054 |             in casmod1 subtreecounter (\y => y + 1) t
1055 |         mergeRebalanceRootSubtreeCounter :  Shift
1056 |                                          -> Ref s Nat
1057 |                                          -> Ref s (SnocList (Tree a))
1058 |                                          -> Ref s (SnocList (Tree a))
1059 |                                          -> F1' s
1060 |         mergeRebalanceRootSubtreeCounter sh subtreecounter newsubtree newroot t =
1061 |           let newsubtree' # t := read1 newsubtree t
1062 |               ()          # t := casmod1 newroot (\y => y :< (computeSizes sh (fromList (cast {to=List (Tree a)} newsubtree')))
1063 |                                                  ) t
1064 |               ()          # t := write1 newsubtree Lin t
1065 |             in write1 subtreecounter Z t
1066 |         mergeRebalanceSubtree''' :  Shift
1067 |                                  -> Ref s Nat
1068 |                                  -> Ref s Nat
1069 |                                  -> Ref s (SnocList (Array (Tree a)))
1070 |                                  -> Ref s (SnocList (Tree a))
1071 |                                  -> Ref s (SnocList (Tree a))
1072 |                                  -> (Array (Tree a) -> Tree a)
1073 |                                  -> Tree a
1074 |                                  -> F1' s
1075 |         mergeRebalanceSubtree''' sh nodecounter subtreecounter newnode newsubtree newroot construct extractedsubtree t =
1076 |           let nodecounter'    # t := read1 nodecounter t
1077 |               ()              # t := when1 (nodecounter' == blocksize) (mergeRebalanceSubtreeNodeCounter nodecounter subtreecounter newnode newsubtree construct) t
1078 |               subtreecounter' # t := read1 subtreecounter t
1079 |               ()              # t := when1 (subtreecounter' == blocksize) (mergeRebalanceRootSubtreeCounter sh subtreecounter newsubtree newroot) t
1080 |               ()              # t := casmod1 newnode (\y => y :< (fill 1 extractedsubtree)
1081 |                                                      ) t
1082 |             in casmod1 nodecounter (\y => y + 1) t        
1083 |         mergeRebalanceSubtree'' :  Shift
1084 |                                 -> Ref s Nat
1085 |                                 -> Ref s Nat
1086 |                                 -> Ref s (SnocList (Array (Tree a)))
1087 |                                 -> Ref s (SnocList (Tree a))
1088 |                                 -> Ref s (SnocList (Tree a))
1089 |                                 -> (Tree a -> Array (Tree a))
1090 |                                 -> (Array (Tree a) -> Tree a)
1091 |                                 -> Tree a
1092 |                                 -> F1' s
1093 |         mergeRebalanceSubtree'' sh nodecounter subtreecounter newnode newsubtree newroot extract construct subtree t =
1094 |           traverse1_ (mergeRebalanceSubtree''' sh nodecounter subtreecounter newnode newsubtree newroot construct) (extract subtree) t
1095 |         mergeRebalanceSubtree' :  Shift
1096 |                                -> Ref s Nat
1097 |                                -> Ref s Nat
1098 |                                -> Ref s (SnocList (Array (Tree a)))
1099 |                                -> Ref s (SnocList (Tree a))
1100 |                                -> Ref s (SnocList (Tree a))
1101 |                                -> (Tree a -> Array (Tree a))
1102 |                                -> (Array (Tree a) -> Tree a)
1103 |                                -> List (Tree a)
1104 |                                -> F1' s
1105 |         mergeRebalanceSubtree' sh nodecounter subtreecounter newnode newsubtree newroot extract construct leftcenterright t =
1106 |           traverse1_ (mergeRebalanceSubtree'' sh nodecounter subtreecounter newnode newsubtree newroot extract construct) leftcenterright t
1107 |     mergeRebalance'' :  Shift
1108 |                      -> Array (Tree a)
1109 |                      -> Array (Tree a)
1110 |                      -> Array (Tree a)
1111 |                      -> (Tree a -> Array a)
1112 |                      -> (Array a -> Tree a)
1113 |                      -> Array (Tree a)
1114 |     mergeRebalance'' sh left center right extract construct =
1115 |       run1 $ \t =>
1116 |         let nodecounter    # t := ref1 Z t
1117 |             subtreecounter # t := ref1 Z t
1118 |             newnode        # t := ref1 Lin t
1119 |             newsubtree     # t := ref1 Lin t
1120 |             newroot        # t := ref1 Lin t
1121 |             ()             # t := mergeRebalanceSubtree' sh nodecounter subtreecounter newnode newsubtree newroot extract construct (toList left ++ toList center ++ toList right) t
1122 |             newnode'       # t := read1 newnode t
1123 |             ()             # t := casmod1 newsubtree (\y => y :< (construct $ A (SnocSize newnode')
1124 |                                                                                 (snocConcat newnode'))
1125 |                                                      ) t                
1126 |             newsubtree'    # t := read1 newsubtree t
1127 |             ()             # t := casmod1 newroot (\y => y :< (computeSizes sh (fromList (cast {to=List (Tree a)} newsubtree')))
1128 |                                                   ) t
1129 |             newroot'       # t := read1 newroot t
1130 |           in fromList (cast {to=List (Tree a)} newroot') # t
1131 |       where
1132 |         mergeRebalanceSubtreeNodeCounter :  Ref s Nat
1133 |                                          -> Ref s Nat
1134 |                                          -> Ref s (SnocList (Array a))
1135 |                                          -> Ref s (SnocList (Tree a))
1136 |                                          -> (Array a -> Tree a)
1137 |                                          -> F1' s
1138 |         mergeRebalanceSubtreeNodeCounter nodecounter subtreecounter newnode newsubtree construct t =
1139 |           let newnode' # t := read1 newnode t
1140 |               ()       # t := casmod1 newsubtree (\y => y :< (construct $ A (SnocSize newnode')
1141 |                                                                             (snocConcat newnode'))
1142 |                                                  ) t
1143 |               ()       # t := write1 newnode Lin t
1144 |               ()       # t := write1 nodecounter Z t
1145 |             in casmod1 subtreecounter (\y => y + 1) t
1146 |         mergeRebalanceRootSubtreeCounter :  Shift
1147 |                                          -> Ref s Nat
1148 |                                          -> Ref s (SnocList (Tree a))
1149 |                                          -> Ref s (SnocList (Tree a))
1150 |                                          -> F1' s
1151 |         mergeRebalanceRootSubtreeCounter sh subtreecounter newsubtree newroot t =
1152 |           let newsubtree' # t := read1 newsubtree t
1153 |               ()          # t := casmod1 newroot (\y => y :< (computeSizes sh (fromList (cast {to=List (Tree a)} newsubtree')))
1154 |                                                  ) t
1155 |               ()          # t := write1 newsubtree Lin t
1156 |             in write1 subtreecounter Z t
1157 |         mergeRebalanceSubtree''' :  Shift
1158 |                                  -> Ref s Nat
1159 |                                  -> Ref s Nat
1160 |                                  -> Ref s (SnocList (Array a))
1161 |                                  -> Ref s (SnocList (Tree a))
1162 |                                  -> Ref s (SnocList (Tree a))
1163 |                                  -> (Array a -> Tree a)
1164 |                                  -> a
1165 |                                  -> F1' s
1166 |         mergeRebalanceSubtree''' sh nodecounter subtreecounter newnode newsubtree newroot construct extractedsubtree t =
1167 |           let nodecounter'    # t := read1 nodecounter t
1168 |               ()              # t := when1 (nodecounter' == blocksize) (mergeRebalanceSubtreeNodeCounter nodecounter subtreecounter newnode newsubtree construct) t
1169 |               subtreecounter' # t := read1 subtreecounter t
1170 |               ()              # t := when1 (subtreecounter' == blocksize) (mergeRebalanceRootSubtreeCounter sh subtreecounter newsubtree newroot) t
1171 |               ()              # t := casmod1 newnode (\y => y :< (fill 1 extractedsubtree)
1172 |                                                      ) t
1173 |             in casmod1 nodecounter (\y => y + 1) t        
1174 |         mergeRebalanceSubtree'' :  Shift
1175 |                                 -> Ref s Nat
1176 |                                 -> Ref s Nat
1177 |                                 -> Ref s (SnocList (Array a))
1178 |                                 -> Ref s (SnocList (Tree a))
1179 |                                 -> Ref s (SnocList (Tree a))
1180 |                                 -> (Tree a -> Array a)
1181 |                                 -> (Array a -> Tree a)
1182 |                                 -> Tree a
1183 |                                 -> F1' s
1184 |         mergeRebalanceSubtree'' sh nodecounter subtreecounter newnode newsubtree newroot extract construct subtree t =
1185 |           traverse1_ (mergeRebalanceSubtree''' sh nodecounter subtreecounter newnode newsubtree newroot construct) (extract subtree) t
1186 |         mergeRebalanceSubtree' :  Shift
1187 |                                -> Ref s Nat
1188 |                                -> Ref s Nat
1189 |                                -> Ref s (SnocList (Array a))
1190 |                                -> Ref s (SnocList (Tree a))
1191 |                                -> Ref s (SnocList (Tree a))
1192 |                                -> (Tree a -> Array a)
1193 |                                -> (Array a -> Tree a)
1194 |                                -> List (Tree a)
1195 |                                -> F1' s
1196 |         mergeRebalanceSubtree' sh nodecounter subtreecounter newnode newsubtree newroot extract construct leftcenterright t =
1197 |           traverse1_ (mergeRebalanceSubtree'' sh nodecounter subtreecounter newnode newsubtree newroot extract construct) leftcenterright t
1198 |     mergeRebalance :  Shift
1199 |                    -> Array (Tree a)
1200 |                    -> Array (Tree a)
1201 |                    -> Array (Tree a)
1202 |                    -> Array (Tree a)
1203 |     mergeRebalance sh left center right =
1204 |       case compare sh blockshift of
1205 |         LT =>
1206 |           assert_total $ mergeRebalance' sh left center right treeToArray (computeSizes (down sh))
1207 |         EQ =>
1208 |           assert_total $ mergeRebalance'' sh left center right (\(Leaf arr) => arr) Leaf
1209 |         GT =>
1210 |           assert_total $ mergeRebalance' sh left center right treeToArray (computeSizes (down sh))
1211 |     mergeTrees :  Tree a
1212 |                -> Nat
1213 |                -> Tree a
1214 |                -> Nat
1215 |                -> Array (Tree a)
1216 |     mergeTrees tree1@(Leaf arr1) _   tree2@(Leaf arr2) _   =
1217 |       case compare arr1.size blocksize of
1218 |         LT =>
1219 |           let arr' = A (plus arr1.size arr2.size) (append arr1.arr arr2.arr)
1220 |             in case compare arr'.size blocksize of
1221 |                  LT =>
1222 |                    singleton $ Leaf arr'
1223 |                  EQ =>
1224 |                    singleton $ Leaf arr'
1225 |                  GT =>
1226 |                    let (left, right) = (take blocksize arr',drop blocksize arr')
1227 |                        lefttree      = Leaf left
1228 |                        righttree     = Leaf right
1229 |                      in A 2 $ fromPairs 2 lefttree [(1,righttree)]
1230 |         EQ =>
1231 |           A 2 $ fromPairs 2 tree1 [(1,tree2)]
1232 |         GT =>
1233 |           let arr' = A (plus arr1.size arr2.size) (append arr1.arr arr2.arr)
1234 |             in case compare arr'.size blocksize of
1235 |                  LT =>
1236 |                    singleton $ Leaf arr'
1237 |                  EQ =>
1238 |                    singleton $ Leaf arr'
1239 |                  GT =>
1240 |                    let (left, right) = (take blocksize arr',drop blocksize arr')
1241 |                        lefttree      = Leaf left
1242 |                        righttree     = Leaf right
1243 |                      in A 2 $ fromPairs 2 lefttree [(1,righttree)]
1244 |     mergeTrees tree1             sh1 tree2             sh2 =
1245 |       case compare sh1 sh2 of
1246 |         LT =>
1247 |           let right                  = treeToArray tree2
1248 |               (righthead, righttail) = viewlArr right
1249 |               merged                 = assert_total $ mergeTrees tree1 sh1 righthead (down sh2)
1250 |             in mergeRebalance sh2 empty merged righttail
1251 |         GT =>
1252 |           let left                 = treeToArray tree1
1253 |               (leftinit, leftlast) = viewrArr left
1254 |               merged               = assert_total $ mergeTrees leftlast (down sh1) tree2 sh2
1255 |             in mergeRebalance sh1 leftinit merged empty
1256 |         EQ =>
1257 |           let left                   = treeToArray tree1
1258 |               right                  = treeToArray tree2
1259 |               (leftinit, leftlast)   = viewrArr left
1260 |               (righthead, righttail) = viewlArr right
1261 |               merged                 = assert_total $ mergeTrees leftlast (down sh1) righthead (down sh2)
1262 |             in mergeRebalance sh1 leftinit merged righttail
1263 |
1264 | ||| Insert an element at the given index, shifting the rest of the vector over.
1265 | ||| If the index is negative, add the element to the left end of the vector.
1266 | ||| If the index is bigger than or equal to the length of the vector, add the element to the right end of the vector. O(log n)
1267 | |||
1268 | export
1269 | insertAt :  Nat
1270 |          -> a
1271 |          -> RRBVector a
1272 |          -> RRBVector a
1273 | insertAt i x v =
1274 |   let (left, right) = splitAt i v
1275 |     in (left |> x) >< right
1276 |
1277 | ||| Delete the element at the given index.
1278 | ||| If the index is out of range, return the original vector. O(log n)
1279 | |||
1280 | export
1281 | deleteAt :  Nat
1282 |          -> RRBVector a
1283 |          -> RRBVector a
1284 | deleteAt i v =
1285 |   let (left, right) = splitAt (plus i 1) v
1286 |     in take i left >< right
1287 |
1288 | --------------------------------------------------------------------------------
1289 | --          Show Utilities (RRB-Vector)
1290 | --------------------------------------------------------------------------------
1291 |
1292 | ||| Show the full representation of the vector.
1293 | |||
1294 | export
1295 | showRRBVectorRep :  Show a
1296 |                  => Show (Tree a)
1297 |                  => Show (RRBVector a)
1298 |                  => RRBVector a
1299 |                  -> String
1300 | showRRBVectorRep Empty            =
1301 |   ""
1302 | showRRBVectorRep (Root size sh t) =
1303 |   "RRBVector "    ++
1304 |   "{ "            ++
1305 |   "Size = "       ++
1306 |   (show size)     ++
1307 |   ", Shift = "    ++
1308 |   (show sh)       ++
1309 |   ", Tree = "     ++
1310 |   (showTreeRep t) ++
1311 |   "}"
1312 |
1313 | --------------------------------------------------------------------------------
1314 | --          Interfaces (RRBVector)
1315 | --------------------------------------------------------------------------------
1316 |
1317 | export
1318 | Eq a => Eq (RRBVector a) where
1319 |   xs == ys = length xs == length ys && Data.RRBVector.toList xs == Data.RRBVector.toList ys
1320 |
1321 | export
1322 | Ord a => Ord (RRBVector a) where
1323 |   compare xs ys = compare (Data.RRBVector.toList xs) (Data.RRBVector.toList ys)
1324 |
1325 | export
1326 | Functor RRBVector where
1327 |   map f v = map f v
1328 |
1329 | export
1330 | Foldable RRBVector where
1331 |   foldl f z           = Data.RRBVector.foldl f z
1332 |   foldr f z           = Data.RRBVector.foldr f z
1333 |   null                = null
1334 |
1335 | export
1336 | Applicative RRBVector where
1337 |   pure      = singleton
1338 |   fs <*> xs = Data.RRBVector.foldl (\acc, f => acc >< map f xs) empty fs
1339 |
1340 | export
1341 | Semigroup (RRBVector a) where
1342 |   (<+>) = (><)
1343 |
1344 | export
1345 | Semigroup (RRBVector a) => Monoid (RRBVector a) where
1346 |   neutral = empty
1347 |
1348 | export
1349 | Monad RRBVector where
1350 |   xs >>= f = Data.RRBVector.foldl (\acc, x => acc >< f x) empty xs
1351 |