1 | module Data.RRBVector.Internal

  3 | import Data.Array

  4 | import Data.Array.Core

  5 | import Data.Array.Index

  6 | import Data.Array.Indexed

  7 | import Data.Bits

  8 | import Data.Fin

  9 | import Data.List

 10 | import Data.Nat

 11 | import Data.String

 12 | import Derive.Prelude

 13 | import Syntax.T1 as T1

 15 | %default total

 16 | %language ElabReflection

 25 | bitSizeOf :  (ty : Type)

 26 |           -> FiniteBits ty

 27 |           => Nat

 28 | bitSizeOf ty = bitSize {a = ty}

 34 | export %inline

 35 | lastAt :  {n : Nat}

 36 |        -> IArray (S n) a

 37 |        -> a

 38 | lastAt arr =

 39 |   atNat arr n

 57 | public export

 58 | record RelaxedIndex (count : Nat) where

 59 |   constructor MkRelaxedIndex

 60 |   child : Fin count

 61 |   offset : Nat

 67 | public export

 68 | Shift : Type

 69 | Shift = Nat

 74 | blockshift : Shift

 75 | blockshift = 4

 80 | blocksize : Nat

 81 | blocksize = integerToNat $ 1 `shiftL` blockshift

 86 | blockmask : Nat

 87 | blockmask = minus blocksize 1

 90 | up :  Shift

 91 |    -> Shift

 92 | up sh = plus sh blockshift

 95 | down :  Shift

 96 |      -> Shift

 97 | down sh = minus sh blockshift

100 | radixIndex :  Nat

101 |            -> Shift

102 |            -> Nat

103 | radixIndex i sh = integerToNat ((natToInteger i) `shiftR` sh .&. (natToInteger blockmask))

111 |   ||| A nonempty collection of child nodes for a balanced RRB tree node.

123 |   public export

124 |   data Children : Type -> Type where

125 |     MkChildren :  {n : Nat}

126 |                -> {auto 0 nonEmpty : LT 0 n}

127 |                -> {auto 0 withinBlock : LTE n Data.RRBVector.Internal.blocksize}

128 |                -> IArray n (Tree a)

129 |                -> Children a

131 |   ||| A nonempty collection of child nodes for a relaxed RRB tree node,

144 |   public export

145 |   data RelaxedChildren : Type -> Type where

146 |     MkRelaxedChildren :  {n : Nat}

147 |                       -> {auto 0 nonEmpty : LT 0 n}

148 |                       -> {auto 0 withinBlock : LTE n Data.RRBVector.Internal.blocksize}

149 |                       -> IArray n (Tree a)

150 |                       -> IArray n Nat

151 |                       -> RelaxedChildren a

153 |   ||| The internal tree representation of an RRB vector.

167 |   public export

168 |   data Tree : Type -> Type where

169 |     Balanced   :  Children a

170 |                -> Tree a

171 |     Unbalanced :  RelaxedChildren a

172 |                -> Tree a

173 |     Leaf       :  Array a

174 |                -> Tree a

186 | export %inline

187 | childrenToArray :  Children a

188 |                 -> Array (Tree a)

189 | childrenToArray (MkChildren {n} arr) =

190 |   A n arr

198 | export %inline

199 | relaxedChildrenToArray :  RelaxedChildren a

200 |                        -> Array (Tree a)

201 | relaxedChildrenToArray (MkRelaxedChildren {n} children _) =

202 |   A n children

207 | export %inline

208 | relaxedSizesToArray :  RelaxedChildren a

209 |                     -> Array Nat

210 | relaxedSizesToArray (MkRelaxedChildren {n} _ sizes) =

211 |   A n sizes

220 | null :  Tree a

221 |      -> Bool

222 | null (Balanced _)   =

223 |   False

224 | null (Unbalanced _) =

225 |   False

226 | null (Leaf arr)     =

227 |   null arr

234 | foldl :  (b -> a -> b)

235 |       -> b

236 |       -> Tree a

237 |       -> b

238 | foldl f acc tree =

239 |   foldlTree acc tree

240 |   where

241 |     foldlTree :  b

242 |               -> Tree a

243 |               -> b

244 |     foldlTree acc' (Balanced (MkChildren {n} arr))            =

