
   0 | module Data.Array.Mutable
   1 | 
   2 | import public Data.Linear.Token
   3 | import public Data.Array.Core
   4 | import public Data.Array.Index
   5 | 
   6 | import Data.List
   7 | import Data.Vect
   8 | 
   9 | %default total
  10 | 
  11 | --------------------------------------------------------------------------------
  12 | --          Reading and writing mutable arrays
  13 | --------------------------------------------------------------------------------
  14 | 
  15 | ||| Set a value in a mutable array corresponding to a position in a list
  16 | ||| used for filling said array.
  17 | export %inline
  18 | setAtSuffix :
  19 |      (r : MArray s (length ys) a)
  20 |   -> Suffix (x::xs) ys
  21 |   -> a
  22 |   -> F1' s
  23 | setAtSuffix r v = set r (suffixToFin v)
  24 | 
  25 | parameters (r : MArray s n a)
  26 | 
  27 |   ||| Set a value at index `n - m` in a mutable array.
  28 |   |||
  29 |   ||| This actually uses the auto implicit argument for accessing the
  30 |   ||| the array. It is mainly useful for iterating over an array from the left
  31 |   ||| by using a natural number for counting down (see also the documentation
  32 |   ||| for `Ix`).
  33 |   export %inline
  34 |   setIx : (0 m : Nat) -> (x : Ix (S m) n) => a -> F1' s
  35 |   setIx _ = set r (ixToFin x)
  36 | 
  37 |   ||| Set a value at index `m` in a mutable array.
  38 |   export %inline
  39 |   setNat : (m : Nat) -> (0 lt : LT m n) => a -> F1' s
  40 |   setNat x = set r (natToFinLT x)
  41 | 
  42 |   ||| Read a value at index `n - m` from a mutable array.
  43 |   |||
  44 |   ||| This actually uses the auto implicit argument for accessing the
  45 |   ||| the array. It is mainly useful for iterating over an array from the left
  46 |   ||| by using a natural number for counting down (see also the documentation
  47 |   ||| for `Ix`).
  48 |   export %inline
  49 |   getIx : (0 m : Nat) -> (x : Ix (S m) n) => F1 s a
  50 |   getIx _ = get r (ixToFin x)
  51 | 
  52 |   ||| Read a value at index `m` from a mutable array.
  53 |   export %inline
  54 |   getNat : (m : Nat) -> (0 lt : LT m n) => F1 s a
  55 |   getNat x = get r (natToFinLT x)
  56 | 
  57 |   ||| Modify a value at index `n - m` in a mutable array.
  58 |   |||
  59 |   ||| This actually uses the auto implicit argument for accessing the
  60 |   ||| the array. It is mainly useful for iterating over an array from the left
  61 |   ||| by using a natural number for counting down (see also the documentation
  62 |   ||| for `Ix`).
  63 |   export %inline
  64 |   modifyIx : (0 m : Nat) -> (x : Ix (S m) n) => (a -> a) -> F1' s
  65 |   modifyIx _ = modify r (ixToFin x)
  66 | 
  67 |   ||| Modify a value at index `m` in a mutable array.
  68 |   export %inline
  69 |   modifyNat : (m : Nat) -> (0 lt : LT m n) => (a -> a) -> F1' s
  70 |   modifyNat m = modify r (natToFinLT m)
  71 | 
  72 | --------------------------------------------------------------------------------
  73 | --          Filling Arrays
  74 | --------------------------------------------------------------------------------
  75 | 
  76 | ||| Writes the data from a list to a mutable array.
  77 | |||
  78 | ||| This uses the `Suffix` proof for safely indexing into the array.
  79 | export
  80 | writeList :
  81 |      (r  : MArray s (length xs) a)
  82 |   -> (ys : List a)
  83 |   -> {auto p : Suffix ys xs}
  84 |   -> F1' s
  85 | writeList r []        t = () # t
  86 | writeList r (y :: ys) t =
  87 |  let _ # t := setAtSuffix r p y t
  88 |   in writeList {xs} r ys t
  89 | 
  90 | ||| Writes the data from a list to a mutable array.
  91 | |||
  92 | ||| This uses the `Suffix` proof for safely indexing into the array.
  93 | export
  94 | writeListWith :
  95 |      (r  : MArray s (length xs) b)
  96 |   -> (ys : List a)
  97 |   -> (f : a -> b)
  98 |   -> {auto p : Suffix ys xs}
  99 |   -> F1' s
 100 | writeListWith r []        f t = () # t
 101 | writeListWith r (y :: ys) f t =
 102 |  let _ # t := setAtSuffix r p (f y) t
 103 |   in writeListWith {xs} r ys f t
 104 | 
 105 | parameters (r : MArray s n a)
 106 | 
 107 |   ||| Writes the data from a vector to a mutable array.
 108 |   export
 109 |   writeVect : Vect k a -> Ix k n => F1' s
 110 |   writeVect           []        t = () # t
 111 |   writeVect {k = S m} (y :: ys) t =
 112 |    let _ # t := setIx r m y t
 113 |     in writeVect ys t
 114 | 
 115 |   ||| Writes the data from a vector to a mutable array in reverse order.
 116 |   export
 117 |   writeVectRev : (m : Nat) -> Vect k a -> (0 _ : LTE m n) => F1' s
 118 |   writeVectRev (S l) (y :: ys) t =
 119 |    let _ # t := setNat r l y t
 120 |     in writeVectRev l ys t
 121 |   writeVectRev _     _         t = () # t
 122 | 
 123 |   ||| Overwrite the values in a mutable array from the
 124 |   ||| given index downward with the result of the given function.
 125 |   export
 126 |   genFrom : (m : Nat) -> (0 _ : LTE m n) => (Fin n -> a) -> F1' s
 127 |   genFrom 0     f t = () # t
 128 |   genFrom (S k) f t =
 129 |     let _ # t := setNat r k (f $ natToFinLT k) t
 130 |      in genFrom k f t
 131 | 
 132 |   ||| Overwrite the values in a mutable array from the
 133 |   ||| given index downward with the result of the given function.
 134 |   export
 135 |   genFrom' : (m : Nat) -> (0 _ : LTE m n) => (Fin n -> F1 s a) -> F1' s
 136 |   genFrom' 0     f t = () # t
 137 |   genFrom' (S k) f t =
 138 |     let f' # t := f (natToFinLT k) t
 139 |         _  # t := setNat r k f' t
 140 |       in genFrom' k f t
 141 | 
 142 |   ||| Overwrite the values in a mutable array from the
 143 |   ||| given index upward with the results of applying the given
 144 |   ||| function repeatedly.
 145 |   export
 146 |   iterateFrom : (m : Nat) -> (ix : Ix m n) => (a -> a) -> a -> F1' s
 147 |   iterateFrom 0     f v t = () # t
 148 |   iterateFrom (S k) f v t =
 149 |    let _ # t := setIx r k v t
 150 |     in iterateFrom k f (f v) t
 151 | 
 152 | export
 153 | allocList : (xs : List a) -> WithMArray (length xs) a b -> b
 154 | allocList xs g =
 155 |   unsafeAlloc (length xs) $ \r,t =>
 156 |    let _ # t := writeList {xs} r xs t
 157 |     in g r t
 158 | 
 159 | export
 160 | allocListWith : (xs : List a) -> (a -> b) -> WithMArray (length xs) b c -> c
 161 | allocListWith xs f g =
 162 |   unsafeAlloc (length xs) $ \r,t =>
 163 |    let _ # t := writeListWith {xs} r xs f t
 164 |     in g r t
 165 | 
 166 | export
 167 | allocVect : {n : _} -> Vect n a -> WithMArray n a b -> b
 168 | allocVect xs g =
 169 |   unsafeAlloc n $ \r,t =>
 170 |    let _ # t := writeVect r xs t
 171 |     in g r t
 172 | 
 173 | export
 174 | allocVectRev : {n : _} -> Vect n a -> WithMArray n a b -> b
 175 | allocVectRev xs g =
 176 |   unsafeAlloc n $ \r,t =>
 177 |    let _ # t := writeVectRev r n xs t
 178 |     in g r t
 179 | 
 180 | export
 181 | allocGen : (n : Nat) -> (Fin n -> a) -> WithMArray n a b -> b
 182 | allocGen n f g =
 183 |   unsafeAlloc n $ \r,t =>
 184 |    let _ # t := genFrom r n f t
 185 |     in g r t
 186 | 
 187 | export
 188 | allocIter : (n : Nat) -> (a -> a) -> a -> WithMArray n a b -> b
 189 | allocIter n f v g =
 190 |   unsafeAlloc n $ \r,t =>
 191 |    let _ # t := iterateFrom r n f v t
 192 |     in g r t
 193 | 
 194 | --------------------------------------------------------------------------------
 195 | --          Growing Arrays
 196 | --------------------------------------------------------------------------------
 197 | 
 198 | parameters {n : Nat}
 199 |            (r : MArray s n a)
 200 | 
 201 |   ||| Allocates a new mutable array and adds the elements from `r`
 202 |   ||| at its beginning.
 203 |   export
 204 |   mgrow : (m : Nat) -> (deflt : a) -> F1 s (MArray s (m+n) a)
 205 |   mgrow m deflt t =
 206 |     let tgt # t := marray1 (m+n) deflt t
 207 |         _   # t := copy r 0 0 n @{reflexive} @{lteAddLeft n} tgt t
 208 |      in tgt # t
 209 | 
 210 | --------------------------------------------------------------------------------
 211 | --          Appending Arrays
 212 | --------------------------------------------------------------------------------
 213 | 
 214 | parameters {m, n : Nat}
 215 |            (p : MArray s m a)
 216 |            (q : MArray s n a)
 217 | 
 218 |   ||| Allocates a new mutable array and adds the elements from `p`
 219 |   ||| at its beginning, followed by adding the elements from `q`.
 220 |   export
 221 |   mappend : F1 s (MArray s (m+n) a)
 222 |   mappend t =
 223 |     let tgt # t := unsafeMArray1 (m+n) t
 224 |         _   # t := copy p 0 0 m @{reflexive} @{lteAddRight m} tgt t
 225 |         _   # t := copy q 0 m n @{reflexive} tgt t
 226 |      in tgt # t
 227 | 
 228 | --------------------------------------------------------------------------------
 229 | --          Sub-Arrays
 230 | --------------------------------------------------------------------------------
 231 | 
 232 | 0 curLTE : (s : Ix m n) -> LTE c (ixToNat s) -> LTE c n
 233 | curLTE s lte = transitive lte $ ixLTE s
 234 | 
 235 | 0 curLT : (s : Ix (S m) n) -> LTE c (ixToNat s) -> LT c n
 236 | curLT s lte = let LTESucc p := ixLT s in LTESucc $ transitive lte p
 237 | 
 238 | parameters {n : Nat}
 239 |            (p : MArray s n a)
 240 | 
 241 |   ||| Filters the values in a mutable array according to the given predicate.
 242 |   export
 243 |   mfilterWithKey : (Fin n -> a -> Bool) -> F1 s (m ** MArray s m a)
 244 |   mfilterWithKey f t =
 245 |     let tft # t := unsafeMArray1 n t
 246 |      in go 0 n tft t
 247 |     where
 248 |       go :  (m, x : Nat)
 249 |          -> (q : MArray s n a)
 250 |          -> {auto v : Ix x n}
 251 |          -> {auto 0 prf : LTE m $ ixToNat v}
 252 |          -> F1 s (m ** MArray s m a)
 253 |       go m Z     q t =
 254 |         let q' # t := mtake q m @{curLTE v prf} t
 255 |          in (m ** q') # t
 256 |       go m (S j) q t =
 257 |         let j' # t := getIx p j t
 258 |          in case f (ixToFin v) j' of
 259 |               True  =>
 260 |                 let () # t := setNat q m @{curLT v prf} j' t
 261 |                  in go (S m) j q t
 262 |               False => go m j q t
 263 | 
 264 |   ||| Filters the values in a mutable array according to the given predicate.
 265 |   export %inline
 266 |   mfilter : (a -> Bool) -> F1 s (m ** MArray s m a)
 267 |   mfilter = mfilterWithKey . const
 268 | 
 269 |   ||| Filters the values in a mutable array according to the given predicate.
 270 |   export
 271 |   mmapMaybeWithKey : (Fin n -> a -> Maybe b) -> F1 s (m ** MArray s m b)
 272 |   mmapMaybeWithKey f t =
 273 |     let tft # t := unsafeMArray1 n t
 274 |      in go 0 n tft t
 275 |     where
 276 |       go :  (m, x : Nat)
 277 |          -> (q : MArray s n b)
 278 |          -> {auto v : Ix x n}
 279 |          -> {auto 0 prf : LTE m $ ixToNat v}
 280 |          -> F1 s (m ** MArray s m b)
 281 |       go m Z     q t =
 282 |         let q' # t := mtake q m @{curLTE v prf} t
 283 |          in (m ** q') # t
 284 |       go m (S j) q t =
 285 |         let j' # t := getIx p j t
 286 |          in case f (ixToFin v) j' of
 287 |               Just y =>
 288 |                 let () # t := setNat q m @{curLT v prf} y t
 289 |                  in go (S m) j q t
 290 |               Nothing => go m j q t
 291 | 
 292 | parameters {n : Nat}
 293 |            (m : Nat)
 294 |            (r : MArray s n a)
 295 | 
 296 |   ||| Drop `m` elements from a mutable array.
 297 |   export
 298 |   mdrop : F1 s (MArray s (n `minus` m) a)
 299 |   mdrop t =
 300 |     let tdt # t := unsafeMArray1 (n `minus` m) t
 301 |         _   # t := copy r (n `minus` (n `minus` m)) 0 (n `minus` m) @{dropLemma m n} tdt t
 302 |      in tdt # t
 303 | 
 304 | --------------------------------------------------------------------------------
 305 | --          Maps and Folds
 306 | --------------------------------------------------------------------------------
 307 | 
 308 | ||| Apply an effectful function `f` to each element of the mutable array.
 309 | |||
 310 | ||| Unlike `mmap1` and `mmap`, this function performs inplace mutation.
 311 | export
 312 | mupdate1 :  {n : _} -> (f : a -> F1 s a) -> MArray s n a -> F1' s
 313 | mupdate1 f r = go n
 314 |   where
 315 |     go :  (k: Nat) -> (x : Ix k n) => F1' s
 316 |     go 0     t = () # t
 317 |     go (S k) t =
 318 |      let v  # t := getIx r k t
 319 |          v2 # t := f v t
 320 |          _  # t := setIx r k v2 t
 321 |       in go k t
 322 | 
 323 | ||| Apply a function `f` to each element of the mutable array.
 324 | |||
 325 | ||| Unlike `mmap1` and `mmap`, this function performs inplace mutation.
 326 | export
 327 | mupdate :  {n : _} -> (f : a -> a) -> MArray s n a -> F1' s
 328 | mupdate f r = go n
 329 |   where
 330 |     go :  (k: Nat) -> (0 lt : LTE k n) => F1' s
 331 |     go 0     t = () # t
 332 |     go (S k) t =
 333 |      let v  # t := getNat r k t
 334 |          _  # t := setNat r k (f v) t
 335 |       in go k t
 336 | 
 337 | ||| Apply a function `f` to each element of the mutable array.
 338 | export
 339 | mmap1 :  {n : _} -> (f : a -> F1 s b) -> MArray s n a -> F1 s (MArray s n b)
 340 | mmap1 f r t =
 341 |   let tmt # t := unsafeMArray1 n t
 342 |       _   # t := genFrom' tmt n go t
 343 |     in tmt # t
 344 |   where
 345 |     go :  Fin n -> F1 s b
 346 |     go x t = let x'  # t := get r x t in f x' t
 347 | 
 348 | ||| Apply a function `f` to each element of the mutable array.
 349 | export %inline
 350 | mmap :  {n : _} -> (f : a -> b) -> MArray s n a -> F1 s (MArray s n b)
 351 | mmap f = mmap1 (\x,t => f x # t)
 352 | 
 353 | --------------------------------------------------------------------------------
 354 | --          Reversing Arrays
 355 | --------------------------------------------------------------------------------
 356 | 
 357 | parameters {k : Nat}
 358 |            (r : MArray s k a)
 359 | 
 360 |   ||| Reverse the order of elements in a mutable array.
 361 |   export
 362 |   mreverse : F1 s (MArray s k a)
 363 |   mreverse t =
 364 |     let trt # t := unsafeMArray1 k t
 365 |      in go k trt t
 366 |     where
 367 |       go :  (x : Nat)
 368 |          -> (q : MArray s k a)
 369 |          -> {auto v : Ix x k}
 370 |          -> {auto 0 _ : LTE x k}
 371 |          -> F1 s (MArray s k a)
 372 |       go Z     q t =
 373 |         q # t
 374 |       go (S j) q t =
 375 |         let j' # t := getIx r j t
 376 |             () # t := setNat q j j' t
 377 |           in go j q t
 378 | 
 379 | --------------------------------------------------------------------------------
 380 | --          Linear Utilities
 381 | --------------------------------------------------------------------------------
 382 | 
 383 |   ||| Accumulate the values stored in a mutable array.
 384 |   export
 385 |   foldrLin : (a -> b -> b) -> b -> F1 s b
 386 |   foldrLin f = go k
 387 |     where
 388 |       go : (n : Nat) -> (0 lt : LTE n k) => b -> F1 s b
 389 |       go 0     v t = v # t
 390 |       go (S k) v t =
 391 |        let el # t := getNat r k t
 392 |         in go k (f el v) t
 393 | 
 394 |   ||| Store the values in a mutable array in a `Vect` of the same size.
 395 |   export
 396 |   toVect : F1 s (Vect k a)
 397 |   toVect = go [] k
 398 |     where
 399 |       go : Vect m a -> (n : Nat) -> (0 lt : LTE n k) => F1 s (Vect (m+n) a)
 400 |       go vs 0     t = rewrite plusCommutative m 0 in vs # t
 401 |       go vs (S x) t =
 402 |         let v # t := getNat r x t
 403 |          in rewrite sym (plusSuccRightSucc m x) in go (v::vs) x t
 404 | 
 405 |   ||| Extract and convert the values stored in a mutable array
 406 |   ||| and store them in a `Vect` of the same size.
 407 |   export
 408 |   toVectWith : (Fin k -> a -> b) -> F1 s (Vect k b)
 409 |   toVectWith f = go [] k
 410 |     where
 411 |       go : Vect m b -> (n : Nat) -> (0 lt : LTE n k) => F1 s (Vect (m+n) b)
 412 |       go vs 0     t = rewrite plusCommutative m 0 in vs # t
 413 |       go vs (S x) t =
 414 |         let v # t := getNat r x t
 415 |             vb       := f (natToFinLT x) v
 416 |          in rewrite sym (plusSuccRightSucc m x) in go (vb::vs) x t
 417 |