0 | {--
  1 | Copyright (C) 2026  Joel Berkeley
  2 |
  3 | This program is free software: you can redistribute it and/or modify
  4 | it under the terms of the GNU Affero General Public License as published
  5 | by the Free Software Foundation, either version 3 of the License, or
  6 | (at your option) any later version.
  7 |
  8 | This program is distributed in the hope that it will be useful,
  9 | but WITHOUT ANY WARRANTY; without even the implied warranty of
 10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 11 | GNU Affero General Public License for more details.
 12 |
 13 | You should have received a copy of the GNU Affero General Public License
 14 | along with this program.  If not, see <https://www.gnu.org/licenses/>.
 15 | --}
 16 | ||| For internal spidr use only.
 17 | module Compiler.Array
 18 |
 19 | import Compiler.FFI
 20 | import Util
 21 |
 22 | ffi : String -> String
 23 | ffi = libxla "c/array.h"
 24 |
 25 | export
 26 | %foreign (ffi "sizeof_ptr")
 27 | sizeofPtr : Bits64
 28 |
 29 | export
 30 | %foreign (ffi "get_array_void_ptr")
 31 | prim__getArrayVoidPtr : AnyPtr -> Bits64 -> AnyPtr
 32 |
 33 | export
 34 | %foreign (ffi "set_array_void_ptr")
 35 | prim__setArrayVoidPtr : AnyPtr -> Bits64 -> AnyPtr -> PrimIO ()
 36 |
 37 | public export
 38 | data Array a = MkArray GCAnyPtr Bits64
 39 |
 40 | public export
 41 | interface ArrayType a where
 42 |   elemSize : Bits64
 43 |   get : Array a -> Nat -> a
 44 |   set : HasIO io => Array a -> Nat -> a -> io ()
 45 |
 46 | %foreign (ffi "sizeof_int32_t")
 47 | sizeofInt32 : Bits64
 48 |
 49 | %foreign (ffi "get_array_int32_t")
 50 | prim__getArrayInt32 : GCAnyPtr -> Bits64 -> Int32
 51 |
 52 | %foreign (ffi "set_array_int32_t")
 53 | prim__setArrayInt32 : GCAnyPtr -> Bits64 -> Int32 -> PrimIO ()
 54 |
 55 | export
 56 | ArrayType Int32 where
 57 |   elemSize = sizeofInt32
 58 |   get (MkArray arr _) idx = prim__getArrayInt32 arr (cast idx)
 59 |   set (MkArray arr _) idx x = primIO $ prim__setArrayInt32 arr (cast idx) x
 60 |
 61 | %foreign (ffi "sizeof_int64_t")
 62 | sizeofInt64 : Bits64
 63 |
 64 | %foreign (ffi "get_array_int64_t")
 65 | prim__getArrayInt64 : GCAnyPtr -> Bits64 -> Int64
 66 |
 67 | %foreign (ffi "set_array_int64_t")
 68 | prim__setArrayInt64 : GCAnyPtr -> Bits64 -> Int64 -> PrimIO ()
 69 |
 70 | export
 71 | ArrayType Int64 where
 72 |   elemSize = sizeofInt64
 73 |   get (MkArray arr _) idx = prim__getArrayInt64 arr (cast idx)
 74 |   set (MkArray arr _) idx x = primIO $ prim__setArrayInt64 arr (cast idx) x
 75 |
 76 | %foreign (ffi "sizeof_uint32_t")
 77 | sizeofUInt32 : Bits64
 78 |
 79 | %foreign (ffi "get_array_uint32_t")
 80 | prim__getArrayUInt32 : GCAnyPtr -> Bits64 -> Bits32
 81 |
 82 | %foreign (ffi "set_array_uint32_t")
 83 | prim__setArrayUInt32 : GCAnyPtr -> Bits64 -> Bits32 -> PrimIO ()
 84 |
 85 | export
 86 | ArrayType Bits32 where
 87 |   elemSize = sizeofUInt32
 88 |   get (MkArray arr _) idx = prim__getArrayUInt32 arr (cast idx)
 89 |   set (MkArray arr _) idx x = primIO $ prim__setArrayUInt32 arr (cast idx) x
 90 |
 91 | %foreign (ffi "sizeof_uint64_t")
 92 | sizeofUInt64 : Bits64
 93 |
 94 | %foreign (ffi "get_array_uint64_t")
 95 | prim__getArrayUInt64 : GCAnyPtr -> Bits64 -> Bits64
 96 |
 97 | %foreign (ffi "set_array_uint64_t")
 98 | prim__setArrayUInt64 : GCAnyPtr -> Bits64 -> Bits64 -> PrimIO ()
 99 |
100 | export
101 | ArrayType Bits64 where
102 |   elemSize = sizeofUInt64
103 |   get (MkArray arr _) idx = prim__getArrayUInt64 arr (cast idx)
104 |   set (MkArray arr _) idx x = primIO $ prim__setArrayUInt64 arr (cast idx) x
105 |
106 | %foreign (ffi "sizeof_double")
107 | sizeofDouble : Bits64
108 |
109 | %foreign (ffi "get_array_double")
110 | prim__getArrayDouble : GCAnyPtr -> Bits64 -> Double
111 |
112 | %foreign (ffi "set_array_double")
113 | prim__setArrayDouble : GCAnyPtr -> Bits64 -> Double -> PrimIO ()
114 |
115 | export
116 | ArrayType Double where
117 |   elemSize = sizeofDouble
118 |   get (MkArray arr _) idx = prim__getArrayDouble arr (cast idx)
119 |   set (MkArray arr _) idx x = primIO $ prim__setArrayDouble arr (cast idx) x
120 |
121 | %foreign (ffi "sizeof_bool")
122 | sizeofBool : Bits64
123 |
124 | %foreign (ffi "get_array_bool")
125 | prim__getArrayBool : GCAnyPtr -> Bits64 -> Int
126 |
127 | %foreign (ffi "set_array_bool")
128 | prim__setArrayBool : GCAnyPtr -> Bits64 -> Int -> PrimIO ()
129 |
130 | export
131 | ArrayType Bool where
132 |   elemSize = sizeofBool
133 |   get (MkArray arr _) idx = cIntToBool $ prim__getArrayBool arr (cast idx)
134 |   set (MkArray arr _) idx x = primIO $ prim__setArrayBool arr (cast idx) (boolToCInt x)
135 |
136 | ||| An uninitialized array. Contents may be anything.
137 | export
138 | mkArray : HasIO io => ArrayType a => Nat -> io (Array a)
139 | mkArray elemCount = do
140 |   let bytes : Bits64 = cast elemCount * elemSize {a}
141 |   arr <- malloc (cast bytes)
142 |   arr <- onCollectAny' arr free
143 |   pure $ MkArray arr (cast elemCount)
144 |
145 | export
146 | fromList : HasIO io => ArrayType a => List a -> io (Array a)
147 | fromList xs = do
148 |   arr <- mkArray (length xs)
149 |   traverse_ (\(idx, x) => set arr (cast idx) x) (enumerate xs)
150 |   pure arr
151 |
152 | export
153 | toList : ArrayType a => Array a -> List a
154 | toList arr@(MkArray _ len) = Prelude.map (\i => get arr i) $ range (cast {to = Nat} len)
155 |