unlab_gpu/
io.rs

1//
2// Copyright (c) 2025-2026 Łukasz Szpakowski
3//
4// This Source Code Form is subject to the terms of the Mozilla Public
5// License, v. 2.0. If a copy of the MPL was not distributed with this
6// file, You can obtain one at https://mozilla.org/MPL/2.0/.
7//
8//! An I/O module.
9use std::collections::BTreeMap;
10use std::collections::HashMap;
11use std::fs::File;
12use std::io;
13use std::io::BufReader;
14use std::io::BufWriter;
15use std::io::ErrorKind;
16use std::io::Read;
17use std::io::Write;
18use std::mem::size_of;
19use std::path::Path;
20use std::sync::Arc;
21use std::sync::RwLock;
22use std::sync::Weak;
23use crate::env::*;
24use crate::error::*;
25use crate::tree::*;
26use crate::utils::*;
27use crate::value::*;
28
29fn read_magic(r: &mut dyn Read) -> Result<()>
30{
31    let mut buf = [0u8; 6];
32    match r.read_exact(&mut buf) {
33        Ok(()) => {
34            if &buf != b"unlab1" {
35                return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid data format")));
36            }
37            Ok(())
38        },
39        Err(err) => Err(Error::Io(err)),
40    }
41}
42
43fn read_u8(r: &mut dyn Read) -> Result<u8>
44{
45    let mut buf = [0u8; 1];
46    match r.read_exact(&mut buf) {
47        Ok(()) => Ok(buf[0]),
48        Err(err) => Err(Error::Io(err)),
49    }
50}
51
52fn read_bool(r: &mut dyn Read) -> Result<bool>
53{ 
54    match read_u8(r)? {
55        0 => Ok(false),
56        1 => Ok(true),
57        _ => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid boolean"))),
58    }
59}
60
61fn read_u64(r: &mut dyn Read) -> Result<u64>
62{
63    let mut buf = [0u8; 8];
64    match r.read_exact(&mut buf) {
65        Ok(()) => Ok(u64::from_le_bytes(buf)),
66        Err(err) => Err(Error::Io(err)),
67    }
68}
69
70fn read_usize(r: &mut dyn Read) -> Result<usize>
71{
72    let n = read_u64(r)?;
73    if n > (usize::MAX as u64) {
74        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large unsigned number")));
75    }
76    Ok(n as usize)
77}
78
79fn read_i64(r: &mut dyn Read) -> Result<i64>
80{
81    let mut buf = [0u8; 8];
82    match r.read_exact(&mut buf) {
83        Ok(()) => Ok(i64::from_le_bytes(buf)),
84        Err(err) => Err(Error::Io(err)),
85    }
86}
87
88fn read_f32(r: &mut dyn Read) -> Result<f32>
89{
90    let mut buf = [0u8; 4];
91    match r.read_exact(&mut buf) {
92        Ok(()) => Ok(f32::from_le_bytes(buf)),
93        Err(err) => Err(Error::Io(err)),
94    }
95}
96
97fn read_string(r: &mut dyn Read) -> Result<String>
98{
99    let len = read_usize(r)?;
100    let mut buf = vec![0u8; len];
101    match r.read_exact(&mut buf) {
102        Ok(()) => {
103            match String::from_utf8(buf) {
104                Ok(s) => Ok(s),
105                Err(err) => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, format!("invalid string: {}", err)))),
106            }
107        },
108        Err(err) => Err(Error::Io(err)),
109    }
110}
111
112fn write_magic(w: &mut dyn Write) -> Result<()>
113{
114    match w.write_all(b"unlab1") {
115        Ok(()) => Ok(()),
116        Err(err) => Err(Error::Io(err)),
117    }
118}
119
120fn write_u8(w: &mut dyn Write, n: u8) -> Result<()>
121{
122    let buf = [n];
123    match w.write_all(&buf) {
124        Ok(()) => Ok(()),
125        Err(err) => Err(Error::Io(err)),
126    }
127}
128
129fn write_bool(w: &mut dyn Write, b: bool) -> Result<()>
130{ write_u8(w, if b { 1 } else { 0 }) }
131
132fn write_u64(w: &mut dyn Write, n: u64) -> Result<()>
133{
134    let buf = n.to_le_bytes();
135    match w.write_all(&buf) {
136        Ok(()) => Ok(()),
137        Err(err) => Err(Error::Io(err)),
138    }
139}
140
141fn write_usize(w: &mut dyn Write, n: usize) -> Result<()>
142{ write_u64(w, n as u64) }
143
144fn write_i64(w: &mut dyn Write, n: i64) -> Result<()>
145{
146    let buf = n.to_le_bytes();
147    match w.write_all(&buf) {
148        Ok(()) => Ok(()),
149        Err(err) => Err(Error::Io(err)),
150    }
151}
152
153fn write_f32(w: &mut dyn Write, n: f32) -> Result<()>
154{
155    let buf = n.to_le_bytes();
156    match w.write_all(&buf) {
157        Ok(()) => Ok(()),
158        Err(err) => Err(Error::Io(err)),
159    }
160}
161
162fn write_str(w: &mut dyn Write, s: &str) -> Result<()>
163{
164    write_usize(w, s.as_bytes().len())?;
165    match w.write_all(s.as_bytes()) {
166        Ok(()) => Ok(()),
167        Err(err) => Err(Error::Io(err)),
168    }
169}
170
171const VALUE_NONE: u8 = 0;
172const VALUE_BOOL: u8 = 1;
173const VALUE_INT: u8 = 2;
174const VALUE_FLOAT: u8 = 3;
175const VALUE_OBJECT: u8 = 4;
176const VALUE_REF: u8 = 5;
177const VALUE_WEAK: u8 = 6;
178const VALUE_WEAK_NONE: u8 = 7;
179const VALUE_OBJECT_INDEX: u8 = 8;
180const VALUE_REF_INDEX: u8 = 9;
181const VALUE_WEAK_INDEX: u8 = 10;
182
183const OBJECT_STRING: u8 = 0;
184const OBJECT_INT_RANGE: u8 = 1;
185const OBJECT_FLOAT_RANGE: u8 = 2;
186const OBJECT_MATRIX: u8 = 3;
187const OBJECT_FUN: u8 = 4;
188const OBJECT_BUILTIN_FUN: u8 = 5;
189const OBJECT_MATRIX_ARRAY: u8 = 6;
190const OBJECT_MATRIX_ROW_SLICE: u8 = 7;
191const OBJECT_ERROR: u8 = 8;
192
193const MUT_OBJECT_ARRAY: u8 = 0;
194const MUT_OBJECT_STRUCT: u8 = 1;
195
196const MATRIX_ARRAY_OBJECT: u8 = 0;
197const MATRIX_ARRAY_INDEX: u8 = 1;
198
199struct ObjectTab<T>
200{
201    indices: HashMap<*const T, usize>,
202    objects: HashMap<usize, Arc<T>>,
203    count: usize,
204}
205
206impl<T> ObjectTab<T>
207{
208    fn new() -> Self
209    { ObjectTab { indices: HashMap::new(), objects: HashMap::new(), count: 0, } }
210    
211    fn index(&self, object: &Arc<T>) -> Option<usize>
212    { self.indices.get(&Arc::as_ptr(object)).map(|i| *i) }
213    
214    fn object(&self, idx: usize) -> Option<&Arc<T>>
215    { self.objects.get(&idx) }
216    
217    fn add_object(&mut self, object: Arc<T>) -> bool
218    {
219        self.indices.insert(Arc::as_ptr(&object), self.count);
220        self.objects.insert(self.count, object);
221        match self.count.checked_add(1) {
222            Some(new_count) => {
223                self.count = new_count;
224                true
225            },
226            None => false, 
227        }
228    }
229}
230
231fn checked_mul_row_count_and_col_count(row_count: usize, col_count: usize) -> Result<usize>
232{
233    if row_count > (isize::MAX as usize) {
234        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large number of rows")));
235    }
236    if col_count > (isize::MAX as usize) {
237        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large number of columns")));
238    }
239    match row_count.checked_mul(col_count) {
240        Some(len) => {
241            if len > (isize::MAX as usize) {
242                return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large number of matrix elements")));
243            }
244            match (len as isize).checked_mul(size_of::<f32>() as isize) {
245                Some(_) => Ok(len as usize),
246                None => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large number of matrix elements"))),
247            }
248        },
249        None => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large number of matrix elements"))),
250    }
251}
252
253fn read_object(r: &mut dyn Read, env: &Env, object_tab: &mut ObjectTab<Object>) -> Result<Arc<Object>>
254{
255    let object = match read_u8(r)? {
256        OBJECT_STRING => Arc::new(Object::String(read_string(r)?)),
257        OBJECT_INT_RANGE => {
258            let from = read_i64(r)?;
259            let to = read_i64(r)?;
260            let step = read_i64(r)?;
261            Arc::new(Object::IntRange(from, to, step))
262        },
263        OBJECT_FLOAT_RANGE => {
264            let from = read_f32(r)?;
265            let to = read_f32(r)?;
266            let step = read_f32(r)?;
267            Arc::new(Object::FloatRange(from, to, step))
268        },
269        OBJECT_MATRIX => {
270            let row_count = read_usize(r)?;
271            let col_count = read_usize(r)?;
272            let is_transposed = read_bool(r)?;
273            let len = checked_mul_row_count_and_col_count(row_count, col_count)?;
274            let mut xs = vec![0.0f32; len];
275            for i in 0..len {
276                xs[i] = read_f32(r)?;
277            }
278            if !is_transposed {
279                Arc::new(Object::Matrix(matrix_create_and_set_elems(row_count, col_count, xs.as_slice())?))
280            } else {
281                Arc::new(Object::Matrix(matrix_create_and_set_elems(col_count, row_count, xs.as_slice())?.transpose()))
282            }
283        },
284        OBJECT_FUN => {
285            let ident_count = read_usize(r)?;
286            let mut idents: Vec<String> = Vec::new();
287            for _ in 0..ident_count {
288                idents.push(read_string(r)?);
289            }
290            let ident = read_string(r)?;
291            let name = Name::Abs(idents, ident);
292            match env.var(&name)? {
293                Some(Value::Object(object)) => object.clone(),
294                Some(_) => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, format!("invalid function type {}", name).as_str()))),
295                None => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, format!("undefined function {}", name).as_str()))),
296            }
297        },
298        OBJECT_BUILTIN_FUN => {
299            let ident = read_string(r)?;
300            let name = Name::Abs(Vec::new(), ident);
301            match env.var(&name)? {
302                Some(Value::Object(object)) => object.clone(),
303                Some(_) => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, format!("invalid built-in function type {}", name).as_str()))),
304                None => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, format!("undefined built-in function {}", name).as_str()))),
305            }
306        },
307        OBJECT_MATRIX_ARRAY => {
308            let row_count = read_usize(r)?;
309            let col_count = read_usize(r)?;
310            let transpose_flag = if read_bool(r)? {
311                TransposeFlag::Transpose
312            } else {
313                TransposeFlag::NoTranspose
314            };
315            let len = checked_mul_row_count_and_col_count(row_count, col_count)?;
316            let mut xs = vec![0.0f32; len];
317            for i in 0..len {
318                xs[i] = read_f32(r)?;
319            }
320            Arc::new(Object::MatrixArray(row_count, col_count, transpose_flag, xs))
321        },
322        OBJECT_MATRIX_ROW_SLICE => {
323            let matrix_array = match read_u8(r)? {
324                MATRIX_ARRAY_OBJECT => read_object(r, env, object_tab)?,
325                MATRIX_ARRAY_INDEX => {
326                    match object_tab.object(read_usize(r)?) {
327                        Some(tmp_matrix_array) => tmp_matrix_array.clone(),
328                        None => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid object index"))),
329                    }
330                },
331                _ => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid matrix array type"))),
332            };
333            let i = read_usize(r)?;
334            Arc::new(Object::MatrixRowSlice(matrix_array, i))
335        },
336        OBJECT_ERROR => {
337            let kind = read_string(r)?;
338            let msg = read_string(r)?;
339            Arc::new(Object::Error(kind, msg))
340        },
341        _ => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid object type"))),
342    };
343    if !object_tab.add_object(object.clone()) {
344        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large index")));
345    }
346    Ok(object)
347}
348
349fn read_mut_object(r: &mut dyn Read, env: &Env, object_tab: &mut ObjectTab<Object>, mut_object_tab: &mut ObjectTab<RwLock<MutObject>>) -> Result<Arc<RwLock<MutObject>>>
350{
351    let object = match read_u8(r)? {
352        MUT_OBJECT_ARRAY => Arc::new(RwLock::new(MutObject::Array(Vec::new()))),
353        MUT_OBJECT_STRUCT => Arc::new(RwLock::new(MutObject::Struct(BTreeMap::new()))),
354        _ => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid mutable object type"))),
355    };
356    if !mut_object_tab.add_object(object.clone()) {
357        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large index")));
358    }
359    let len = read_usize(r)?;
360    {
361        let mut object_g = rw_lock_write(&object)?;
362        match &mut *object_g {
363            MutObject::Array(elems) => {
364                for _ in 0..len {
365                    elems.push(read_value(r, env, object_tab, mut_object_tab)?);
366                }
367            },
368            MutObject::Struct(fields) => {
369                for _ in 0..len {
370                    let ident = read_string(r)?;
371                    let field = read_value(r, env, object_tab, mut_object_tab)?;
372                    fields.insert(ident, field);
373                }
374            },
375        }
376    }
377    Ok(object)
378}
379
380fn read_value(r: &mut dyn Read, env: &Env, object_tab: &mut ObjectTab<Object>, mut_object_tab: &mut ObjectTab<RwLock<MutObject>>) -> Result<Value>
381{
382    match read_u8(r)? {
383        VALUE_NONE => Ok(Value::None),
384        VALUE_BOOL => Ok(Value::Bool(read_bool(r)?)),
385        VALUE_INT => Ok(Value::Int(read_i64(r)?)),
386        VALUE_FLOAT => Ok(Value::Float(read_f32(r)?)),
387        VALUE_OBJECT => Ok(Value::Object(read_object(r, env, object_tab)?)),
388        VALUE_REF => Ok(Value::Ref(read_mut_object(r, env, object_tab, mut_object_tab)?)),
389        VALUE_WEAK => Ok(Value::Weak(Arc::downgrade(&read_mut_object(r, env, object_tab, mut_object_tab)?))),
390        VALUE_WEAK_NONE => Ok(Value::Weak(Weak::new())),
391        VALUE_OBJECT_INDEX => {
392            match object_tab.object(read_usize(r)?) {
393                Some(object) => Ok(Value::Object(object.clone())),
394                None => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid object index"))),
395            }
396        },
397        VALUE_REF_INDEX => {
398            match mut_object_tab.object(read_usize(r)?) {
399                Some(object) => Ok(Value::Ref(object.clone())),
400                None => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid mutable object index"))),
401            }
402        },
403        VALUE_WEAK_INDEX => {
404            match mut_object_tab.object(read_usize(r)?) {
405                Some(object) => Ok(Value::Weak(Arc::downgrade(object))),
406                None => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid mutable object index"))),
407            }
408        },
409        _ => Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "invalid value type"))),
410    }
411}
412
413/// Reads values from the reader.
414pub fn read_values(r: &mut dyn Read, env: &Env) -> Result<Vec<Value>>
415{ 
416    let mut object_tab: ObjectTab<Object> = ObjectTab::new();
417    let mut mut_object_tab: ObjectTab<RwLock<MutObject>> = ObjectTab::new();
418    read_magic(r)?;
419    let count = read_usize(r)?;
420    let mut values: Vec<Value> = Vec::new();
421    for _ in 0..count {
422        values.push(read_value(r, env, &mut object_tab, &mut mut_object_tab)?);
423    }
424    Ok(values)
425}
426
427fn write_object(w: &mut dyn Write, object: &Arc<Object>, object_tab: &mut ObjectTab<Object>) -> Result<()>
428{
429    match &**object {
430        Object::String(s) => {
431            write_u8(w, OBJECT_STRING)?;
432            write_str(w, s.as_str())?;
433        },
434        Object::IntRange(from, to, step) => {
435            write_u8(w, OBJECT_INT_RANGE)?;
436            write_i64(w, *from)?;
437            write_i64(w, *to)?;
438            write_i64(w, *step)?;
439        },
440        Object::FloatRange(from, to, step) => {
441            write_u8(w, OBJECT_FLOAT_RANGE)?;
442            write_f32(w, *from)?;
443            write_f32(w, *to)?;
444            write_f32(w, *step)?;
445        },
446        Object::Matrix(a) => {
447            let xs = matrix_elems_and_transpose_flag(a)?.0;
448            write_u8(w, OBJECT_MATRIX)?;
449            write_usize(w, a.row_count())?;
450            write_usize(w, a.col_count())?;
451            write_bool(w, a.is_transposed())?;
452            for x in &xs {
453                write_f32(w, *x)?;
454            }
455        },
456        Object::Fun(idents, ident, _) => {
457            write_u8(w, OBJECT_FUN)?;
458            write_usize(w, idents.len())?;
459            for ident2 in idents {
460                write_str(w, ident2.as_str())?;
461            }
462            write_str(w, ident.as_str())?;
463        },
464        Object::BuiltinFun(ident, _) => {
465            write_u8(w, OBJECT_BUILTIN_FUN)?;
466            write_str(w, ident.as_str())?;
467        },
468        Object::MatrixArray(row_count, col_count, transpose_flag, xs) => {
469            write_u8(w, OBJECT_MATRIX_ARRAY)?;
470            write_usize(w, *row_count)?;
471            write_usize(w, *col_count)?;
472            write_bool(w, *transpose_flag == TransposeFlag::Transpose)?;
473            for x in xs {
474                write_f32(w, *x)?;
475            }
476        },
477        Object::MatrixRowSlice(matrix_array, i) => {
478            write_u8(w, OBJECT_MATRIX_ROW_SLICE)?;
479            match object_tab.index(matrix_array) {
480                Some(idx) => {
481                    write_u8(w, MATRIX_ARRAY_INDEX)?;
482                    write_usize(w, idx)?;
483                },
484                None => {
485                    write_u8(w, MATRIX_ARRAY_OBJECT)?;
486                    write_object(w, matrix_array, object_tab)?;
487                },
488            }
489            write_usize(w, *i)?;
490        },
491        Object::Error(kind, msg) => {
492            write_u8(w, OBJECT_ERROR)?;
493            write_str(w, kind.as_str())?;
494            write_str(w, msg.as_str())?;
495        },
496        Object::WindowId(_) => return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "can't write window identifier"))),
497    }
498    if !object_tab.add_object(object.clone()) {
499        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large index")));
500    }
501    Ok(())
502}
503
504fn write_mut_object(w: &mut dyn Write, object: &Arc<RwLock<MutObject>>, object_tab: &mut ObjectTab<Object>, mut_object_tab: &mut ObjectTab<RwLock<MutObject>>) -> Result<()>
505{
506    if !mut_object_tab.add_object(object.clone()) {
507        return Err(Error::Io(io::Error::new(ErrorKind::InvalidData, "too large index")));
508    }
509    let object_g = rw_lock_read(object)?;
510    match &*object_g {
511        MutObject::Array(elems) => {
512            write_u8(w, MUT_OBJECT_ARRAY)?;
513            write_usize(w, elems.len())?;
514            for elem in elems {
515                write_value(w, elem, object_tab, mut_object_tab)?;
516            }
517        },
518        MutObject::Struct(fields) => {
519            write_u8(w, MUT_OBJECT_STRUCT)?;
520            write_usize(w, fields.len())?;
521            for (ident, field) in fields {
522                write_str(w, ident.as_str())?;
523                write_value(w, field, object_tab, mut_object_tab)?;
524            }
525        },
526    }
527    Ok(())
528}
529
530fn write_value(w: &mut dyn Write, value: &Value, object_tab: &mut ObjectTab<Object>, mut_object_tab: &mut ObjectTab<RwLock<MutObject>>) -> Result<()>
531{
532    match value {
533        Value::None => write_u8(w, VALUE_NONE)?,
534        Value::Bool(b) => {
535            write_u8(w, VALUE_BOOL)?;
536            write_bool(w, *b)?;
537        },
538        Value::Int(n) => {
539            write_u8(w, VALUE_INT)?;
540            write_i64(w, *n)?;
541        },
542        Value::Float(n) => {
543            write_u8(w, VALUE_FLOAT)?;
544            write_f32(w, *n)?;
545        },
546        Value::Object(object) => {
547            match object_tab.index(object) {
548                Some(idx) => {
549                    write_u8(w, VALUE_OBJECT_INDEX)?;
550                    write_usize(w, idx)?;
551                },
552                None => {
553                    write_u8(w, VALUE_OBJECT)?;
554                    write_object(w, object, object_tab)?;
555                },
556            }
557        },
558        Value::Ref(object) => {
559            match mut_object_tab.index(object) {
560                Some(idx) => {
561                    write_u8(w, VALUE_REF_INDEX)?;
562                    write_usize(w, idx)?;
563                },
564                None => {
565                    write_u8(w, VALUE_REF)?;
566                    write_mut_object(w, object, object_tab, mut_object_tab)?;
567                },
568            }
569        },
570        Value::Weak(object) => {
571            match object.upgrade() {
572                Some(object) => {
573                    match mut_object_tab.index(&object) {
574                        Some(idx) => {
575                            write_u8(w, VALUE_WEAK_INDEX)?;
576                            write_usize(w, idx)?;
577                        },
578                        None => {
579                            write_u8(w, VALUE_WEAK)?;
580                            write_mut_object(w, &object, object_tab, mut_object_tab)?;
581                        },
582                    }
583                },
584                None => write_u8(w, VALUE_WEAK_NONE)?,
585            }
586        },
587    }
588    Ok(())
589}
590
591/// Writes the values to the writer.
592pub fn write_values(w: &mut dyn Write, values: &[Value]) -> Result<()>
593{ 
594    let mut object_tab: ObjectTab<Object> = ObjectTab::new();
595    let mut mut_object_tab: ObjectTab<RwLock<MutObject>> = ObjectTab::new();
596    write_magic(w)?;
597    write_usize(w, values.len())?;
598    for value in values {
599        write_value(w, value, &mut object_tab, &mut mut_object_tab)?;
600    }
601    Ok(())
602}
603
604/// Loads values from the file.
605pub fn load_values<P: AsRef<Path>>(path: P, env: &Env) -> Result<Vec<Value>>
606{
607    match File::open(path) {
608        Ok(file) => {
609            let mut r = BufReader::new(file);
610            read_values(&mut r, env)
611        },
612        Err(err) => Err(Error::Io(err)),
613    }
614}
615
616/// Saves the values to a file.
617pub fn save_values<P: AsRef<Path>>(path: P, values: &[Value]) -> Result<()>
618{
619    match File::create(path) {
620        Ok(file) => {
621            let mut w = BufWriter::new(file);
622            write_values(&mut w, values)
623        },
624        Err(err) => Err(Error::Io(err)),
625    }
626}
627
628#[cfg(test)]
629mod tests;
unlab_gpu/io.rs

unlab_gpu/
io.rs
