probl_engine/data/
mod.rs

1//! Loading a program's data (docs/data-input.md).
2//!
3//! A program's manifest (`Program::inputs`) says what it reads. The host
4//! opens each input through a [`Resolver`], which decides what a path means
5//! and whether it may be read at all. This module reads the bytes, within
6//! [`InputLimits`], and turns them into values of the declared types. The
7//! engine never touches files: it gets the values as [`Inputs`], which only
8//! [`load`] makes, and checks they're for the program it runs.
9
10mod csv;
11mod json;
12mod schema;
13mod text;
14
15pub use schema::suggest;
16
17use crate::error::{ErrorKind, RuntimeError};
18use crate::report::thousands;
19use crate::value::Value;
20use probl_sema::ir::{DataFormat, Input, Program, TypeSpec};
21use rustc_hash::FxHashMap;
22use sha2::{Digest, Sha256};
23use std::io::{ErrorKind as IoErrorKind, Read};
24use std::sync::Arc;
25use std::sync::atomic::{AtomicBool, Ordering};
26
27/// How a host gives a program its data. A program only names paths: what
28/// they refer to, and whether they may be read, is the host's policy, and a
29/// program can't widen it.
30pub trait Resolver {
31    /// What `path` refers to, as written in the program (`-` is standard
32    /// input): an identity, such as a full file name. Two inputs with the
33    /// same identity read the same bytes. Or why it may not be read.
34    fn resolve(&mut self, path: &str) -> Result<String, String>;
35    /// Open the data with this identity.
36    fn open(&mut self, identity: &str) -> Result<Box<dyn Read + Send>, String>;
37}
38
39/// Limits on loading data, for all of a program's inputs together.
40#[derive(Clone, Debug)]
41pub struct InputLimits {
42    /// Maximum magnitude size for each integer.
43    pub max_integer_bits: u64,
44    /// Cumulative allowance for large integer payloads across all inputs.
45    pub max_integer_bytes: u64,
46    /// Bytes read.
47    pub max_bytes: u64,
48    /// Values made: every number, string, record and element.
49    pub max_values: u64,
50    /// Elements of one list, map or bag.
51    pub max_collection: usize,
52    /// How deeply JSON arrays and objects may nest.
53    pub max_depth: usize,
54}
55
56impl Default for InputLimits {
57    fn default() -> InputLimits {
58        InputLimits {
59            max_integer_bits: probl_number::MAX_INTEGER_BITS,
60            max_integer_bytes: 256 * 1024 * 1024,
61            max_bytes: 64 * 1024 * 1024,
62            max_values: 10_000_000,
63            max_collection: 10_000_000,
64            max_depth: 64,
65        }
66    }
67}
68
69/// The bytes read so far, by identity. A load reads each identity once;
70/// kept between loads, snapshots keep the data the same. The REPL, which
71/// runs its whole session again after each input, keeps them until
72/// `:reload`, so that earlier bindings don't change under it.
73#[derive(Debug, Default)]
74pub struct Snapshots {
75    bytes: FxHashMap<String, Arc<[u8]>>,
76}
77
78impl Snapshots {
79    /// Forget everything read, so the next load reads it again.
80    pub fn clear(&mut self) {
81        self.bytes.clear();
82    }
83}
84
85/// What was read for a program: the value of each input, and where they
86/// came from. Only [`load`] makes these, and the engine checks that they're
87/// for the program it runs.
88#[derive(Debug)]
89pub struct Inputs {
90    max_integer_bits: u64,
91    max_string_bytes: usize,
92    manifest: Vec<Input>,
93    values: Vec<Value>,
94    sources: Vec<SourceInfo>,
95}
96
97impl Inputs {
98    pub(crate) fn max_string_bytes(&self) -> usize {
99        self.max_string_bytes
100    }
101    pub(crate) fn max_integer_bits(&self) -> u64 {
102        self.max_integer_bits
103    }
104
105    /// The value of each input, in the order of the program's manifest.
106    pub fn values(&self) -> &[Value] {
107        &self.values
108    }
109
110    /// Each file (or standard input) read, once.
111    pub fn sources(&self) -> &[SourceInfo] {
112        &self.sources
113    }
114
115    /// Whether these are the inputs of `program`.
116    pub fn fit(&self, program: &Program) -> bool {
117        self.manifest == program.inputs
118    }
119}
120
121/// A file, or standard input, that was read.
122#[derive(Clone, Debug, PartialEq)]
123pub struct SourceInfo {
124    pub identity: String,
125    pub bytes: u64,
126    /// The SHA-256 of the bytes, in hex: what the results were computed
127    /// from.
128    pub sha256: String,
129}
130
131/// Read every input of `program`, through `resolver`. Errors are about the
132/// input's `read(…)` call, and say where in the data they are.
133pub fn load(
134    program: &Program,
135    resolver: &mut dyn Resolver,
136    snapshots: &mut Snapshots,
137    limits: &InputLimits,
138    cancel: Option<&AtomicBool>,
139) -> Result<Inputs, RuntimeError> {
140    let mut cx = Cx {
141        program,
142        budget: Budget {
143            values_left: limits.max_values,
144            integer_bytes_left: limits.max_integer_bytes,
145            max_integer_bits_seen: 0,
146            max_string_bytes_seen: 0,
147            limits,
148            cancel,
149            ticks: 0,
150        },
151        names: FxHashMap::default(),
152        depth: 0,
153        path: Vec::new(),
154        problem: None,
155    };
156    // Each identity once, in the order they're first read.
157    let mut read: Vec<(String, Arc<[u8]>)> = Vec::new();
158    let mut total: u64 = 0;
159    let mut values = Vec::with_capacity(program.inputs.len());
160    for input in &program.inputs {
161        let identity = resolver
162            .resolve(&input.path)
163            .map_err(|why| failure(input, Problem::new(format!("can't read it: {why}"))))?;
164        let bytes = match read.iter().find(|(id, _)| *id == identity) {
165            Some((_, bytes)) => bytes.clone(),
166            None => {
167                let bytes = match snapshots.bytes.get(&identity) {
168                    Some(bytes) => bytes.clone(),
169                    None => {
170                        let reader = resolver
171                            .open(&identity)
172                            .map_err(|why| failure(input, Problem::new(format!("can't read it: {why}"))))?;
173                        let left = limits.max_bytes.saturating_sub(total);
174                        let bytes: Arc<[u8]> = read_bounded(reader, left, limits.max_bytes, cancel)
175                            .map_err(|p| failure(input, p))?
176                            .into();
177                        snapshots.bytes.insert(identity.clone(), bytes.clone());
178                        bytes
179                    }
180                };
181                total += bytes.len() as u64;
182                if total > limits.max_bytes {
183                    return Err(failure(input, too_many_bytes(limits.max_bytes)));
184                }
185                read.push((identity, bytes.clone()));
186                bytes
187            }
188        };
189        let value = decode(&bytes, input, &mut cx).map_err(|p| failure(input, p))?;
190        values.push(value);
191    }
192    let sources = read
193        .iter()
194        .map(|(identity, bytes)| SourceInfo {
195            identity: identity.clone(),
196            bytes: bytes.len() as u64,
197            sha256: Sha256::digest(bytes).iter().map(|b| format!("{b:02x}")).collect(),
198        })
199        .collect();
200    Ok(Inputs {
201        max_integer_bits: cx.budget.max_integer_bits_seen,
202        max_string_bytes: cx.budget.max_string_bytes_seen,
203        manifest: program.inputs.clone(),
204        values,
205        sources,
206    })
207}
208
209/// Read all of `reader`, if it has at most `left` bytes, checking for
210/// cancellation as it goes.
211fn read_bounded(
212    mut reader: Box<dyn Read + Send>,
213    left: u64,
214    max: u64,
215    cancel: Option<&AtomicBool>,
216) -> Result<Vec<u8>, Problem> {
217    let mut bytes = Vec::new();
218    let mut chunk = vec![0u8; 64 * 1024];
219    loop {
220        if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
221            return Err(Problem::limit("reading it was cancelled"));
222        }
223        let n = match reader.read(&mut chunk) {
224            Ok(0) => return Ok(bytes),
225            Ok(n) => n,
226            Err(e) if e.kind() == IoErrorKind::Interrupted => continue,
227            Err(_) if cancel.is_some_and(|c| c.load(Ordering::Relaxed)) => {
228                return Err(Problem::limit("reading it was cancelled"));
229            }
230            Err(e) => return Err(Problem::new(format!("reading it failed: {e}"))),
231        };
232        if (bytes.len() + n) as u64 > left {
233            return Err(too_many_bytes(max));
234        }
235        bytes.extend_from_slice(&chunk[..n]);
236    }
237}
238
239fn too_many_bytes(max: u64) -> Problem {
240    Problem::limit(format!(
241        "the data is more than {} MiB, the limit for all the data together",
242        max.div_ceil(1024 * 1024)
243    ))
244    .help("raise the limit with `--max-input`")
245}
246
247/// Read one input's bytes as its declared type.
248fn decode(bytes: &[u8], input: &Input, cx: &mut Cx) -> Result<Value, Problem> {
249    let bytes = bytes.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(bytes);
250    cx.depth = 0;
251    cx.path.clear();
252    cx.problem = None;
253    match input.format {
254        DataFormat::Csv => csv::read(bytes, &input.ty, cx),
255        DataFormat::Json => json::read(bytes, &input.ty, cx),
256        DataFormat::Lines => lines(bytes, &input.ty, cx),
257    }
258}
259
260/// One value per non-blank line.
261fn lines(bytes: &[u8], ty: &TypeSpec, cx: &mut Cx) -> Result<Value, Problem> {
262    let TypeSpec::List(item) = ty else {
263        return Err(Problem::new("lines read as a list"));
264    };
265    let text = std::str::from_utf8(bytes).map_err(|e| {
266        let line = bytes[..e.valid_up_to()].iter().filter(|&&b| b == b'\n').count() + 1;
267        Problem::new("isn't text: it isn't valid UTF-8").at(format!("line {line}"))
268    })?;
269    let mut items = Vec::new();
270    for (i, line) in text.lines().enumerate() {
271        if line.trim().is_empty() {
272            continue;
273        }
274        cx.budget.collection(items.len() + 1)?;
275        cx.budget.value()?;
276        let v = text::plain(line, item, cx.program, &mut cx.budget).map_err(|p| p.at(format!("line {}", i + 1)))?;
277        items.push(v);
278    }
279    Ok(Value::list(items))
280}
281
282/// What decoding needs: the program's types, the budget, and where it is.
283pub(crate) struct Cx<'a> {
284    program: &'a Program,
285    budget: Budget<'a>,
286    /// Record type and field names, made once.
287    names: FxHashMap<String, Arc<str>>,
288    /// JSON nesting.
289    depth: usize,
290    /// Where in a JSON document: keys and indices.
291    path: Vec<PathPart>,
292    /// The first problem found inside the JSON parser, which only passes on
293    /// its own errors.
294    problem: Option<Problem>,
295}
296
297impl Cx<'_> {
298    fn name(&mut self, name: &str) -> Arc<str> {
299        if let Some(n) = self.names.get(name) {
300            return n.clone();
301        }
302        let n: Arc<str> = Arc::from(name);
303        self.names.insert(name.to_string(), n.clone());
304        n
305    }
306}
307
308enum PathPart {
309    Key(String),
310    Index(usize),
311}
312
313/// Limits being spent while decoding.
314struct Budget<'a> {
315    values_left: u64,
316    integer_bytes_left: u64,
317    max_integer_bits_seen: u64,
318    max_string_bytes_seen: usize,
319    limits: &'a InputLimits,
320    cancel: Option<&'a AtomicBool>,
321    ticks: u32,
322}
323
324impl Budget<'_> {
325    fn integer(&mut self, bits: u64) -> Result<(), Problem> {
326        let limit = self.limits.max_integer_bits.min(probl_number::MAX_INTEGER_BITS);
327        if bits > limit {
328            return Err(Problem::limit(format!(
329                "integer size exceeds the limit of {limit} bits"
330            )));
331        }
332        self.max_integer_bits_seen = self.max_integer_bits_seen.max(bits);
333        if bits > 63 {
334            let bytes = bits.div_ceil(64) * 8 + 48;
335            self.integer_bytes_left = self
336                .integer_bytes_left
337                .checked_sub(bytes)
338                .ok_or_else(|| Problem::limit("the data used up its large integer memory allowance"))?;
339        }
340        Ok(())
341    }
342
343    /// Count one more value, before making it.
344    fn value(&mut self) -> Result<(), Problem> {
345        if self.values_left == 0 {
346            return Err(Problem::limit(format!(
347                "the data has more than {} values, the limit for all the data together",
348                thousands(self.limits.max_values.min(i64::MAX as u64) as i64)
349            )));
350        }
351        self.values_left -= 1;
352        self.ticks = self.ticks.wrapping_add(1);
353        if self.ticks % 4096 == 0 && self.cancel.is_some_and(|c| c.load(Ordering::Relaxed)) {
354            return Err(Problem::limit("reading it was cancelled"));
355        }
356        Ok(())
357    }
358
359    /// Check that a collection may have `n` elements.
360    fn collection(&self, n: usize) -> Result<(), Problem> {
361        if n > self.limits.max_collection {
362            return Err(Problem::limit(format!(
363                "a list, map or bag has more than {} elements",
364                thousands(self.limits.max_collection as i64)
365            )));
366        }
367        Ok(())
368    }
369}
370
371/// What's wrong with the data, and where.
372#[derive(Clone, Debug)]
373pub(crate) struct Problem {
374    message: String,
375    help: Option<String>,
376    notes: Vec<String>,
377    /// Where in the data: "line 12, column `signups`".
378    at: Option<String>,
379    limit: bool,
380}
381
382impl Problem {
383    fn new(message: impl Into<String>) -> Problem {
384        Problem {
385            message: message.into(),
386            help: None,
387            notes: Vec::new(),
388            at: None,
389            limit: false,
390        }
391    }
392
393    fn limit(message: impl Into<String>) -> Problem {
394        Problem {
395            limit: true,
396            ..Problem::new(message)
397        }
398    }
399
400    fn help(mut self, help: impl Into<String>) -> Problem {
401        self.help = Some(help.into());
402        self
403    }
404
405    fn note(mut self, note: impl Into<String>) -> Problem {
406        self.notes.push(note.into());
407        self
408    }
409
410    /// Say where, unless it's already said.
411    fn at(mut self, at: impl Into<String>) -> Problem {
412        if self.at.is_none() {
413            self.at = Some(at.into());
414        }
415        self
416    }
417}
418
419/// A problem with an input, as an error at its `read(…)` call.
420fn failure(input: &Input, problem: Problem) -> RuntimeError {
421    let place = match &problem.at {
422        Some(at) => format!("{}, {at}", source_name(input)),
423        None => source_name(input),
424    };
425    let mut error = RuntimeError::new(input.span, format!("{place}: {}", problem.message));
426    error.notes = problem.notes;
427    error.help = problem.help;
428    if problem.limit {
429        error.kind = ErrorKind::Limit;
430    }
431    error
432}
433
434fn source_name(input: &Input) -> String {
435    match input.path.as_str() {
436        "-" => "standard input".to_string(),
437        path => path.to_string(),
438    }
439}
440
441/// A short, quoted rendering of text from the data, for messages.
442fn quoted(text: &str) -> String {
443    const MAX: usize = 40;
444    match text.char_indices().nth(MAX) {
445        Some((cut, _)) => format!("`{}…`", &text[..cut]),
446        None => format!("`{text}`"),
447    }
448}
449
450#[cfg(test)]
451mod tests;
probl_engine/data/mod.rs

probl_engine/data/
mod.rs
