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