
   0 | module Data.Array.Indexed
   1 | 
   2 | import Data.Array.Mutable
   3 | import Data.List
   4 | import Data.Vect
   5 | import Data.Linear.Traverse1
   6 | import Syntax.PreorderReasoning
   7 | import Syntax.T1
   8 | 
   9 | %default total
  10 | 
  11 | ||| An immutable array paired with its size (= number of values).
  12 | |||
  13 | ||| This is the dependent pair version of `IArray size a`.
  14 | public export
  15 | record Array a where
  16 |   constructor A
  17 |   size : Nat
  18 |   arr  : IArray size a
  19 | 
  20 | --------------------------------------------------------------------------------
  21 | --          Accessing Data
  22 | --------------------------------------------------------------------------------
  23 | 
  24 | ||| Safely access a value in an immutable array at position `n - m`.
  25 | export %inline
  26 | ix : IArray n a -> (0 m : Nat) -> {auto x: Ix (S m) n} -> a
  27 | ix arr _ = at arr (ixToFin x)
  28 | 
  29 | ||| Safely access a value in an array at the given position.
  30 | export %inline
  31 | atNat : IArray n a -> (m : Nat) -> {auto 0 lt : LT m n} -> a
  32 | atNat arr x = at arr (natToFinLT x)
  33 | 
  34 | --------------------------------------------------------------------------------
  35 | --          Initializing Arrays
  36 | --------------------------------------------------------------------------------
  37 | 
  38 | ||| The empty array.
  39 | export %inline
  40 | empty : IArray 0 a
  41 | empty = unsafeAlloc 0 unsafeFreeze
  42 | 
  43 | ||| Copy the values in a list to an immutable array of the same length.
  44 | export %inline
  45 | arrayL : (ls : List a) -> IArray (length ls) a
  46 | arrayL xs = allocList xs unsafeFreeze
  47 | 
  48 | ||| Copy the values in a vector to an immutable array of the same length.
  49 | export %inline
  50 | array : {n : _} -> Vect n a -> IArray n a
  51 | array xs = allocVect xs unsafeFreeze
  52 | 
  53 | ||| Copy the values in a vector to an immutable array of the same length
  54 | ||| in reverse order.
  55 | |||
  56 | ||| This is useful if the values in the array have been collected
  57 | ||| from tail to head for instance when parsing some data.
  58 | export %inline
  59 | revArray : {n : _} -> Vect n a -> IArray n a
  60 | revArray xs = allocVectRev xs unsafeFreeze
  61 | 
  62 | ||| Fill an immutable array of the given size with the given value.
  63 | export %inline
  64 | fill : (n : Nat) -> a -> IArray n a
  65 | fill n v = alloc n v unsafeFreeze
  66 | 
  67 | ||| Generate an immutable array of the given size using
  68 | ||| the given iteration function.
  69 | export %inline
  70 | generate : (n : Nat) -> (Fin n -> a) -> IArray n a
  71 | generate n f = allocGen n f unsafeFreeze
  72 | 
  73 | ||| Generate an immutable array of the given size by filling it with the
  74 | ||| results of repeatedly applying `f` to the initial value.
  75 | export %inline
  76 | iterate : (n : Nat) -> (f : a -> a) -> a -> IArray n a
  77 | iterate n f v = allocIter n f v unsafeFreeze
  78 | 
  79 | ||| Copy the content of an array to a new immutable array.
  80 | |||
  81 | ||| This is mainly useful for reducing memory consumption, in case the
  82 | ||| original array is actually backed by a much larger array, for
  83 | ||| instance after taking a smaller prefix of a large array with `take`.
  84 | export
  85 | force : {n : _} -> IArray n a -> IArray n a
  86 | force arr = generate n (at arr)
  87 | 
  88 | ||| Allocate an immutable array, fill it with the given default value, and use a list
  89 | ||| of pairs to replace specific positions.
  90 | export
  91 | fromPairs : (n : Nat) -> a -> List (Nat,a) -> IArray n a
  92 | fromPairs n v ps = alloc n v (go ps)
  93 |   where
  94 |     go : List (Nat,a) -> WithMArray n a (IArray n a)
  95 |     go []            r = unsafeFreeze r
  96 |     go ((x,v) :: xs) r =
  97 |       case tryNatToFin x of
  98 |         Just y  => T1.do set r y v; go xs r
  99 |         Nothing => go xs r
 100 | 
 101 | ||| Like `fromPairs` but with `Fin n` as indices.
 102 | export
 103 | fromFinPairs : (n : Nat) -> a -> List (Fin n,a) -> IArray n a
 104 | fromFinPairs n v ps = alloc n v (go ps)
 105 |   where
 106 |     go : List (Fin n,a) -> WithMArray n a (IArray n a)
 107 |     go []            r = unsafeFreeze r
 108 |     go ((x,v) :: xs) r = T1.do set r x v; go xs r
 109 | 
 110 | ||| Wraps a pointer to a mutable JS array in an immutable array
 111 | ||| by copying the values first.
 112 | |||
 113 | ||| Use this for wrapping arrays that can still be mutated in JS land, or if
 114 | ||| you are uncertain about whether the pointer points to a mutable array or
 115 | ||| not. The array will first be copied before being wrapped in an `Array a`.
 116 | export
 117 | unsafeJSArrayOf1 : (0 a : Type) -> AnyPtr -> F1 s (Array a)
 118 | unsafeJSArrayOf1 a ptr t =
 119 |  let (n ** marr) # t := unsafeJSMArrayOf1 a ptr t
 120 |      iarr        # t := freeze marr t
 121 |   in  A n iarr # t
 122 | 
 123 | ||| Like `unsafeJSArrayOf1` but runs in an `IO` monad.
 124 | export
 125 | unsafeJSArrayOf : HasIO io => (0 a : Type) -> AnyPtr -> io (Array a)
 126 | unsafeJSArrayOf a ptr = runIO (unsafeJSArrayOf1 a ptr)
 127 | 
 128 | ||| Like `unsafeJSArrayOf1` but wraps the pointer without copying
 129 | ||| it first.
 130 | |||
 131 | ||| Use this only if the pointer points to an array that cannot be
 132 | ||| inadvertently mutated later on.
 133 | export
 134 | unsafeWrapJSArray : (0 a : Type) -> AnyPtr -> Array a
 135 | unsafeWrapJSArray a ptr =
 136 |   run1 $ \t =>
 137 |    let (n ** marr) # t := unsafeJSMArrayOf1 a ptr t
 138 |        iarr        # t := unsafeFreeze marr t
 139 |     in  A n iarr # t
 140 | 
 141 | --------------------------------------------------------------------------------
 142 | --          Eq and Ord
 143 | --------------------------------------------------------------------------------
 144 | 
 145 | ||| Lexicographic comparison of immutable arrays of distinct length.
 146 | export
 147 | hcomp : {m,n : Nat} -> Ord a => IArray m a -> IArray n a -> Ordering
 148 | hcomp a1 a2 = go m n
 149 | 
 150 |   where
 151 |     go : (k,l : Nat) -> {auto _ : Ix k m} -> {auto _ : Ix l n} -> Ordering
 152 |     go 0     0     = EQ
 153 |     go 0     (S _) = LT
 154 |     go (S _) 0     = GT
 155 |     go (S k) (S j) = case compare (ix a1 k) (ix a2 j) of
 156 |       EQ => go k j
 157 |       r  => r
 158 | 
 159 | ||| Heterogeneous equality for immutable arrays.
 160 | export
 161 | heq : {m,n : Nat} -> Eq a => IArray m a -> IArray n a -> Bool
 162 | heq a1 a2 = go m n
 163 | 
 164 |   where
 165 |     go : (k,l : Nat) -> {auto _ : Ix k m} -> {auto _ : Ix l n} -> Bool
 166 |     go 0     0     = True
 167 |     go (S k) (S j) = if ix a1 k == ix a2 j then go k j else False
 168 |     go _     _     = False
 169 | 
 170 | export
 171 | {n : Nat} -> Eq a => Eq (IArray n a) where
 172 |   a1 == a2 = go n
 173 | 
 174 |     where
 175 |       go : (k : Nat) -> {auto 0 _ : LTE k n} -> Bool
 176 |       go 0     = True
 177 |       go (S k) = if atNat a1 k == atNat a2 k then go k else False
 178 | 
 179 | export
 180 | {n : Nat} -> Ord a => Ord (IArray n a) where
 181 |   compare a1 a2 = go n
 182 | 
 183 |     where
 184 |       go : (k : Nat) -> {auto _ : Ix k n} -> Ordering
 185 |       go 0     = EQ
 186 |       go (S k) = case compare (ix a1 k) (ix a2 k) of
 187 |         EQ => go k
 188 |         c  => c
 189 | 
 190 | --------------------------------------------------------------------------------
 191 | --          Maps and Folds
 192 | --------------------------------------------------------------------------------
 193 | 
 194 | ontoList : List a -> (m : Nat) -> (0 lte : LTE m n) => IArray n a -> List a
 195 | ontoList xs 0     arr = xs
 196 | ontoList xs (S k) arr = ontoList (atNat arr k :: xs) k arr
 197 | 
 198 | ontoVect :
 199 |      Vect k a
 200 |   -> (m : Nat)
 201 |   -> {auto 0 lte : LTE m n}
 202 |   -> IArray n a
 203 |   -> Vect (k + m) a
 204 | ontoVect xs 0     arr = rewrite plusZeroRightNeutral k in xs
 205 | ontoVect xs (S v) arr =
 206 |   rewrite sym (plusSuccRightSucc k v) in ontoVect (atNat arr v :: xs) v arr
 207 | 
 208 | ontoVectWithIndex :
 209 |      Vect k (Fin n, a)
 210 |   -> (m : Nat)
 211 |   -> {auto 0 lte : LTE m n}
 212 |   -> IArray n a
 213 |   -> Vect (k + m) (Fin n, a)
 214 | ontoVectWithIndex xs 0     arr = rewrite plusZeroRightNeutral k in xs
 215 | ontoVectWithIndex xs (S v) arr =
 216 |   rewrite sym (plusSuccRightSucc k v)
 217 |   in let x := natToFinLT v in ontoVectWithIndex ((x, at arr x) :: xs) v arr
 218 | 
 219 | ||| Convert an immutable array to a vector of the same length.
 220 | export %inline
 221 | toVect : {n : _} -> IArray n a -> Vect n a
 222 | toVect = ontoVect [] n
 223 | 
 224 | ||| Convert an immutable array to a vector of the same length
 225 | ||| pairing all values with their index.
 226 | export %inline
 227 | toVectWithIndex : {n : _} -> IArray n a -> Vect n (Fin n, a)
 228 | toVectWithIndex = ontoVectWithIndex [] n
 229 | 
 230 | foldrI : (m : Nat) -> (0 _ : LTE m n) => (e -> a -> a) -> a -> IArray n e -> a
 231 | foldrI 0     _ x arr = x
 232 | foldrI (S k) f x arr = foldrI k f (f (atNat arr k) x) arr
 233 | 
 234 | foldrKV_ :
 235 |      (m : Nat)
 236 |   -> {auto 0 prf : LTE m n}
 237 |   -> (Fin n -> e -> a -> a)
 238 |   -> a
 239 |   -> IArray n e -> a
 240 | foldrKV_ 0     _ x arr = x
 241 | foldrKV_ (S k) f x arr =
 242 |   let fin := natToFinLT k @{prf} in foldrKV_ k f (f fin (at arr fin) x) arr
 243 | 
 244 | foldlI : (m : Nat) -> (x : Ix m n) => (a -> e -> a) -> a -> IArray n e -> a
 245 | foldlI 0     _ v arr = v
 246 | foldlI (S k) f v arr = foldlI k f (f v (ix arr k)) arr
 247 | 
 248 | foldlKV_ :
 249 |      (m : Nat)
 250 |   -> {auto x : Ix m n}
 251 |   -> (Fin n -> a -> e -> a)
 252 |   -> a
 253 |   -> IArray n e
 254 |   -> a
 255 | foldlKV_ 0     _ v arr = v
 256 | foldlKV_ (S k) f v arr =
 257 |   let fin := ixToFin x in foldlKV_ k f (f fin v (at arr fin)) arr
 258 | 
 259 | ||| Right fold over the values of an immutable array plus their indices.
 260 | export %inline
 261 | foldrKV : {n : _} -> (Fin n -> e -> a -> a) -> a -> IArray n e -> a
 262 | foldrKV = foldrKV_ n
 263 | 
 264 | ||| Left fold over the values of an immutable array plus their indices.
 265 | export %inline
 266 | foldlKV : {n : _} -> (Fin n -> a -> e -> a) -> a -> IArray n e -> a
 267 | foldlKV = foldlKV_ n
 268 | 
 269 | export %inline
 270 | {n : Nat} -> Foldable (IArray n) where
 271 |   foldr = foldrI n
 272 |   foldl = foldlI n
 273 |   toList = ontoList [] n
 274 |   null _ = n == Z
 275 | 
 276 | export %inline
 277 | {n : Nat} -> Functor (IArray n) where
 278 |   map f arr = generate n (f . at arr)
 279 | 
 280 | export
 281 | {n : Nat} -> Applicative (IArray n) where
 282 |   pure = fill n
 283 |   af <*> av = generate n (\x => at af x (at av x))
 284 | 
 285 | export
 286 | {n : Nat} -> Monad (IArray n) where
 287 |   arr >>= f = generate n (\x => at (f $ at arr x) x)
 288 | 
 289 | export
 290 | {n : Nat} -> Show a => Show (IArray n a) where
 291 |   showPrec p arr = showCon p "array" (showArg $ ontoList [] n arr)
 292 | 
 293 | ||| Mapping over the values of an immutable array together with their indices.
 294 | export
 295 | mapWithIndex : {n : _} -> (Fin n -> a -> b) -> IArray n a -> IArray n b
 296 | mapWithIndex f arr = generate n (\x => f x (at arr x))
 297 | 
 298 | ||| Update a single position in an immutable array by applying the given
 299 | ||| function.
 300 | |||
 301 | ||| This will have to copy the whole array, so it runs in O(n).
 302 | export
 303 | updateAt : {n : _} -> Fin n -> (a -> a) -> IArray n a -> IArray n a
 304 | updateAt x f = mapWithIndex (\k,v => if x == k then f v else v)
 305 | 
 306 | ||| Set a single position in an immutable array.
 307 | |||
 308 | ||| This will have to copy the whole array, so it runs in O(n).
 309 | export
 310 | setAt : {n : _} -> Fin n -> a -> IArray n a -> IArray n a
 311 | setAt x y = mapWithIndex (\k,v => if x == k then y else v)
 312 | 
 313 | --------------------------------------------------------------------------------
 314 | --          Traversals
 315 | --------------------------------------------------------------------------------
 316 | 
 317 | ||| Effectful traversal of the values in an immutable array together with
 318 | ||| their corresponding indices.
 319 | export
 320 | traverseWithIndex :
 321 |      {n : _}
 322 |   -> {auto app : Applicative f}
 323 |   -> (Fin n -> a -> f b)
 324 |   -> IArray n a
 325 |   -> f (IArray n b)
 326 | traverseWithIndex f arr =
 327 |   array <$> traverse (\(x,v) => f x v) (toVectWithIndex arr)
 328 | 
 329 | export
 330 | {n : _} -> Traversable (IArray n) where
 331 |   traverse = traverseWithIndex . const
 332 | 
 333 | --------------------------------------------------------------------------------
 334 | --          Sub-Arrays
 335 | --------------------------------------------------------------------------------
 336 | 
 337 | 0 curLTE : (s : Ix m n) -> LTE c (ixToNat s) -> LTE c n
 338 | curLTE s lte = transitive lte $ ixLTE s
 339 | 
 340 | 0 curLT : (s : Ix (S m) n) -> LTE c (ixToNat s) -> LT c n
 341 | curLT s lte = let LTESucc p := ixLT s in LTESucc $ transitive lte p
 342 | 
 343 | ||| Drop n elements from a immutable array. O(n)
 344 | export
 345 | drop : {n : _} -> (m : Nat) -> IArray n a -> IArray (n `minus` m) a
 346 | drop m arr = generate (n `minus` m) (\f => at arr (inc f))
 347 | 
 348 | ||| Filter the values in an immutable array together with their corresponding
 349 | ||| indices according to the given predicate.
 350 | export
 351 | filterWithKey :
 352 |      {n : Nat}
 353 |   -> (Fin n -> a -> Bool)
 354 |   -> IArray n a
 355 |   -> Array a
 356 | filterWithKey f arr = run1 $ T1.do
 357 |   (n ** m) <- mfilterWithKey (unsafeThaw arr) f
 358 |   res      <- unsafeFreeze m
 359 |   pure (A n res)
 360 | 
 361 | ||| Filters the values in an immutable array according to the given predicate.
 362 | export %inline
 363 | filter : {n : Nat} -> (a -> Bool) -> IArray n a -> Array a
 364 | filter = filterWithKey . const
 365 | 
 366 | ||| Map the values in an immutable array together with their corresponding indices
 367 | ||| over a function that might not return a result for all values.
 368 | export
 369 | mapMaybeWithKey :
 370 |      {n : Nat}
 371 |   -> (Fin n -> a -> Maybe b)
 372 |   -> IArray n a
 373 |   -> Array b
 374 | mapMaybeWithKey f arr = run1 $ T1.do
 375 |   (n ** m) <- mmapMaybeWithKey (unsafeThaw arr) f
 376 |   res      <- unsafeFreeze m
 377 |   pure (A n res)
 378 | 
 379 | ||| Map the values in an immutable array together with their corresponding indices
 380 | ||| over a function that might not return a result for all values.
 381 | export %inline
 382 | mapMaybe : {n : Nat} -> (a -> Maybe b) -> IArray n a -> Array b
 383 | mapMaybe = mapMaybeWithKey . const
 384 | 
 385 | --------------------------------------------------------------------------------
 386 | --          Concatenating Arrays
 387 | --------------------------------------------------------------------------------
 388 | 
 389 | ||| Size of an array after concatenating a SnocList of arrays.
 390 | |||
 391 | ||| It is easier to implement this and keep the indices correct,
 392 | ||| therefore, this is the default for concatenating arrays.
 393 | public export
 394 | SnocSize : SnocList (Array a) -> Nat
 395 | SnocSize [<]           = 0
 396 | SnocSize (xs :< A s _) = SnocSize xs + s
 397 | 
 398 | ||| Size of an immutable array after concatenating a List of arrays.
 399 | public export
 400 | ListSize : List (Array a) -> Nat
 401 | ListSize = SnocSize . ([<] <><)
 402 | 
 403 | -- snocConcat implementation
 404 | sconc :
 405 |      (pos         : Nat)
 406 |   -> (cur         : Nat)
 407 |   -> (x           : IArray m a)
 408 |   -> (arrs        : SnocList (Array a))
 409 |   -> {auto 0 lte1 : LTE pos n}
 410 |   -> {auto 0 lte2 : LTE cur m}
 411 |   -> WithMArray n a (IArray n a)
 412 | sconc pos   0     _   (sx :< A s x) r t = sconc pos s x   sx r t
 413 | sconc (S j) (S k) x   sx            r t =
 414 |   let _ # t := setNat r j (atNat x k) t
 415 |    in sconc j k x sx r t
 416 | sconc _     _     _   _             r t = unsafeFreeze r t
 417 | 
 418 | ||| Concatenate a SnocList of arrays.
 419 | |||
 420 | ||| This allocates a large enough array in advance, and therefore runs in
 421 | ||| O(SnocSize sa).
 422 | export
 423 | snocConcat : (sa : SnocList (Array a)) -> IArray (SnocSize sa) a
 424 | snocConcat [<]                 = empty
 425 | snocConcat (sa :< A 0 _)       =
 426 |   rewrite plusZeroRightNeutral (SnocSize sa) in snocConcat sa
 427 | snocConcat (sa :< A (S k) arr) with (SnocSize sa + S k)
 428 |   _ | n = alloc n (at arr 0) (sconc n (S k) arr sa)
 429 | 
 430 | ||| Concatenate a List of arrays.
 431 | |||
 432 | ||| This allocates a large enough array in advance, and therefore runs in
 433 | ||| O(ListSize as).
 434 | export
 435 | listConcat : (as : List (Array a)) -> IArray (ListSize as) a
 436 | listConcat as = snocConcat ([<] <>< as)
 437 | 
 438 | ||| Concatenate two immutable arrays in O(m+n) runtime.
 439 | export
 440 | append : {m,n : Nat} -> IArray m a -> IArray n a -> IArray (m + n) a
 441 | append src1 src2 =
 442 |   unsafeAlloc (m+n) $ \r,t =>
 443 |     let _ # t := icopy {n = m+n} src1 0 0 m @{reflexive} @{lteAddRight _} r t
 444 |         _ # t := icopy src2 0 m n @{reflexive} @{reflexive} r t
 445 |      in unsafeFreeze r t
 446 | 
 447 | --------------------------------------------------------------------------------
 448 | --          Linear Folds and Traversals
 449 | --------------------------------------------------------------------------------
 450 | 
 451 | parameters (arr : IArray n e)
 452 | 
 453 |   foldl1A : (k : Nat) -> Ix k n => (a -> e -> F1 s a) -> a -> F1 s a
 454 |   foldl1A 0     f x t = x # t
 455 |   foldl1A (S k) f x t = let y # t := f x (arr `ix` k) t in foldl1A k f y t
 456 | 
 457 |   foldr1A : (k : Nat) -> (0 p : LTE k n) => (e -> a -> F1 s a) -> a -> F1 s a
 458 |   foldr1A 0     f x t = x # t
 459 |   foldr1A (S k) f x t = let y # t := f (arr `atNat` k) x t in foldr1A k f y t
 460 | 
 461 |   fm1A : Semigroup a => (k : Nat) -> (0 p : LTE k n) => (e -> F1 s a) -> a -> F1 s a
 462 |   fm1A 0     f x t = x # t
 463 |   fm1A (S k) f x t = let y # t := f (arr `atNat` k) t in fm1A k f (y<+>x) t
 464 | 
 465 |   tr1A_ : (k : Nat) -> Ix k n => (e -> F1' s) -> F1' s
 466 |   tr1A_ 0     f t = () # t
 467 |   tr1A_ (S k) f t = let _ # t := f (arr `ix` k) t in tr1A_ k f t
 468 | 
 469 |   tr1A : (k : Nat) -> Ix k n => (e -> F1 s b) -> MArray s n b -> F1 s (IArray n b)
 470 |   tr1A 0     f m t = unsafeFreeze m t
 471 |   tr1A (S k) f m t =
 472 |    let y # t := f (arr `ix` k) t
 473 |        _ # t := setIx m k y t
 474 |     in tr1A k f m t
 475 | 
 476 | ||| Runs a linear effect over the values plus their indices in an array.
 477 | export
 478 | traverseKV1_ : {n : _} -> (Fin n -> a -> F1' q) -> IArray n a -> F1' q
 479 | traverseKV1_ f arr = go n
 480 |   where
 481 |     go : (k : Nat) -> (x : Ix k n) => F1' q
 482 |     go 0     t = () # t
 483 |     go (S k) t = let _ # t := f (ixToFin x) (arr `ix` k) t in go k t
 484 | 
 485 | export %inline
 486 | {n : _} -> Foldable1 (IArray n) where
 487 |   foldl1     f x arr = foldl1A arr n f x
 488 |   foldr1     f x arr = foldr1A arr n f x
 489 |   foldMap1   f   arr = fm1A arr n f neutral
 490 |   traverse1_ f   arr = tr1A_ arr n f
 491 | 
 492 | export %inline
 493 | {n : _} -> Traversable1 (IArray n) where
 494 |   traverse1 f arr = unsafeMArray1 n >>= tr1A arr n f
 495 | 
 496 | export %inline
 497 | Foldable1 Array where
 498 |   foldl1     f x (A n arr) = foldl1A arr n f x
 499 |   foldr1     f x (A n arr) = foldr1A arr n f x
 500 |   foldMap1   f   (A n arr) = fm1A arr n f neutral
 501 |   traverse1_ f   (A n arr) = tr1A_ arr n f
 502 | 
 503 | export %inline
 504 | Traversable1 Array where
 505 |   traverse1 f (A n arr) t = let a # t := traverse1 f arr t in A n a # t
 506 |