0 | module Data.IOArray
 1 |
 2 | import Data.IOArray.Prims
 3 |
 4 | %default total
 5 |
 6 | export
 7 | record IOArray elem where
 8 |   constructor MkIOArray
 9 |   maxSize : Int
10 |   content : ArrayData (Maybe elem)
11 |
12 | export
13 | max : IOArray elem -> Int
14 | max = maxSize
15 |
16 | export
17 | newArray : HasIO io => Int -> io (IOArray elem)
18 | newArray size
19 |     = pure (MkIOArray size !(primIO (prim__newArray size Nothing)))
20 |
21 | export
22 | writeArray : HasIO io => IOArray elem -> Int -> elem -> io Bool
23 | writeArray arr pos el
24 |     = if pos < 0 || pos >= max arr
25 |          then pure False
26 |          else do primIO (prim__arraySet (content arr) pos (Just el))
27 |                  pure True
28 |
29 | export
30 | readArray : HasIO io => IOArray elem -> Int -> io (Maybe elem)
31 | readArray arr pos
32 |     = if pos < 0 || pos >= max arr
33 |          then pure Nothing
34 |          else primIO (prim__arrayGet (content arr) pos)
35 |
36 | -- Make a new array of the given size with the elements copied from the
37 | -- other array
38 | export
39 | newArrayCopy : HasIO io =>
40 |                (newsize : Int) -> IOArray elem -> io (IOArray elem)
41 | newArrayCopy newsize arr
42 |     = do let newsize' = if newsize < max arr then max arr else newsize
43 |          arr' <- newArray newsize'
44 |          copyFrom (content arr) (content arr') (max arr - 1)
45 |          pure arr'
46 |   where
47 |     copyFrom : ArrayData (Maybe elem) ->
48 |                ArrayData (Maybe elem) ->
49 |                Int -> io ()
50 |     copyFrom old new pos
51 |         = if pos < 0
52 |              then pure ()
53 |              else do el <- primIO $ prim__arrayGet old pos
54 |                      primIO $ prim__arraySet new pos el
55 |                      copyFrom old new $ assert_smaller pos (pos - 1)
56 |
57 | export
58 | toList : HasIO io => IOArray elem -> io (List (Maybe elem))
59 | toList arr = iter 0 (max arr) []
60 |   where
61 |     iter : Int -> Int -> List (Maybe elem) -> io (List (Maybe elem))
62 |     iter pos end acc
63 |          = if pos >= end
64 |               then pure (reverse acc)
65 |               else do el <- readArray arr pos
66 |                       assert_total (iter (pos + 1) end (el :: acc))
67 |
68 | export
69 | fromList : HasIO io => List (Maybe elem) -> io (IOArray elem)
70 | fromList ns
71 |     = do arr <- newArray (cast (length ns))
72 |          addToArray 0 ns arr
73 |          pure arr
74 |   where
75 |     addToArray : Int -> List (Maybe elem) -> IOArray elem -> io ()
76 |     addToArray loc [] arr = pure ()
77 |     addToArray loc (Nothing :: ns) arr = addToArray (loc + 1) ns arr
78 |     addToArray loc (Just el :: ns) arr
79 |         = do primIO $ prim__arraySet (content arr) loc (Just el)
80 |              addToArray (loc + 1) ns arr
81 |