245 |       assert_total (foldl foldlTree acc' (A n arr))

246 |     foldlTree acc' (Unbalanced (MkRelaxedChildren {n} arr _)) =

247 |       assert_total (foldl foldlTree acc' (A n arr))

248 |     foldlTree acc' (Leaf arr)                                 =

249 |       assert_total (foldl f acc' arr)

252 | foldr :  (a -> b -> b)

253 |       -> b

254 |       -> Tree a

255 |       -> b

256 | foldr f acc tree =

257 |   foldrTree tree acc

258 |   where

259 |     foldrTree :  Tree a

260 |               -> b

261 |               -> b

262 |     foldrTree (Balanced (MkChildren {n} arr)) acc'            =

263 |       assert_total (foldr foldrTree acc' (A n arr))

264 |     foldrTree (Unbalanced (MkRelaxedChildren {n} arr _)) acc' =

265 |       assert_total (foldr foldrTree acc' (A n arr))

266 |     foldrTree (Leaf arr) acc'                                 =

267 |       assert_total (foldr f acc' arr)

274 | toList :  Tree a

275 |        -> List a

276 | toList (Balanced (MkChildren {n} arr))            =

277 |   assert_total (concat $ map toList $ toList (A n arr))

278 | toList (Unbalanced (MkRelaxedChildren {n} arr _)) =

279 |   assert_total (concat $ map toList $ toList (A n arr))

280 | toList (Leaf arr)                                 =

281 |   toList arr

287 | public export

288 | Show a => Show (Tree a) where

289 |   show (Balanced children)   =

290 |     assert_total ("Balanced " ++ show (childrenToArray children))

291 |   show (Unbalanced children) =

292 |     assert_total ("Unbalanced " ++ show (relaxedChildrenToArray children))

293 |   show (Leaf arr)            =

294 |     "Leaf " ++ show arr

296 | public export

297 | Foldable Tree where

298 |   foldl f z = Data.RRBVector.Internal.foldl f z

299 |   foldr f z = Data.RRBVector.Internal.foldr f z

300 |   toList    = Data.RRBVector.Internal.toList

301 |   null      = Data.RRBVector.Internal.null

303 | public export

304 | Eq a => Eq (Tree a) where

305 |   Balanced xs == Balanced ys     =

306 |     assert_total (childrenToArray xs == childrenToArray ys)

307 |   Unbalanced xs == Unbalanced ys =

308 |     assert_total (relaxedChildrenToArray xs == relaxedChildrenToArray ys)

309 |   Leaf xs == Leaf ys             =

310 |     xs == ys

311 |   _ == _                         =

312 |     False

314 | public export

315 | Ord a => Ord (Tree a) where

316 |   compare tree1 tree2 =

317 |     compare (Data.RRBVector.Internal.toList tree1) (Data.RRBVector.Internal.toList tree2)

323 | public export

324 | showTreeRep :  Show a

325 |             => Show (Tree a)

326 |             => Tree a

327 |             -> String

328 | showTreeRep (Balanced children)   =

329 |   assert_total ("Balanced " ++ show (toList $ childrenToArray children))

330 | showTreeRep (Unbalanced children) =

331 |   assert_total ("Unbalanced " ++ show (toList $ relaxedChildrenToArray children))

332 | showTreeRep (Leaf elems)          =

333 |   assert_total ("Leaf " ++ show (toList elems))

340 | singleton :  a

341 |           -> Array a

342 | singleton x =

343 |   A 1 $ fill 1 x

346 | treeToArray :  Tree a

347 |             -> Array (Tree a)

348 | treeToArray (Balanced children)   =

349 |   childrenToArray children

350 | treeToArray (Unbalanced children) =

351 |   relaxedChildrenToArray children

352 | treeToArray (Leaf _)              =

353 |   assert_total (idris_crash "Data.RRBVector.Internal.treeToArray: leaf")

356 | treeBalanced :  Tree a

357 |              -> Bool

358 | treeBalanced (Balanced _)   =

359 |   True

360 | treeBalanced (Unbalanced _) =

361 |   False

362 | treeBalanced (Leaf _)       =

363 |   True

368 | treeSize :  Shift

369 |          -> Tree a

370 |          -> Nat

371 | treeSize =

372 |   go 0

373 |   where

374 |     go :  Shift

375 |        -> Shift

376 |        -> Tree a

377 |        -> Nat

378 |     go acc _ (Leaf arr)                                         =

379 |       plus acc arr.size

380 |     go acc _ (Unbalanced (MkRelaxedChildren {n = S k} _ sizes)) =

381 |       plus acc (lastAt sizes)

382 |     go acc sh (Balanced (MkChildren {n = S k} children))        =

383 |       let subtreeSize : Nat

384 |           subtreeSize = integerToNat (1 `shiftL` sh)

385 |           acc'        : Nat

386 |           acc'        = plus acc (mult k subtreeSize)

387 |           child       : Tree a

388 |           child       = lastAt children

389 |        in go acc' (down sh) (assert_smaller children child)

414 | relaxedRadixIndex :  {n : Nat}

415 |                   -> {auto 0 nonEmpty : LT 0 n}

416 |                   -> IArray n Nat

417 |                   -> Nat

418 |                   -> Shift

419 |                   -> RelaxedIndex n

420 | relaxedRadixIndex {n = Z} {nonEmpty} sizes i sh impossible

421 | relaxedRadixIndex {n = S k} sizes i sh =

422 |   let guess     : Nat

423 |       guess     = radixIndex i sh

424 |       0 guessLT : LT guess (S k)

425 |       guessLT   = believe_me ()

426 |       child     : Fin (S k)

427 |       child     = natToFinLT guess @{guessLT}

428 |     in assert_total (loop child)

429 |   where

430 |     ||| Compute the logical index relative to the selected child.

436 |     childOffset :  Fin (S k)

437 |                 -> Nat

438 |     childOffset FZ =

439 |       i

440 |     childOffset (FS previous) =

441 |       minus i (at sizes (weaken previous))

442 |     ||| Search forward through the cumulative size table for the first child

448 |     loop :  Fin (S k)

449 |          -> RelaxedIndex (S k)

450 |     loop child =

451 |       let current : Nat

452 |           current = at sizes child

453 |        in case i < current of

454 |             True  =>

455 |               MkRelaxedIndex child (childOffset child)

456 |             False =>

457 |               let next      : Nat

458 |                   next      = S (finToNat child)

459 |                   0 nextLT  : LT next (S k)

460 |                   nextLT    = believe_me ()

461 |                   nextchild : Fin (S k)

462 |                   nextchild = natToFinLT next @{nextLT}

463 |                 in assert_total (loop nextchild)

475 | computeSizes :  Shift

476 |              -> Children a

477 |              -> Tree a

478 | computeSizes sh children@(MkChildren {n} {nonEmpty} {withinBlock} trees) =

479 |   case isBalanced n of

480 |     True =>

481 |       Balanced children

482 |     False =>

483 |       let sizes : IArray n Nat

484 |           sizes = unsafeAlloc n (loop n 0)

485 |         in Unbalanced (MkRelaxedChildren {nonEmpty = nonEmpty} {withinBlock = withinBlock} trees sizes)

486 |   where

487 |     ||| Fill the cumulative subtree-size table from left to right.

492 |     loop :  (remaining : Nat)

493 |          -> {auto pos : Ix remaining n}

494 |          -> Nat

495 |          -> WithMArray n Nat (IArray n Nat)

496 |     loop Z acc r = T1.do

497 |       unsafeFreeze r

498 |     loop (S k) {pos} acc r =

499 |       let subtree : Tree a

500 |           subtree = ix trees k

501 |           acc'    : Nat

502 |           acc'    = plus acc (treeSize (down sh) subtree)

503 |           dst     : Fin n

504 |           dst     = ixToFin pos

505 |        in T1.do

506 |             set r dst acc'

507 |             assert_total $ loop k acc' r

508 |     ||| Maximum logical size of a full child subtree at this level.

510 |     maxsize : Integer

511 |     maxsize = 1 `shiftL` sh

512 |     ||| Determine whether the children can use the compact balanced-node

515 |     isBalanced :  (remaining : Nat)

516 |                -> {auto pos : Ix remaining n}

517 |                -> Bool

518 |     isBalanced Z     =

519 |       True

520 |     isBalanced (S Z) =

521 |       treeBalanced (ix trees Z)

522 |     isBalanced (S (S k)) =

523 |       let subtree : Tree a

524 |           subtree = ix trees (S k)

525 |        in assert_total ((natToInteger $ treeSize (down sh) subtree) == maxsize && isBalanced (S k))

537 | countTrailingZeros :  Nat

538 |                    -> Nat

539 | countTrailingZeros x =

540 |   go (bitSizeOf Int)

541 |   where

542 |     value : Int

543 |     value = cast x

544 |     ||| Scan bit positions from least significant to most significant.

546 |     go :  (remaining : Nat)

547 |        -> {auto pos : Ix remaining (bitSizeOf Int)}

548 |        -> Nat

549 |     go Z           =

550 |       bitSizeOf Int

551 |     go (S k) {pos} =

552 |       let bit : Fin (bitSizeOf Int)

553 |           bit = ixToFin pos

554 |        in case testBit value bit of

555 |             True  =>

556 |               finToNat bit

557 |             False =>

558 |               assert_total (go k)

574 | log2 :  Nat

575 |      -> Nat

576 | log2 x =

577 |   go (bitSizeOf Int)

578 |   where

579 |     value : Int

580 |     value = cast x

581 |     ||| Scan bit positions from most significant to least significant.

587 |     go :  (remaining : Nat)

588 |        -> {auto 0 valid : LTE remaining (bitSizeOf Int)}

589 |        -> Nat

590 |     go Z =

591 |       Z

592 |     go (S k) {valid} =

593 |       let bit : Fin (bitSizeOf Int)

594 |           bit = natToFinLT k @{valid}

595 |        in case testBit value bit of

596 |             True  =>

597 |               k

598 |             False =>

599 |               assert_total (go k {valid = lteSuccLeft valid})

608 | public export

609 | data RRBVector a

610 |   = Root Nat   -- size

611 |          Shift -- shift (blockshift * height)

612 |          (Tree a)

613 |   | Empty

615 | %runElab derive "RRBVector" [Show]