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.List

  9 | import Data.Nat

 10 | import Data.String

 11 | import Derive.Prelude

 12 | import Syntax.T1 as T1

 14 | %default total

 15 | %language ElabReflection

 24 | bitSizeOf :  (ty : Type)

 25 |           -> FiniteBits ty

 26 |           => Nat

 27 | bitSizeOf ty = bitSize {a = ty}

 33 | public export

 34 | Shift : Type

 35 | Shift = Nat

 40 | blockshift : Shift

 41 | blockshift = 4

 46 | blocksize : Nat

 47 | blocksize = integerToNat $ 1 `shiftL` blockshift

 52 | blockmask : Nat

 53 | blockmask = minus blocksize 1

 56 | up :  Shift

 57 |    -> Shift

 58 | up sh = plus sh blockshift

 61 | down :  Shift

 62 |      -> Shift

 63 | down sh = minus sh blockshift

 66 | radixIndex :  Nat

 67 |            -> Shift

 68 |            -> Nat

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

 72 | relaxedRadixIndex :  Array Nat

 73 |                   -> Nat

 74 |                   -> Shift

 75 |                   -> (Nat, Nat)

 76 | relaxedRadixIndex sizes i sh =

 77 |   let guess  = radixIndex i sh -- guess <= idx

 78 |       idx    = loop sizes guess

 79 |       subIdx = case idx == 0 of

 80 |                  True  =>

 81 |                    i

 82 |                  False =>

 83 |                    let idx' = case tryNatToFin $ minus idx 1 of

 84 |                                 Nothing    =>

 85 |                                   assert_total $ idris_crash "Data.RRBVector.Internal.relaxedRadixIndex: index out of bounds"

 86 |                                 Just idx'' =>

 87 |                                   idx''

 88 |                      in minus i (at sizes.arr idx')

 89 |     in (idx, subIdx)

 90 |   where

 91 |     loop :  Array Nat

 92 |          -> Nat

 93 |          -> Nat

 94 |     loop sizes idx =

 95 |       let current = case tryNatToFin idx of

 96 |                       Nothing       =>

 97 |                         assert_total $ idris_crash "Data.RRBVector.Internal.relaxedRadixIndex.loop: index out of bounds"

 98 |                       Just idx' =>

 99 |                         at sizes.arr idx' -- idx will always be in range for a well-formed tree

100 |         in case i < current of

101 |              True  =>

102 |                idx

103 |              False =>

104 |                assert_total $ loop sizes (plus idx 1)

112 | public export

113 | data Tree a

114 |   = Balanced (Array (Tree a))

115 |   | Unbalanced (Array (Tree a)) (Array Nat)

116 |   | Leaf (Array a)

125 | null :  Tree a

126 |      -> Bool

127 | null (Balanced arr)     =

128 |   null arr

129 | null (Unbalanced arr _) =

130 |   null arr

131 | null (Leaf arr)         =

132 |   null arr

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

140 |       -> b

141 |       -> Tree a

142 |       -> b

143 | foldl f acc tree =

144 |   foldlTree acc tree

145 |   where

146 |     foldlTree :  b

147 |               -> Tree a

148 |               -> b

149 |     foldlTree acc' (Balanced arr)     =

150 |       assert_total $ foldl foldlTree acc' arr

151 |     foldlTree acc' (Unbalanced arr _) =

152 |       assert_total $ foldl foldlTree acc' arr

153 |     foldlTree acc' (Leaf arr)         =

154 |       assert_total $ foldl f acc' arr

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

158 |       -> b

159 |       -> Tree a

160 |       -> b

161 | foldr f acc tree =

162 |   foldrTree tree acc

163 |   where

164 |     foldrTree :  Tree a

165 |               -> b

166 |               -> b

167 |     foldrTree (Balanced arr) acc'     =

168 |       assert_total $ foldr foldrTree acc' arr

169 |     foldrTree (Unbalanced arr _) acc' =

170 |       assert_total $ foldr foldrTree acc' arr

171 |     foldrTree (Leaf arr) acc'         =

172 |       assert_total $ foldr f acc' arr

179 | toList :  Tree a

180 |        -> List a

181 | toList (Balanced arr)     =

182 |   assert_total $ concat (map toList (toList arr))

183 | toList (Unbalanced arr _) =

184 |   assert_total $ concat (map toList (toList arr))

185 | toList (Leaf arr)         =

186 |   toList arr

192 | public export

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

194 |   show (Balanced arr)     =

195 |     assert_total $ "Balanced " ++ show arr

196 |   show (Unbalanced arr _) =

197 |     assert_total $ "Unbalanced " ++ show arr

198 |   show (Leaf arr)         =

199 |     "Leaf " ++ show arr

201 | public export

202 | Foldable Tree where

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

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

205 |   toList    = Data.RRBVector.Internal.toList

206 |   null      = Data.RRBVector.Internal.null

208 | public export

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

210 |   (Balanced arr1) == (Balanced arr2)         =

211 |     assert_total $ arr1 == arr2

212 |   (Unbalanced arr1 _) == (Unbalanced arr2 _) =

213 |     assert_total $ arr1 == arr2

214 |   (Leaf arr1) == (Leaf arr2)                 =

215 |     arr1 == arr2

216 |   _                        == _              =

217 |     False

219 | public export

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

221 |   compare tree1 tree2 =

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

228 | public export

229 | showTreeRep :  Show a

230 |             => Show (Tree a)

231 |             => Tree a

232 |             -> String

233 | showTreeRep (Balanced trees)     =

234 |   assert_total $ "Balanced " ++ (show $ toList trees)

235 | showTreeRep (Unbalanced trees _) =

236 |   assert_total $ "Unbalanced " ++ (show $ toList trees)

237 | showTreeRep (Leaf elems)         =

238 |   assert_total $ "Leaf " ++ (show $ toList elems)

245 | singleton :  a

246 |           -> Array a

247 | singleton x =

248 |   A 1 $ fill 1 x

251 | treeToArray :  Tree a

252 |             -> Array (Tree a)

253 | treeToArray (Balanced arr)     =

254 |   arr

255 | treeToArray (Unbalanced arr _) =

256 |   arr

257 | treeToArray (Leaf _)           =

258 |   assert_total $ idris_crash "Data.RRBVector.Internal.treeToArray: leaf"

261 | treeBalanced :  Tree a

262 |              -> Bool

263 | treeBalanced (Balanced _)     =

264 |   True

265 | treeBalanced (Unbalanced _ _) =

266 |   False

267 | treeBalanced (Leaf _)         =

268 |   True

273 | treeSize :  Shift

274 |          -> Tree a

275 |          -> Nat

276 | treeSize = go 0

277 |   where

278 |     go :  Shift

279 |        -> Shift

280 |        -> Tree a

281 |        -> Nat

282 |     go acc _  (Leaf arr)             =

283 |       plus acc arr.size

284 |     go acc _  (Unbalanced arr sizes) =

285 |       let i = case tryNatToFin $ minus arr.size 1 of

286 |                 Nothing =>

287 |                   assert_total $ idris_crash "Data.RRBVector.Internal.treeSize: index out of bounds"

288 |                 Just i' =>

289 |                   i'

290 |         in plus acc (at sizes.arr i)

291 |     go acc sh (Balanced arr)         =

292 |       let i  = minus arr.size 1

293 |           i' = case tryNatToFin i of

294 |                  Nothing  =>

295 |                    assert_total $ idris_crash "Data.RRBVector.Internal.treeSize: index out of bounds"

296 |                  Just i'' =>

297 |                    i''

298 |         in go (plus acc (mult i (integerToNat (1 `shiftL` sh))))

299 |               (down sh)

300 |               (assert_smaller arr (at arr.arr i'))

306 | computeSizes :  Shift

307 |              -> Array (Tree a)

308 |              -> Tree a

309 | computeSizes sh arr =

310 |   case isBalanced of

311 |     True  =>

312 |       Balanced arr

313 |     False =>

314 |       let arrnat = unsafeAlloc arr.size (loop sh 0 0 arr.size (toList arr))

315 |         in Unbalanced arr arrnat

316 |   where

317 |     loop :  (sh,cur,acc,n : Nat)

318 |          -> List (Tree a)

319 |          -> WithMArray n Nat (Array Nat)

320 |     loop sh _   acc n []        r = T1.do

321 |       res <- unsafeFreeze r

322 |       pure $ A n res

323 |     loop sh cur acc n (x :: xs) r =

324 |       case tryNatToFin cur of

325 |         Nothing   =>

326 |           assert_total $ idris_crash "Data.RRBVector.Internal.computeSizes.go: can't convert Nat to Fin"

327 |         Just cur' =>

328 |           let acc' = plus acc (treeSize (down sh) x)

329 |             in T1.do set r cur' acc'

330 |                      assert_total $ loop sh (S cur) acc' n xs r

331 |     maxsize : Integer

332 |     maxsize = 1 `shiftL` sh -- the maximum size of a subtree

333 |     len : Nat

334 |     len = arr.size

335 |     lenM1 : Nat

336 |     lenM1 = minus len 1

337 |     isBalanced : Bool

338 |     isBalanced = go 0

339 |       where

340 |         go :  Nat

341 |            -> Bool

342 |         go i =

343 |           let subtree = case tryNatToFin i of

344 |                           Nothing =>

345 |                             assert_total $ idris_crash "Data.RRBVector.Internal.computeSizes.isBalanced: can't convert Nat to Fin"

346 |                           Just i' =>

347 |                             at arr.arr i'

348 |             in case i < lenM1 of

349 |                  True  =>

350 |                    assert_total $ (natToInteger $ treeSize (down sh) subtree) == maxsize && go (plus i 1)

351 |                  False =>

352 |                    treeBalanced subtree

355 | countTrailingZeros :  Nat

356 |                    -> Nat

357 | countTrailingZeros x =

358 |   go 0

359 |   where

360 |     w : Nat

361 |     w = bitSizeOf Int

362 |     go : Nat -> Nat

363 |     go i =

364 |       case i >= w of

365 |         True  =>

366 |           i

367 |         False =>

368 |           case tryNatToFin i of

369 |             Nothing =>

370 |               assert_total $ idris_crash "Data.RRBVector.Internal.countTrailingZeros: can't convert Nat to Fin"

371 |             Just i' =>

372 |               case testBit (the Int (cast x)) i' of

373 |                 True  =>

374 |                   i

375 |                 False =>

376 |                   assert_total $ go (plus i 1)

381 | log2 :  Nat

382 |      -> Nat

383 | log2 x =

384 |   let bitSizeMinus1 = minus (bitSizeOf Int) 1

385 |     in minus bitSizeMinus1 (countLeadingZeros x)

386 |   where

387 |     countLeadingZeros : Nat -> Nat

388 |     countLeadingZeros x =

389 |       minus (minus w 1) (go (minus w 1))

390 |       where

391 |         w : Nat

392 |         w = bitSizeOf Int

393 |         go : Nat -> Nat

394 |         go i =

395 |           case i < 0 of

396 |             True  =>

397 |               i

398 |             False =>

399 |               case tryNatToFin i of

400 |                 Nothing =>

401 |                   assert_total $ idris_crash "Data.RRBVector.Internal.log2: can't convert Nat to Fin"

402 |                 Just i' =>

403 |                   case testBit (the Int (cast x)) i' of

404 |                     True  =>

405 |                       i

406 |                     False =>

407 |                       assert_total $ go (minus i 1)

416 | public export

417 | data RRBVector a

418 |   = Root Nat   -- size

419 |          Shift -- shift (blockshift * height)

420 |          (Tree a)

421 |   | Empty

423 | %runElab derive "RRBVector" [Show]