1 | module Data.RRBVector.Unsized.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

 23 | bitSizeOf :  (ty : Type)

 24 |           -> FiniteBits ty

 25 |           => Nat

 26 | bitSizeOf ty = bitSize {a = ty}

 32 | public export

 33 | Shift : Type

 34 | Shift = Nat

 38 | blockshift : Shift

 39 | blockshift = 4

 43 | blocksize : Nat

 44 | blocksize = integerToNat $ 1 `shiftL` blockshift

 48 | blockmask : Nat

 49 | blockmask = minus blocksize 1

 52 | up :  Shift

 53 |    -> Shift

 54 | up sh = plus sh blockshift

 57 | down :  Shift

 58 |      -> Shift

 59 | down sh = minus sh blockshift

 62 | radixIndex :  Nat

 63 |            -> Shift

 64 |            -> Nat

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

 68 | relaxedRadixIndex :  Array Nat

 69 |                   -> Nat

 70 |                   -> Shift

 71 |                   -> (Nat, Nat)

 72 | relaxedRadixIndex sizes i sh =

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

 74 |       idx    = loop sizes guess

 75 |       subIdx = case idx == 0 of

 76 |                  True  =>

 77 |                    i

 78 |                  False =>

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

 80 |                                 Nothing    =>

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

 82 |                                 Just idx'' =>

 83 |                                   idx''

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

 85 |     in (idx, subIdx)

 86 |   where

 87 |     loop :  Array Nat

 88 |          -> Nat

 89 |          -> Nat

 90 |     loop sizes idx =

 91 |       let current = case tryNatToFin idx of

 92 |                       Nothing       =>

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

 94 |                       Just idx' =>

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

 96 |         in case i < current of

 97 |              True  =>

 98 |                idx

 99 |              False =>

100 |                assert_total $ loop sizes (plus idx 1)

107 | public export

108 | data Tree a = Balanced (Array (Tree a))

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

110 |             | Leaf (Array a)

118 | null :  Tree a

119 |      -> Bool

120 | null (Balanced arr)     =

121 |   null arr

122 | null (Unbalanced arr _) =

123 |   null arr

124 | null (Leaf arr)         =

125 |   null arr

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

133 |       -> b

134 |       -> Tree a

135 |       -> b

136 | foldl f acc tree =

137 |   foldlTree acc tree

138 |   where

139 |     foldlTree :  b

140 |               -> Tree a

141 |               -> b

142 |     foldlTree acc' (Balanced arr)     =

143 |       assert_total $ foldl foldlTree acc' arr

144 |     foldlTree acc' (Unbalanced arr _) =

145 |       assert_total $ foldl foldlTree acc' arr

146 |     foldlTree acc' (Leaf arr)         =

147 |       assert_total $ foldl f acc' arr

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

151 |       -> b

152 |       -> Tree a

153 |       -> b

154 | foldr f acc tree =

155 |   foldrTree tree acc

156 |   where

157 |     foldrTree :  Tree a

158 |               -> b

159 |               -> b

160 |     foldrTree (Balanced arr) acc'     =

161 |       assert_total $ foldr foldrTree acc' arr

162 |     foldrTree (Unbalanced arr _) acc' =

163 |       assert_total $ foldr foldrTree acc' arr

164 |     foldrTree (Leaf arr) acc'         =

165 |       assert_total $ foldr f acc' arr

172 | toList :  Tree a

173 |        -> List a

174 | toList (Balanced arr)     =

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

176 | toList (Unbalanced arr _) =

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

178 | toList (Leaf arr)         =

179 |   toList arr

185 | public export

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

187 |   show (Balanced arr)     =

188 |     assert_total $ "Balanced " ++ show arr

189 |   show (Unbalanced arr _) =

190 |     assert_total $ "Unbalanced " ++ show arr

191 |   show (Leaf arr)         =

192 |     "Leaf " ++ show arr

194 | public export

195 | Foldable Tree where

196 |   foldl f z = Data.RRBVector.Unsized.Internal.foldl f z

197 |   foldr f z = Data.RRBVector.Unsized.Internal.foldr f z

198 |   toList    = Data.RRBVector.Unsized.Internal.toList

199 |   null      = Data.RRBVector.Unsized.Internal.null

201 | public export

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

203 |   (Balanced arr1) == (Balanced arr2)         =

204 |     assert_total $ arr1 == arr2

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

206 |     assert_total $ arr1 == arr2

207 |   (Leaf arr1) == (Leaf arr2)                 =

208 |     arr1 == arr2

209 |   _                        == _              =

210 |     False

212 | public export

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

214 |   compare tree1 tree2 =

