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