0 | module Data.DArray
 1 |
 2 | import Data.Array.Core
 3 | import Data.Linear.Token
 4 | import Data.List.Quantifiers
 5 | import public Data.Enum
 6 |
 7 | %default total
 8 |
 9 | export
10 | record DArray (i : Type) (p : i -> Type) where
11 |   constructor DA
12 |   arr : AnyPtr
13 |
14 | export %inline
15 | at : Enum i n => DArray i p -> (x : i) -> p x
16 | at (DA ad) x = believe_me $ prim__arrayGet ad (enumToInteger x)
17 |
18 | export
19 | record MDArray (s : Type) (i : Type) (p : i -> Type) where
20 |   constructor MDA
21 |   arr : AnyPtr
22 |
23 | export %inline
24 | dget : Enum i n => MDArray s i p -> (x : i) -> F1 s (p x)
25 | dget (MDA ad) x t =
26 |   believe_me (prim__arrayGet ad $ enumToInteger x) # t
27 |
28 | export %inline
29 | dset : Enum i n => MDArray s i p -> (x : i) -> p x -> F1' s
30 | dset (MDA ad) x v =
31 |   ffi (prim__arraySet ad (enumToInteger x) (believe_me v))
32 |
33 | export %inline
34 | unsafeFreeze : MDArray s i p -> F1 s (DArray i p)
35 | unsafeFreeze (MDA ad) t = DA ad # t
36 |
37 | export
38 | freeze : {n : _} -> Enum i n => MDArray s i p -> F1 s (DArray i p)
39 | freeze (MDA src) t =
40 |   let dst # t := ffi (prim__emptyArray $ cast n) t
41 |       _   # t := ffi (prim__copyArray src 0 (cast n) dst 0) t
42 |    in DA dst # t
43 |
44 | parameters (0 i       : Type)
45 |            (0 p       : i -> Type)
46 |            {n         : Bits32}
47 |            {auto enum : Enum i n}
48 |
49 |   export %inline
50 |   unsafeMDArray1 : F1 s (MDArray s i p)
51 |   unsafeMDArray1 t =
52 |     let p # t := ffi (prim__emptyArray $ cast n) t in MDA p # t
53 |
54 |   ||| A safe constructor for mutable dependent arrays.
55 |   export
56 |   mdarrayAll1 : All p Finite.values -> F1 s (MDArray s i p)
57 |   mdarrayAll1 vs t = let md # t := unsafeMDArray1 t in go values vs md t
58 |     where
59 |       go : (is : List i) -> All p is -> MDArray s i p -> F1 s (MDArray s i p)
60 |       go []      []      m t = m # t
61 |       go (x::xs) (v::vs) m t = let _ # t := dset m x v t in go xs vs m t
62 |
63 |   ||| Like `mdarrayAll1` but accepts the list of values as an auto implicit.
64 |   export %inline
65 |   mdarrayAuto1 : All p Finite.values => F1 s (MDArray s i p)
66 |   mdarrayAuto1 = mdarrayAll1 %search
67 |
68 |   ||| A safe constructor for mutable dependent arrays.
69 |   export
70 |   mdarray1 : (val : (v : i) -> p v) -> F1 s (MDArray s i p)
71 |   mdarray1 val t = let md # t := unsafeMDArray1 t in go values md t
72 |     where
73 |       go : List i -> MDArray s i p -> F1 s (MDArray s i p)
74 |       go []        m t = m # t
75 |       go (x :: xs) m t = let _ # t := dset m x (val x) t in go xs m t
76 |
77 |   export %inline
78 |   darray : (val : (v : i) -> p v) -> DArray i p
79 |   darray val = run1 $ \t => let m # t := mdarray1 val t in freeze m t
80 |
81 |   export %inline
82 |   darrayAll : All p Finite.values -> DArray i p
83 |   darrayAll vs = run1 $ \t => let m # t := mdarrayAll1 vs t in freeze m t
84 |
85 |   export %inline
86 |   darrayAuto : All p Finite.values => DArray i p
87 |   darrayAuto = darrayAll %search
88 |