215 |     compare (Data.RRBVector.Unsized.Internal.toList tree1) (Data.RRBVector.Unsized.Internal.toList tree2)

221 | public export

222 | showTreeRep :  Show a

223 |             => Show (Tree a)

224 |             => Tree a

225 |             -> String

226 | showTreeRep (Balanced trees)     =

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

228 | showTreeRep (Unbalanced trees _) =

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

230 | showTreeRep (Leaf elems)         =

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

238 | singleton :  a

239 |           -> Array a

240 | singleton x =

241 |   A 1 $ fill 1 x

244 | treeToArray :  Tree a

245 |             -> Array (Tree a)

246 | treeToArray (Balanced arr)     =

247 |   arr

248 | treeToArray (Unbalanced arr _) =

249 |   arr

250 | treeToArray (Leaf _)           =

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

254 | treeBalanced :  Tree a

255 |              -> Bool

256 | treeBalanced (Balanced _)     =

257 |   True

258 | treeBalanced (Unbalanced _ _) =

259 |   False

260 | treeBalanced (Leaf _)         =

261 |   True

265 | treeSize :  Shift

266 |          -> Tree a

267 |          -> Nat

268 | treeSize = go 0

269 |   where

270 |     go :  Shift

271 |        -> Shift

272 |        -> Tree a

273 |        -> Nat

274 |     go acc _  (Leaf arr)             =

275 |       plus acc arr.size

276 |     go acc _  (Unbalanced arr sizes) =

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

278 |                 Nothing =>

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

280 |                 Just i' =>

281 |                   i'

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

283 |     go acc sh (Balanced arr)         =

284 |       let i  = minus arr.size 1

285 |           i' = case tryNatToFin i of

286 |                  Nothing  =>

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

288 |                  Just i'' =>

289 |                    i''

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

291 |               (down sh)

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

297 | computeSizes :  Shift

298 |              -> Array (Tree a)

299 |              -> Tree a

300 | computeSizes sh arr =

301 |   case isBalanced of

302 |     True  =>

303 |       Balanced arr

304 |     False =>

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

306 |         in Unbalanced arr arrnat

307 |   where

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

309 |          -> List (Tree a)

310 |          -> WithMArray n Nat (Array Nat)

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

312 |       res <- unsafeFreeze r

313 |       pure $ A n res

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

315 |       case tryNatToFin cur of

316 |         Nothing   =>

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

318 |         Just cur' =>

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

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

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

322 |     maxsize : Integer

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

324 |     len : Nat

325 |     len = arr.size

326 |     lenM1 : Nat

327 |     lenM1 = minus len 1

328 |     isBalanced : Bool

329 |     isBalanced = go 0

330 |       where

331 |         go :  Nat

332 |            -> Bool

333 |         go i =

334 |           let subtree = case tryNatToFin i of

335 |                           Nothing =>

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

337 |                           Just i' =>

338 |                             at arr.arr i'

339 |             in case i < lenM1 of

340 |                  True  =>

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

342 |                  False =>

343 |                    treeBalanced subtree

346 | countTrailingZeros :  Nat

347 |                    -> Nat

348 | countTrailingZeros x =

349 |   go 0

350 |   where

351 |     w : Nat

352 |     w = bitSizeOf Int

353 |     go : Nat -> Nat

354 |     go i =

355 |       case i >= w of

356 |         True  =>

357 |           i

358 |         False =>

359 |           case tryNatToFin i of

360 |             Nothing =>

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

362 |             Just i' =>

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

364 |                 True  =>

365 |                   i

366 |                 False =>

367 |                   assert_total $ go (plus i 1)

371 | log2 :  Nat

372 |      -> Nat

373 | log2 x =

374 |   let bitSizeMinus1 = minus (bitSizeOf Int) 1

375 |     in minus bitSizeMinus1 (countLeadingZeros x)

376 |   where

377 |     countLeadingZeros : Nat -> Nat

378 |     countLeadingZeros x =

379 |       minus (minus w 1) (go (minus w 1))

380 |       where

381 |         w : Nat

382 |         w = bitSizeOf Int

383 |         go : Nat -> Nat

384 |         go i =

385 |           case i < 0 of

386 |             True  =>

387 |               i

388 |             False =>

389 |               case tryNatToFin i of

390 |                 Nothing =>

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

392 |                 Just i' =>

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

394 |                     True  =>

395 |                       i

396 |                     False =>

397 |                       assert_total $ go (minus i 1)

405 | public export

406 | data RRBVector a = Root Nat   -- size

407 |                         Shift -- shift (blockshift * height)

408 |                         (Tree a)

409 |                  | Empty

411 | %runElab derive "RRBVector" [Show]