]> Witch of Git - ivy/blob - rt/src/lib.rs
[rt] Improve memory ordering of incref/decref
[ivy] / rt / src / lib.rs
1 use std::sync::atomic::{Ordering, AtomicU32};
2
3 const _STDIN: i32 = 0;
4 const _STDOUT: i32 = 1;
5 const STDERR: i32 = 2;
6
7 macro_rules! trace {
8 ($fmt:literal $(, $arg:expr)* $(,)?) => {
9 if std::env::var("IVY_RT_TRACE").is_ok() {
10 eprintln!($fmt, $($arg),*);
11 }
12 }
13 }
14
15 #[repr(u8)]
16 #[derive(PartialEq, Eq)]
17 pub enum ObjTag {
18 Lam = 0,
19 Int = 1,
20 }
21
22 #[repr(C)]
23 pub struct ObjHeader {
24 tag: ObjTag,
25 rc: AtomicU32,
26 }
27
28 #[repr(C)]
29 pub struct ObjInt {
30 tag: ObjTag,
31 rc: AtomicU32,
32 value: i64,
33 }
34
35 #[repr(C)]
36 pub struct ObjLam {
37 tag: ObjTag,
38 upvars: u16,
39 rc: AtomicU32,
40 func: extern "C" fn(&ObjLam) -> Obj,
41 params: u16,
42 filled: u16,
43 }
44
45 #[derive(Clone, Copy)]
46 #[repr(C)]
47 pub union Obj {
48 int: i64,
49 header: *mut ObjHeader,
50 box_lam: *mut ObjLam,
51 }
52
53 mod sys {
54 extern "C" {
55 pub fn write(fd: i32, buf: *const u8, len: usize) -> isize;
56 pub fn exit(code: i32) -> !;
57 pub fn malloc(size: usize) -> *mut u8;
58 pub fn free(ptr: *mut u8);
59 }
60 }
61
62 #[no_mangle]
63 pub unsafe extern "C" fn ivy_debug(obj: Obj) -> Obj {
64 println!("DEBUG {:016x}", obj.int);
65 obj
66 }
67
68 #[no_mangle]
69 pub unsafe extern "C" fn ivy_abort(msg: *const u8, len: usize) -> ! {
70 sys::write(STDERR, msg, len);
71 sys::exit(1);
72 }
73
74 #[no_mangle]
75 pub unsafe extern "C" fn ivy_exit(code: i32) -> ! {
76 sys::exit(code)
77 }
78
79 #[no_mangle]
80 pub unsafe extern "C" fn ivy_check_int(obj: Obj) {
81 if !obj.is_int() {
82 panic!("ivy_check_int called with non-integer object {:016x}.", obj.int);
83 }
84 }
85
86 #[no_mangle]
87 pub unsafe extern "C" fn ivy_check_lam(obj: Obj) {
88 if !obj.is_lam() {
89 panic!("ivy_check_lam called with non-lambda object {:016x}.", obj.int);
90 }
91 }
92
93 // This should probably be a macro rather than a call?
94 // But it might be good to have it for completeness.
95 // Or maybe it's valuable if we want to support big integers.
96 #[no_mangle]
97 pub unsafe extern "C" fn ivy_make_int(value: i64) -> Obj {
98 Obj { int: value << 1 }
99 }
100
101 #[no_mangle]
102 pub unsafe extern "C" fn ivy_make_lam(func: extern "C" fn(&ObjLam) -> Obj, params: u16, upvars: u16) -> Obj {
103 let size = ObjLam::size_of(params, upvars);
104 let box_lam = sys::malloc(size) as *mut ObjLam;
105 box_lam.write(ObjLam {
106 tag: ObjTag::Lam,
107 upvars,
108 rc: AtomicU32::new(0),
109 func,
110 params,
111 filled: 0,
112 });
113 (*box_lam)
114 .raw_fields_mut()
115 .write_bytes(0, (params + upvars) as usize);
116 trace!("MAKE {:016x} {:016x}", box_lam as usize, func as usize);
117 Obj { box_lam }
118 }
119
120 #[no_mangle]
121 pub unsafe extern "C" fn ivy_free(obj: Obj) {
122 if !obj.is_box() {
123 return;
124 }
125 sys::free(obj.header as *mut u8)
126 }
127
128 #[no_mangle]
129 pub unsafe extern "C" fn ivy_incref(obj: Obj) {
130 obj.incref();
131 }
132
133 #[no_mangle]
134 pub unsafe extern "C" fn ivy_decref(obj: Obj) {
135 obj.decref();
136 }
137
138 #[no_mangle]
139 pub unsafe extern "C" fn ivy_clone(obj: Obj) -> Obj {
140 if obj.is_null() || !obj.is_box() {
141 return obj;
142 }
143 if obj.is_int() {
144 unimplemented!("copying boxed integers")
145 }
146 let lam = &*obj.box_lam;
147 let size = lam.size();
148 let data = sys::malloc(size);
149 core::ptr::copy(obj.box_lam as *const u8, data, size);
150 let box_lam = data as *mut ObjLam;
151 *(*box_lam).rc.get_mut() = 0;
152 trace!("COPY {:016x} {:016x}", obj.int, box_lam as usize);
153 Obj { box_lam }
154 }
155
156 #[no_mangle]
157 pub unsafe extern "C" fn ivy_app(fun: Obj, arg: Obj) -> Obj {
158 ivy_app_mut(ivy_clone(fun), arg)
159 }
160
161 #[no_mangle]
162 pub unsafe extern "C" fn ivy_app_mut(fun: Obj, arg: Obj) -> Obj {
163 trace!("APP {:016x} {:016x}", fun.int, arg.int);
164 if !fun.is_lam() {
165 panic!("ivy_app called with a non-lam as the function: {:016x}.", fun.int);
166 }
167 let lam = &mut *fun.box_lam;
168 if lam.filled < lam.params {
169 if arg.is_null() {
170 println!(
171 "Lam @ {:016x} ({:016x}) has {} of {} arguments filled.",
172 fun.int, lam.func as usize, lam.filled, lam.params
173 );
174 panic!("ivy_app called with a null arg.");
175 }
176 arg.incref();
177 let idx = lam.filled as usize;
178 lam.params_mut()[idx] = arg;
179 lam.filled += 1;
180 } else if lam.params == lam.filled {
181 if !arg.is_null() {
182 panic!("ivy_app called for a 0-arity application with a non-null arg.");
183 }
184 }
185
186 if lam.params == lam.filled {
187 trace!("RUN {:016x}", fun.int);
188 (lam.func)(lam)
189 } else {
190 trace!("UPD8 {:016x}", fun.int);
191 fun.incref();
192 fun
193 }
194 }
195
196 impl Obj {
197 fn is_null(self) -> bool {
198 unsafe { self.int == 0 }
199 }
200
201 fn is_box(self) -> bool {
202 !self.is_null() && unsafe { self.int & 1 == 0 }
203 }
204
205 unsafe fn is_int(self) -> bool {
206 !self.is_null() && (!self.is_box() || (*self.header).tag == ObjTag::Int)
207 }
208
209 unsafe fn is_lam(self) -> bool {
210 self.is_box() && (*self.header).tag == ObjTag::Lam
211 }
212
213 unsafe fn incref(self) {
214 if !self.is_box() {
215 return;
216 }
217 // Ordering::Relaxed is appropriate here, since we assume that each thread with access to a
218 // reference owns at least one reference (rather than simply borrowing it). Therefore,
219 // another thread cannot decrement it to 0 while we are performing this increment (since we
220 // own a reference), so we only need consistency and not ordering.
221 (*self.header).rc.fetch_add(1, Ordering::Relaxed);
222 }
223
224 unsafe fn decref(self) {
225 if !self.is_box() {
226 return;
227 }
228 // Ordering::AcqRel is appropriate here. I believe we need the Acquire in order to ensure
229 // we see all previous increments/decrements, so we can properly see that the decref is
230 // decrementing to 0, and we need the Release in order to ensure that we see all writes to
231 // the memory before we deallocate.
232 // (Check against 1 instead of 0 since we're loading the old refcount.)
233 if (*self.header).rc.fetch_sub(1, Ordering::AcqRel) == 1 {
234 self.dealloc();
235 }
236 }
237
238 unsafe fn dealloc(self) {
239 if !self.is_box() {
240 return;
241 }
242 if self.is_lam() {
243 let lam = &mut *self.box_lam;
244 for param in lam.params_mut() {
245 param.decref();
246 }
247 for upvar in lam.upvars_mut() {
248 upvar.decref();
249 }
250 }
251 sys::free(self.header as *mut u8);
252 }
253 }
254
255 impl ObjLam {
256 fn size_of(params: u16, upvars: u16) -> usize {
257 core::mem::size_of::<ObjLam>() + params as usize * 8 + upvars as usize * 8
258 }
259
260 fn size(&self) -> usize {
261 ObjLam::size_of(self.params, self.upvars)
262 }
263
264 unsafe fn raw_fields_mut(&mut self) -> *mut Obj {
265 (self as *mut ObjLam).add(1) as *mut Obj
266 }
267
268 unsafe fn params_mut(&mut self) -> &mut [Obj] {
269 let ptr = self.raw_fields_mut();
270 core::slice::from_raw_parts_mut(ptr, self.params as usize)
271 }
272
273 unsafe fn upvars_mut(&mut self) -> &mut [Obj] {
274 let ptr = self.raw_fields_mut().add(self.params as usize);
275 core::slice::from_raw_parts_mut(ptr, self.upvars as usize)
276 }
277 }