use crate::kernel::domain::Domain;
use crate::kernel::expr::{BigFloat, BigInt, BigRat, ExprData, ExprId, PredicateKind};
use crate::kernel::pool::ExprPool;
use std::fs::{self, File};
use std::io::{self, BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
const MAGIC: &[u8; 4] = b"ALKP";
const POOL_FORMAT_V1: u32 = 1;
const POOL_FORMAT_V2: u32 = 2;
const POOL_FORMAT_V3: u32 = 3;
const POOL_FORMAT_V4: u32 = 4;
const POOL_FORMAT_V5: u32 = 5;
const POOL_FORMAT_WRITE: u32 = POOL_FORMAT_V5;
#[derive(Debug)]
pub enum IoError {
Io(io::Error),
BadMagic,
UnsupportedVersion(u32),
Truncated,
BadUtf8,
BadDomain(u8),
BadTag(u8),
BadPredicateKind(u8),
BadNumeric(String),
}
impl std::fmt::Display for IoError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
IoError::Io(e) => write!(f, "io error: {e}"),
IoError::BadMagic => write!(f, "not an alkahest pool file (bad magic)"),
IoError::UnsupportedVersion(v) => {
write!(
f,
"unsupported pool file version {v}; run `alkahest migrate-pool`"
)
}
IoError::Truncated => write!(f, "pool file truncated or incomplete"),
IoError::BadUtf8 => write!(f, "pool file contains invalid UTF-8"),
IoError::BadDomain(b) => write!(f, "pool file has unknown domain tag {b}"),
IoError::BadTag(b) => write!(f, "pool file has unknown node tag {b}"),
IoError::BadPredicateKind(b) => {
write!(f, "pool file has unknown predicate kind {b}")
}
IoError::BadNumeric(s) => write!(f, "pool file has invalid numeric: {s}"),
}
}
}
impl std::error::Error for IoError {}
impl From<io::Error> for IoError {
fn from(e: io::Error) -> Self {
IoError::Io(e)
}
}
impl crate::errors::AlkahestError for IoError {
fn code(&self) -> &'static str {
match self {
IoError::Io(_) => "E-IO-001",
IoError::BadMagic => "E-IO-002",
IoError::UnsupportedVersion(_) => "E-IO-003",
IoError::Truncated => "E-IO-004",
IoError::BadUtf8 => "E-IO-005",
IoError::BadDomain(_) => "E-IO-006",
IoError::BadTag(_) => "E-IO-007",
IoError::BadPredicateKind(_) => "E-IO-008",
IoError::BadNumeric(_) => "E-IO-009",
}
}
fn remediation(&self) -> Option<&'static str> {
match self {
IoError::BadMagic => Some(
"file is not an alkahest pool; check the path or regenerate with ExprPool::checkpoint()",
),
IoError::UnsupportedVersion(_) => Some(
"run the `alkahest migrate-pool` CLI to upgrade the file, or regenerate from source",
),
IoError::Truncated => Some(
"file was truncated (likely a crash during checkpoint); rerun from source and checkpoint again",
),
_ => None,
}
}
}
#[deprecated(since = "2.0.0", note = "renamed to IoError with E-IO-* codes")]
pub type PoolPersistError = IoError;
fn write_u8(w: &mut impl Write, v: u8) -> io::Result<()> {
w.write_all(&[v])
}
fn write_u32(w: &mut impl Write, v: u32) -> io::Result<()> {
w.write_all(&v.to_le_bytes())
}
fn write_u64(w: &mut impl Write, v: u64) -> io::Result<()> {
w.write_all(&v.to_le_bytes())
}
fn write_str(w: &mut impl Write, s: &str) -> io::Result<()> {
let bytes = s.as_bytes();
write_u32(w, bytes.len() as u32)?;
w.write_all(bytes)
}
fn write_ids(w: &mut impl Write, ids: &[ExprId]) -> io::Result<()> {
write_u32(w, ids.len() as u32)?;
for id in ids {
write_u32(w, id.0)?;
}
Ok(())
}
const MAX_FILE_FLOAT_PREC: u32 = 1 << 24;
const MIN_NODE_BYTES: u64 = 5;
struct Bounded<R> {
inner: R,
remaining: u64,
}
impl<R: Read> Bounded<R> {
fn new(inner: R, len: u64) -> Self {
Bounded {
inner,
remaining: len,
}
}
fn ensure(&self, count: u64, size: u64) -> Result<(), IoError> {
match count.checked_mul(size) {
Some(n) if n <= self.remaining => Ok(()),
_ => Err(IoError::Truncated),
}
}
fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), IoError> {
self.ensure(buf.len() as u64, 1)?;
self.inner.read_exact(buf).map_err(|_| IoError::Truncated)?;
self.remaining -= buf.len() as u64;
Ok(())
}
}
fn read_u8(r: &mut Bounded<impl Read>) -> Result<u8, IoError> {
let mut b = [0u8; 1];
r.read_exact(&mut b)?;
Ok(b[0])
}
fn read_u32(r: &mut Bounded<impl Read>) -> Result<u32, IoError> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(u32::from_le_bytes(b))
}
fn read_u64(r: &mut Bounded<impl Read>) -> Result<u64, IoError> {
let mut b = [0u8; 8];
r.read_exact(&mut b)?;
Ok(u64::from_le_bytes(b))
}
fn read_str(r: &mut Bounded<impl Read>) -> Result<String, IoError> {
let len = u64::from(read_u32(r)?);
r.ensure(len, 1)?;
let mut buf = vec![0u8; len as usize];
r.read_exact(&mut buf)?;
String::from_utf8(buf).map_err(|_| IoError::BadUtf8)
}
fn read_ids(r: &mut Bounded<impl Read>) -> Result<Vec<ExprId>, IoError> {
let arity = u64::from(read_u32(r)?);
r.ensure(arity, 4)?;
let mut out = Vec::with_capacity(arity as usize);
for _ in 0..arity {
out.push(ExprId(read_u32(r)?));
}
Ok(out)
}
fn domain_to_u8(d: &Domain) -> u8 {
match d {
Domain::Real => 0,
Domain::Complex => 1,
Domain::Integer => 2,
Domain::Positive => 3,
Domain::NonNegative => 4,
Domain::NonZero => 5,
}
}
fn u8_to_domain(b: u8) -> Result<Domain, IoError> {
match b {
0 => Ok(Domain::Real),
1 => Ok(Domain::Complex),
2 => Ok(Domain::Integer),
3 => Ok(Domain::Positive),
4 => Ok(Domain::NonNegative),
5 => Ok(Domain::NonZero),
b => Err(IoError::BadDomain(b)),
}
}
fn pred_to_u8(k: &PredicateKind) -> u8 {
match k {
PredicateKind::Eq => 0,
PredicateKind::Ne => 1,
PredicateKind::Lt => 2,
PredicateKind::Le => 3,
PredicateKind::Gt => 4,
PredicateKind::Ge => 5,
PredicateKind::And => 6,
PredicateKind::Or => 7,
PredicateKind::Not => 8,
PredicateKind::True => 9,
PredicateKind::False => 10,
}
}
fn u8_to_pred(b: u8) -> Result<PredicateKind, IoError> {
match b {
0 => Ok(PredicateKind::Eq),
1 => Ok(PredicateKind::Ne),
2 => Ok(PredicateKind::Lt),
3 => Ok(PredicateKind::Le),
4 => Ok(PredicateKind::Gt),
5 => Ok(PredicateKind::Ge),
6 => Ok(PredicateKind::And),
7 => Ok(PredicateKind::Or),
8 => Ok(PredicateKind::Not),
9 => Ok(PredicateKind::True),
10 => Ok(PredicateKind::False),
b => Err(IoError::BadPredicateKind(b)),
}
}
fn write_node(w: &mut impl Write, node: &ExprData) -> io::Result<()> {
match node {
ExprData::Symbol {
name,
domain,
commutative,
} => {
write_u8(w, 0)?;
write_u8(w, domain_to_u8(domain))?;
write_u8(w, u8::from(*commutative))?;
write_str(w, name)
}
ExprData::Integer(BigInt(n)) => {
write_u8(w, 1)?;
write_str(w, &n.to_string())
}
ExprData::Rational(BigRat(r)) => {
write_u8(w, 2)?;
write_str(w, &r.numer().to_string())?;
write_str(w, &r.denom().to_string())
}
ExprData::Float(BigFloat { inner, prec }) => {
write_u8(w, 3)?;
write_u32(w, *prec)?;
write_str(w, &inner.to_string_radix(16, None))
}
ExprData::Add(children) => {
write_u8(w, 4)?;
write_ids(w, children)
}
ExprData::Mul(children) => {
write_u8(w, 5)?;
write_ids(w, children)
}
ExprData::Pow { base, exp } => {
write_u8(w, 6)?;
write_u32(w, base.0)?;
write_u32(w, exp.0)
}
ExprData::Func { name, args } => {
write_u8(w, 7)?;
write_str(w, name)?;
write_ids(w, args)
}
ExprData::Piecewise { branches, default } => {
write_u8(w, 8)?;
write_u32(w, branches.len() as u32)?;
for (c, v) in branches {
write_u32(w, c.0)?;
write_u32(w, v.0)?;
}
write_u32(w, default.0)
}
ExprData::Predicate { kind, args } => {
write_u8(w, 9)?;
write_u8(w, pred_to_u8(kind))?;
write_ids(w, args)
}
ExprData::Forall { var, body } => {
write_u8(w, 10)?;
write_u32(w, var.0)?;
write_u32(w, body.0)
}
ExprData::Exists { var, body } => {
write_u8(w, 11)?;
write_u32(w, var.0)?;
write_u32(w, body.0)
}
ExprData::BigO(inner) => {
write_u8(w, 12)?;
write_u32(w, inner.0)
}
ExprData::RootSum { poly, var, body } => {
write_u8(w, 13)?;
write_u32(w, poly.0)?;
write_u32(w, var.0)?;
write_u32(w, body.0)
}
}
}
fn read_node(r: &mut Bounded<impl Read>, format_version: u32) -> Result<ExprData, IoError> {
let tag = read_u8(r)?;
match tag {
0 => {
let domain = u8_to_domain(read_u8(r)?)?;
let commutative = if format_version >= POOL_FORMAT_V4 {
read_u8(r)? != 0
} else {
true
};
let name = read_str(r)?;
Ok(ExprData::Symbol {
name,
domain,
commutative,
})
}
1 => {
let s = read_str(r)?;
let n: rug::Integer = s
.parse()
.map_err(|_| IoError::BadNumeric(format!("integer: {s}")))?;
Ok(ExprData::Integer(BigInt(n)))
}
2 => {
let nstr = read_str(r)?;
let dstr = read_str(r)?;
let n: rug::Integer = nstr
.parse()
.map_err(|_| IoError::BadNumeric(format!("numer: {nstr}")))?;
let d: rug::Integer = dstr
.parse()
.map_err(|_| IoError::BadNumeric(format!("denom: {dstr}")))?;
if d == 0 {
return Err(IoError::BadNumeric(format!("rational {nstr}/0")));
}
Ok(ExprData::Rational(BigRat(rug::Rational::from((n, d)))))
}
3 => {
let prec = read_u32(r)?;
if prec == 0 || prec > MAX_FILE_FLOAT_PREC {
return Err(IoError::BadNumeric(format!(
"float precision {prec} bits is outside 1..={MAX_FILE_FLOAT_PREC}"
)));
}
let s = read_str(r)?;
let f = rug::Float::parse_radix(&s, 16)
.map_err(|_| IoError::BadNumeric(format!("float: {s}")))?;
let inner = rug::Float::with_val(prec, f);
Ok(ExprData::Float(BigFloat { inner, prec }))
}
4 => Ok(ExprData::Add(read_ids(r)?)),
5 => Ok(ExprData::Mul(read_ids(r)?)),
6 => {
let base = ExprId(read_u32(r)?);
let exp = ExprId(read_u32(r)?);
Ok(ExprData::Pow { base, exp })
}
7 => {
let name = read_str(r)?;
let args = read_ids(r)?;
Ok(ExprData::Func { name, args })
}
8 => {
let n = u64::from(read_u32(r)?);
r.ensure(n, 8)?;
let mut branches = Vec::with_capacity(n as usize);
for _ in 0..n {
let c = ExprId(read_u32(r)?);
let v = ExprId(read_u32(r)?);
branches.push((c, v));
}
let default = ExprId(read_u32(r)?);
Ok(ExprData::Piecewise { branches, default })
}
9 => {
let kind = u8_to_pred(read_u8(r)?)?;
let args = read_ids(r)?;
Ok(ExprData::Predicate { kind, args })
}
10 => {
if format_version < POOL_FORMAT_V2 {
return Err(IoError::BadTag(10));
}
let var = ExprId(read_u32(r)?);
let body = ExprId(read_u32(r)?);
Ok(ExprData::Forall { var, body })
}
11 => {
if format_version < POOL_FORMAT_V2 {
return Err(IoError::BadTag(11));
}
let var = ExprId(read_u32(r)?);
let body = ExprId(read_u32(r)?);
Ok(ExprData::Exists { var, body })
}
12 => {
if format_version < POOL_FORMAT_V3 {
return Err(IoError::BadTag(12));
}
let inner = ExprId(read_u32(r)?);
Ok(ExprData::BigO(inner))
}
13 => {
if format_version < POOL_FORMAT_V5 {
return Err(IoError::BadTag(13));
}
let poly = ExprId(read_u32(r)?);
let var = ExprId(read_u32(r)?);
let body = ExprId(read_u32(r)?);
Ok(ExprData::RootSum { poly, var, body })
}
b => Err(IoError::BadTag(b)),
}
}
pub fn save_to(pool: &ExprPool, path: impl AsRef<Path>) -> Result<(), IoError> {
let path = path.as_ref();
let tmp: PathBuf = {
let mut p = path.to_path_buf();
let mut name = p
.file_name()
.map(|s| s.to_os_string())
.unwrap_or_else(|| std::ffi::OsString::from("pool"));
name.push(".tmp");
p.set_file_name(name);
p
};
{
let f = File::create(&tmp)?;
let mut w = BufWriter::new(f);
w.write_all(MAGIC)?;
write_u32(&mut w, POOL_FORMAT_WRITE)?;
write_u32(&mut w, 0u32)?;
let count = pool.len();
write_u64(&mut w, count as u64)?;
for i in 0..count {
let data = pool.get(ExprId(i as u32));
write_node(&mut w, &data)?;
}
w.flush()?;
w.get_ref().sync_all()?;
}
fs::rename(&tmp, path)?;
Ok(())
}
pub fn load_from(path: impl AsRef<Path>) -> Result<Option<ExprPool>, IoError> {
let path = path.as_ref();
if !path.exists() {
return Ok(None);
}
let f = File::open(path)?;
let len = f.metadata()?.len();
let mut r = Bounded::new(BufReader::new(f), len);
let mut magic = [0u8; 4];
r.read_exact(&mut magic)?;
if &magic != MAGIC {
return Err(IoError::BadMagic);
}
let pool = ExprPool::new();
read_body(&mut r, &pool)?;
Ok(Some(pool))
}
fn read_body(r: &mut Bounded<impl Read>, pool: &ExprPool) -> Result<Vec<ExprId>, IoError> {
let version = read_u32(r)?;
if version != POOL_FORMAT_V1
&& version != POOL_FORMAT_V2
&& version != POOL_FORMAT_V3
&& version != POOL_FORMAT_V4
&& version != POOL_FORMAT_V5
{
return Err(IoError::UnsupportedVersion(version));
}
let _flags = read_u32(r)?;
let count = read_u64(r)?;
r.ensure(count, MIN_NODE_BYTES)?;
if count > u64::from(u32::MAX) {
return Err(IoError::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!("pool file declares {count} nodes, more than an ExprId can address"),
)));
}
let count = count as usize;
let mut ids: Vec<ExprId> = Vec::new();
for index in 0..count {
let data = read_node(r, version)?;
let id = match remap_children(data, &ids, index)? {
ExprData::Add(args) => pool.add(args),
ExprData::Mul(args) => pool.mul(args),
other => pool.intern(other),
};
ids.push(id);
}
Ok(ids)
}
fn remap_children(data: ExprData, ids: &[ExprId], index: usize) -> Result<ExprData, IoError> {
map_children(data, |c| {
ids.get(c.0 as usize).copied().ok_or_else(|| {
IoError::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"pool file node {index} refers to node {}, which is not an earlier node",
c.0
),
))
})
})
}
fn map_children(
data: ExprData,
mut f: impl FnMut(ExprId) -> Result<ExprId, IoError>,
) -> Result<ExprData, IoError> {
fn map_all(
v: Vec<ExprId>,
f: &mut impl FnMut(ExprId) -> Result<ExprId, IoError>,
) -> Result<Vec<ExprId>, IoError> {
v.into_iter().map(f).collect()
}
Ok(match data {
ExprData::Symbol { .. }
| ExprData::Integer(_)
| ExprData::Rational(_)
| ExprData::Float(_) => data,
ExprData::Add(args) => ExprData::Add(map_all(args, &mut f)?),
ExprData::Mul(args) => ExprData::Mul(map_all(args, &mut f)?),
ExprData::Pow { base, exp } => ExprData::Pow {
base: f(base)?,
exp: f(exp)?,
},
ExprData::Func { name, args } => ExprData::Func {
name,
args: map_all(args, &mut f)?,
},
ExprData::Piecewise { branches, default } => ExprData::Piecewise {
branches: branches
.into_iter()
.map(|(c, v)| Ok((f(c)?, f(v)?)))
.collect::<Result<_, IoError>>()?,
default: f(default)?,
},
ExprData::Predicate { kind, args } => ExprData::Predicate {
kind,
args: map_all(args, &mut f)?,
},
ExprData::Forall { var, body } => ExprData::Forall {
var: f(var)?,
body: f(body)?,
},
ExprData::Exists { var, body } => ExprData::Exists {
var: f(var)?,
body: f(body)?,
},
ExprData::BigO(inner) => ExprData::BigO(f(inner)?),
ExprData::RootSum { poly, var, body } => ExprData::RootSum {
poly: f(poly)?,
var: f(var)?,
body: f(body)?,
},
})
}
const DAG_MAGIC: &[u8; 4] = b"ALKD";
pub fn dag_to_bytes(pool: &ExprPool, roots: &[ExprId]) -> Vec<u8> {
use std::collections::HashMap;
let mut index: HashMap<ExprId, u32> = HashMap::new();
let mut order: Vec<ExprData> = Vec::new();
let mut stack: Vec<(ExprId, bool)> = roots.iter().rev().map(|&r| (r, false)).collect();
while let Some((id, expanded)) = stack.pop() {
if index.contains_key(&id) {
continue;
}
let data = pool.get(id);
if expanded {
let local = map_children(data, |c| Ok(ExprId(index[&c])))
.expect("every child was emitted before its parent");
index.insert(id, order.len() as u32);
order.push(local);
continue;
}
stack.push((id, true));
let mut children = Vec::new();
let _ = map_children(data, |c| {
children.push(c);
Ok(c)
});
for c in children.into_iter().rev() {
if !index.contains_key(&c) {
stack.push((c, false));
}
}
}
let mut out: Vec<u8> = Vec::new();
out.extend_from_slice(DAG_MAGIC);
out.extend_from_slice(&POOL_FORMAT_WRITE.to_le_bytes());
out.extend_from_slice(&0u32.to_le_bytes());
out.extend_from_slice(&(order.len() as u64).to_le_bytes());
for node in &order {
write_node(&mut out, node).expect("writing to a Vec cannot fail");
}
out.extend_from_slice(&(roots.len() as u32).to_le_bytes());
for r in roots {
out.extend_from_slice(&index[r].to_le_bytes());
}
out
}
pub fn dag_from_bytes(pool: &ExprPool, bytes: &[u8]) -> Result<Vec<ExprId>, IoError> {
let mut r = Bounded::new(bytes, bytes.len() as u64);
let mut magic = [0u8; 4];
r.read_exact(&mut magic)?;
if &magic != DAG_MAGIC {
return Err(IoError::BadMagic);
}
let ids = read_body(&mut r, pool)?;
let n_roots = u64::from(read_u32(&mut r)?);
r.ensure(n_roots, 4)?;
let mut roots = Vec::with_capacity(n_roots as usize);
for _ in 0..n_roots {
let i = read_u32(&mut r)?;
let id = ids.get(i as usize).copied().ok_or_else(|| {
IoError::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"expression DAG root {i} is not one of its {} nodes",
ids.len()
),
))
})?;
roots.push(id);
}
Ok(roots)
}
pub fn open_persistent(path: impl AsRef<Path>) -> Result<ExprPool, IoError> {
match load_from(path)? {
Some(p) => Ok(p),
None => Ok(ExprPool::new()),
}
}
impl ExprPool {
pub fn checkpoint(&self, path: impl AsRef<Path>) -> Result<(), IoError> {
save_to(self, path)
}
pub fn open_persistent(path: impl AsRef<Path>) -> Result<Self, IoError> {
open_persistent(path)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::kernel::{Domain, ExprData};
fn tempfile() -> PathBuf {
static SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let mut p = std::env::temp_dir();
p.push(format!(
"alkahest_pool_{}_{}_{}.akp",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos(),
SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
));
p
}
#[test]
fn round_trip_small_pool() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let y = p.symbol("y", Domain::Positive);
let two = p.integer(2_i32);
let three_halves = p.rational(3, 2);
let f = p.float(1.5_f64, 53);
let xp = p.pow(x, two);
let fn_node = p.func("sin", vec![xp]);
let _sum = p.add(vec![fn_node, y, three_halves, f]);
let path = tempfile();
p.checkpoint(&path).unwrap();
let q = ExprPool::open_persistent(&path).unwrap();
assert_eq!(q.len(), p.len(), "node count must match");
for i in 0..p.len() {
let id = ExprId(i as u32);
assert_eq!(p.get(id), q.get(id), "node {i} mismatch after round-trip");
}
let q_x = q.symbol("x", Domain::Real);
assert_eq!(q_x, x, "symbol id drifted across checkpoint");
let q_two = q.integer(2_i32);
assert_eq!(q_two, two);
let _ = fs::remove_file(&path);
}
#[test]
fn round_trip_root_sum() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let t = p.symbol("t", Domain::Complex);
let poly = p.add(vec![p.pow(t, p.integer(2_i32)), p.integer(1_i32)]);
let body = p.mul(vec![t, p.func("log", vec![p.add(vec![x, t])])]);
let rs = p.root_sum(poly, t, body);
assert!(matches!(p.get(rs), ExprData::RootSum { .. }));
let path = tempfile();
p.checkpoint(&path).unwrap();
let q = ExprPool::open_persistent(&path).unwrap();
assert_eq!(q.len(), p.len());
for i in 0..p.len() {
let id = ExprId(i as u32);
assert_eq!(p.get(id), q.get(id), "node {i} mismatch after round-trip");
}
let _ = fs::remove_file(&path);
}
#[test]
fn bad_magic_rejected() {
let path = tempfile();
std::fs::write(&path, b"nope1234").unwrap();
match load_from(&path) {
Err(IoError::BadMagic) => {}
other => panic!("expected BadMagic, got {:?}", other.err()),
}
let _ = fs::remove_file(&path);
}
#[test]
fn missing_file_returns_fresh() {
let path = tempfile();
assert!(!path.exists());
let p = ExprPool::open_persistent(&path).unwrap();
assert_eq!(p.len(), 0);
}
#[test]
fn predicate_and_piecewise_round_trip() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let zero = p.integer(0_i32);
let one = p.integer(1_i32);
let neg_one = p.integer(-1_i32);
let cond = p.intern(ExprData::Predicate {
kind: PredicateKind::Gt,
args: vec![x, zero],
});
let pc = p.intern(ExprData::Piecewise {
branches: vec![(cond, one)],
default: neg_one,
});
let path = tempfile();
p.checkpoint(&path).unwrap();
let q = ExprPool::open_persistent(&path).unwrap();
assert_eq!(p.get(pc), q.get(pc));
let _ = fs::remove_file(&path);
}
fn write_raw(path: &Path, nodes: &[ExprData]) {
let mut w = BufWriter::new(File::create(path).unwrap());
w.write_all(MAGIC).unwrap();
write_u32(&mut w, POOL_FORMAT_WRITE).unwrap();
write_u32(&mut w, 0).unwrap();
write_u64(&mut w, nodes.len() as u64).unwrap();
for n in nodes {
write_node(&mut w, n).unwrap();
}
w.flush().unwrap();
}
#[test]
fn non_canonical_nodes_collapse_without_drift() {
let x = ExprData::Symbol {
name: "x".into(),
domain: Domain::Real,
commutative: true,
};
let y = ExprData::Symbol {
name: "y".into(),
domain: Domain::Real,
commutative: true,
};
let nodes = vec![
x, ExprData::Rational(BigRat(rug::Rational::from((4, 2)))), ExprData::Integer(BigInt(rug::Integer::from(2))), ExprData::Add(vec![ExprId(0)]), y, ExprData::Add(vec![ExprId(4), ExprId(3)]), ExprData::Pow {
base: ExprId(5),
exp: ExprId(2),
}, ExprData::Mul(vec![]), ExprData::Pow {
base: ExprId(4),
exp: ExprId(7),
}, ];
let path = tempfile();
write_raw(&path, &nodes);
let q = load_from(&path).unwrap().unwrap();
let _ = fs::remove_file(&path);
let x = q.symbol("x", Domain::Real);
let y = q.symbol("y", Domain::Real);
let two = q.integer(2_i32);
let sum = q.add(vec![x, y]);
let square = q.pow(sum, two);
let y_one = q.pow(y, q.integer(1_i32));
assert_eq!(q.len(), 7, "a non-canonical node survived as a duplicate");
for id in [x, y, two, sum, square, y_one] {
assert!(
(id.0 as usize) < 7,
"{} was not in the loaded pool",
q.display(id)
);
}
}
#[test]
fn a_forward_child_reference_is_refused() {
let nodes = vec![ExprData::Add(vec![ExprId(0), ExprId(1)])];
let path = tempfile();
write_raw(&path, &nodes);
let r = load_from(&path);
let _ = fs::remove_file(&path);
assert!(
r.is_err(),
"a child reference past the node itself must be refused"
);
}
#[test]
fn big_o_round_trip() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let o = p.big_o(p.pow(x, p.integer(6)));
let path = tempfile();
p.checkpoint(&path).unwrap();
let q = ExprPool::open_persistent(&path).unwrap();
assert_eq!(q.get(o), p.get(o));
let _ = fs::remove_file(&path);
}
fn header(count: u64) -> Vec<u8> {
let mut v = MAGIC.to_vec();
v.extend_from_slice(&POOL_FORMAT_WRITE.to_le_bytes());
v.extend_from_slice(&0u32.to_le_bytes());
v.extend_from_slice(&count.to_le_bytes());
v
}
fn load_bytes(bytes: &[u8]) -> Result<Option<ExprPool>, IoError> {
let path = tempfile();
std::fs::write(&path, bytes).unwrap();
let r = load_from(&path);
let _ = fs::remove_file(&path);
r
}
fn le32(v: u32) -> [u8; 4] {
v.to_le_bytes()
}
#[test]
fn hostile_headers_are_refused_not_allocated() {
let mut cases: Vec<(&str, Vec<u8>, &str)> = Vec::new();
let mut b = header(1);
b.extend_from_slice(&[0, 0, 1]);
b.extend_from_slice(&le32(0xFFFF_FFF0));
cases.push(("huge string length", b, "E-IO-004"));
let mut b = header(1);
b.push(4);
b.extend_from_slice(&le32(0xFFFF_FFF0));
cases.push(("huge arity", b, "E-IO-004"));
let mut b = header(1);
b.push(8);
b.extend_from_slice(&le32(0xFFFF_FFF0));
cases.push(("huge branch count", b, "E-IO-004"));
let mut b = header(1 << 62);
b.extend_from_slice(&[0, 0, 1]);
b.extend_from_slice(&le32(1));
b.push(b'x');
cases.push(("huge node count", b, "E-IO-004"));
let mut b = header(1);
b.push(2);
b.extend_from_slice(&le32(1));
b.push(b'1');
b.extend_from_slice(&le32(1));
b.push(b'0');
cases.push(("zero denominator", b, "E-IO-009"));
for prec in [0u32, u32::MAX, MAX_FILE_FLOAT_PREC + 1] {
let mut b = header(1);
b.push(3);
b.extend_from_slice(&le32(prec));
b.extend_from_slice(&le32(1));
b.push(b'1');
cases.push(("float precision", b, "E-IO-009"));
}
use crate::errors::AlkahestError;
for (what, bytes, code) in cases {
let r = std::panic::catch_unwind(|| load_bytes(&bytes))
.unwrap_or_else(|_| panic!("{what}: loading panicked"));
match r {
Err(e) => assert_eq!(e.code(), code, "{what}: {e}"),
Ok(_) => panic!("{what}: a corrupt file loaded"),
}
}
}
#[test]
fn duplicate_nodes_do_not_drift_later_references() {
let sym = |n: &str| ExprData::Symbol {
name: n.into(),
domain: Domain::Real,
commutative: true,
};
let nodes = vec![
sym("x"),
sym("x"),
sym("y"),
ExprData::Pow {
base: ExprId(1),
exp: ExprId(2),
},
];
let path = tempfile();
write_raw(&path, &nodes);
let q = load_from(&path).unwrap().unwrap();
let _ = fs::remove_file(&path);
let x = q.symbol("x", Domain::Real);
let y = q.symbol("y", Domain::Real);
let before = q.len();
let _ = q.pow(x, y);
assert_eq!(q.len(), before, "x^y was not the node the file meant");
assert_eq!(q.len(), 3);
}
fn every_tag_file() -> Vec<u8> {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let t = p.symbol("t", Domain::Complex);
let big = p.integer(rug::Integer::from(rug::Integer::u_pow_u(10, 30)));
let r = p.rational(-3, 7);
let f = p.float(2.5, 80);
let s = p.add(vec![x, big, r, f]);
let m = p.mul(vec![x, t]);
let pw = p.pow(s, m);
let sn = p.func("sin", vec![pw]);
let g = p.func("f", vec![x, t, sn]);
let c = p.pred_gt(x, p.integer(0));
let pc = p.piecewise(vec![(c, g)], r);
let fa = p.forall(x, c);
let ex = p.exists(t, fa);
let o = p.big_o(p.pow(x, p.integer(3)));
let poly = p.add(vec![p.pow(t, p.integer(2)), p.integer(1)]);
let rs = p.root_sum(poly, t, p.func("log", vec![p.add(vec![x, t])]));
let _ = p.add(vec![pc, ex, o, rs]);
let path = tempfile();
p.checkpoint(&path).unwrap();
let bytes = std::fs::read(&path).unwrap();
let _ = fs::remove_file(&path);
bytes
}
fn load_and_touch(bytes: &[u8]) -> bool {
std::panic::catch_unwind(|| {
if let Ok(Some(q)) = load_bytes(bytes) {
for i in 0..q.len() {
let _ = q.display(ExprId(i as u32)).to_string();
}
}
})
.is_ok()
}
#[test]
fn truncated_and_bit_flipped_files_never_panic() {
let good = every_tag_file();
assert!(load_bytes(&good).unwrap().is_some());
let mut bad = Vec::new();
for n in 0..good.len() {
if !load_and_touch(&good[..n]) {
bad.push(format!("truncated to {n}"));
}
}
for i in 0..good.len() {
for bit in 0..8 {
let mut b = good.clone();
b[i] ^= 1 << bit;
if !load_and_touch(&b) {
bad.push(format!("byte {i} bit {bit}"));
}
}
}
assert!(bad.is_empty(), "panicked on: {bad:?}");
}
fn show(pool: &ExprPool, id: ExprId) -> String {
pool.display(id).to_string()
}
#[test]
fn dag_round_trips_into_a_fresh_pool_and_carries_only_what_is_reachable() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let y = p.symbol("y", Domain::Positive);
for i in 0..200 {
let _ = p.add(vec![x, p.integer(i)]);
}
let s = p.add(vec![x, y]);
let e1 = p.pow(s, p.rational(3, 7));
let e2 = p.func("sin", vec![p.mul(vec![s, p.float(0.25, 80)])]);
let bytes = dag_to_bytes(&p, &[e1, e2, e1]);
assert!(
bytes.len() < 400,
"{} bytes for two small expressions",
bytes.len()
);
let q = ExprPool::new();
let got = dag_from_bytes(&q, &bytes).unwrap();
assert_eq!(got.len(), 3);
assert_eq!(got[0], got[2]);
assert_eq!(show(&q, got[0]), show(&p, e1));
assert_eq!(show(&q, got[1]), show(&p, e2));
assert_eq!(q.get(q.symbol("y", Domain::Positive)), p.get(y));
let qy = q.symbol("y", Domain::Positive);
let qx = q.symbol("x", Domain::Real);
assert_eq!(q.add(vec![qx, qy]), q.add(vec![qy, qx]));
let before = q.len();
let _ = q.add(vec![qx, qy]);
assert_eq!(q.len(), before, "the shared x + y was not reused");
}
#[test]
fn dag_loaded_into_its_own_pool_returns_the_original_ids() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let e = p.add(vec![p.pow(x, p.integer(2)), p.func("exp", vec![x])]);
let before = p.len();
let got = dag_from_bytes(&p, &dag_to_bytes(&p, &[e])).unwrap();
assert_eq!(got, vec![e]);
assert_eq!(p.len(), before);
}
#[test]
fn dag_of_a_very_deep_expression_serialises_iteratively() {
let p = ExprPool::new();
let mut e = p.symbol("x", Domain::Real);
for _ in 0..200_000 {
e = p.func("sin", vec![e]);
}
let bytes = dag_to_bytes(&p, &[e]);
let q = ExprPool::new();
let got = dag_from_bytes(&q, &bytes).unwrap();
assert_eq!(q.len(), p.len());
assert!(matches!(q.get(got[0]), ExprData::Func { .. }));
}
#[test]
fn corrupt_dag_bytes_are_errors_not_panics() {
let p = ExprPool::new();
let x = p.symbol("x", Domain::Real);
let e = p.add(vec![p.pow(x, p.rational(1, 3)), p.float(1.5, 64), x]);
let good = dag_to_bytes(&p, &[e]);
let mut bad = Vec::new();
let try_load = |b: &[u8]| {
std::panic::catch_unwind(|| {
let q = ExprPool::new();
if let Ok(roots) = dag_from_bytes(&q, b) {
for r in roots {
let _ = q.display(r).to_string();
}
}
})
.is_ok()
};
for n in 0..good.len() {
if !try_load(&good[..n]) {
bad.push(format!("truncated to {n}"));
}
}
for i in 0..good.len() {
for bit in 0..8 {
let mut b = good.clone();
b[i] ^= 1 << bit;
if !try_load(&b) {
bad.push(format!("byte {i} bit {bit}"));
}
}
}
assert!(bad.is_empty(), "panicked on: {bad:?}");
assert!(matches!(
dag_from_bytes(&p, b"ALKP"),
Err(IoError::BadMagic)
));
}
}