use std::collections::HashMap;
use std::sync::Arc;
use crate::array::{Array, Data};
use crate::error::{Error, ErrorKind, Result, Span};
use crate::fmt::{format_array, FmtOpts};
use crate::frontend::Rules;
use crate::fuse::FusedKernel;
use crate::verb::{arrays_match, Agreement, Ctx, Env, EvalCfg, Verb};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Scope {
Local,
Global,
LocalDefault,
}
#[derive(Clone, Debug)]
pub enum Expr {
Const(Array, Span),
Param(usize, Span),
Name(String, Span),
Assign { name: String, value: Box<Expr>, scope: Scope, span: Span },
AmendIndex {
name: String,
slots: Vec<Option<Expr>>,
value: Box<Expr>,
origin: i64,
scope: Scope,
span: Span,
},
Control(Box<Control>, Span),
Monad { verb: Verb, y: Box<Expr>, span: Span },
Dyad { verb: Verb, x: Box<Expr>, y: Box<Expr>, span: Span },
PrintPass { value: Box<Expr>, bare: bool, span: Span },
Input { eval: bool, span: Span },
Fused { kernel: FusedKernel, inputs: Vec<Expr>, orig: Box<Expr>, span: Span },
Elided { orig: Vec<Expr>, span: Span },
VerbDef { name: String, verb: Verb, span: Span },
ModDef { name: String, spelling: String, conjunction: bool, span: Span },
}
#[derive(Clone, Debug)]
pub enum Control {
If { arms: Vec<Branch>, otherwise: Option<Vec<Expr>> },
While { test: Vec<Expr>, body: Vec<Expr>, body_first: bool, until: bool },
For { name: Option<String>, source: Box<Expr>, body: Vec<Expr> },
Select { subject: Box<Expr>, cases: Vec<Branch> },
Try { body: Vec<Expr>, catch: Vec<Expr> },
Return,
Break,
Branch(Box<Expr>),
Continue,
}
#[derive(Clone, Debug)]
pub struct Branch {
pub test: Option<Vec<Expr>>,
pub body: Vec<Expr>,
pub fall_through: bool,
}
pub const NILADIC: &str = "(no argument)";
#[derive(Debug)]
pub struct ExplicitDef {
pub name: String,
pub left: Option<String>,
pub right: String,
pub dyad_only: bool,
pub result: Option<String>,
pub locals: Vec<String>,
pub body: Vec<Expr>,
pub empty: Option<Array>,
pub labels: Vec<(String, usize)>,
pub pure: bool,
}
impl Expr {
pub(crate) fn depth(&self) -> usize {
let mut deepest = 0usize;
let mut stack: Vec<(&Expr, usize)> = vec![(self, 1)];
while let Some((e, d)) = stack.pop() {
deepest = deepest.max(d);
let kids: Vec<&Expr> = match e {
Expr::Const(..)
| Expr::Param(..)
| Expr::Name(..)
| Expr::Control(..)
| Expr::VerbDef { .. }
| Expr::Input { .. }
| Expr::ModDef { .. } => Vec::new(),
Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => vec![value],
Expr::AmendIndex { slots, value, .. } => {
slots.iter().flatten().chain(std::iter::once(&**value)).collect()
}
Expr::Monad { y, .. } => vec![y],
Expr::Dyad { x, y, .. } => vec![x, y],
Expr::Fused { inputs, orig, .. } => {
inputs.iter().chain(std::iter::once(&**orig)).collect()
}
Expr::Elided { orig, .. } => orig.iter().collect(),
};
stack.extend(kids.into_iter().map(|c| (c, d + 1)));
}
deepest
}
pub fn span(&self) -> Span {
match self {
Expr::Const(_, s) | Expr::Param(_, s) | Expr::Name(_, s) => *s,
Expr::Control(_, s) => *s,
Expr::AmendIndex { span, .. } | Expr::Input { span, .. } => *span,
Expr::Assign { span, .. }
| Expr::Monad { span, .. }
| Expr::Dyad { span, .. }
| Expr::PrintPass { span, .. }
| Expr::Fused { span, .. }
| Expr::Elided { span, .. }
| Expr::VerbDef { span, .. }
| Expr::ModDef { span, .. } => *span,
}
}
pub fn set_span(&mut self, to: Span) {
match self {
Expr::Const(_, s) | Expr::Param(_, s) | Expr::Name(_, s) => *s = to,
Expr::Control(_, s) => *s = to,
Expr::AmendIndex { span, .. } | Expr::Input { span, .. } => *span = to,
Expr::Assign { span, .. }
| Expr::Monad { span, .. }
| Expr::Dyad { span, .. }
| Expr::PrintPass { span, .. }
| Expr::Fused { span, .. }
| Expr::Elided { span, .. }
| Expr::VerbDef { span, .. }
| Expr::ModDef { span, .. } => *span = to,
}
}
fn is_silent(&self) -> bool {
matches!(
self,
Expr::Assign { .. }
| Expr::AmendIndex { .. }
| Expr::PrintPass { .. }
| Expr::Elided { .. }
| Expr::VerbDef { .. }
| Expr::ModDef { .. }
)
}
}
#[derive(Clone, Debug)]
pub struct ParamSpec {
pub name: String,
}
#[derive(Clone, Debug)]
pub struct Program {
pub stmts: Vec<Expr>,
pub params: Vec<ParamSpec>,
pub display_src: String,
pub agreement: Agreement,
pub fmt: FmtOpts,
pub rules: Rules,
}
#[derive(Clone, Debug)]
pub(crate) struct Note {
pub shape: Vec<usize>,
pub dtype: crate::dtype::DType,
pub layout: crate::array::Layout,
pub kernel_ran: Option<bool>,
pub decline: Option<crate::fuse::Decline>,
pub placement: crate::device::Placement,
}
pub(crate) type Trace = HashMap<usize, Note>;
pub(crate) fn key(e: &Expr) -> usize {
std::ptr::from_ref(e) as usize
}
impl Program {
pub fn run(&self, args: &[Array], out: &mut dyn FnMut(&str)) -> Result<Option<Array>> {
self.exec(args, out, None, &mut None, None)
}
pub fn run_io(
&self,
args: &[Array],
out: &mut dyn FnMut(&str),
inp: &mut dyn FnMut() -> Option<String>,
) -> Result<Option<Array>> {
self.exec(args, out, Some(inp), &mut None, None)
}
pub fn run_on_io(
&self,
device: &crate::device::Device,
args: &[Array],
out: &mut dyn FnMut(&str),
inp: &mut dyn FnMut() -> Option<String>,
) -> Result<Option<Array>> {
self.exec(args, out, Some(inp), &mut None, Some(device))
}
pub fn run_on(
&self,
device: &crate::device::Device,
args: &[Array],
out: &mut dyn FnMut(&str),
) -> Result<Option<Array>> {
self.exec(args, out, None, &mut None, Some(device))
}
pub(crate) fn trace(
&self,
args: &[Array],
out: &mut dyn FnMut(&str),
device: Option<&crate::device::Device>,
) -> (Result<Option<Array>>, Trace) {
let mut rec = Some(Trace::new());
let r = self.exec(args, out, None, &mut rec, device);
(r, rec.expect("the recorder stays in place"))
}
fn exec(
&self,
args: &[Array],
out: &mut dyn FnMut(&str),
inp: crate::verb::InputFn<'_>,
rec: &mut Option<Trace>,
device: Option<&crate::device::Device>,
) -> Result<Option<Array>> {
if args.len() != self.params.len() {
let names: Vec<&str> = self.params.iter().map(|p| p.name.as_str()).collect();
let wanted = if names.is_empty() {
"no arguments".to_string()
} else {
format!("one value for each of {}", names.join(", "))
};
return Err(Error::new(
ErrorKind::Value,
format!("this program takes {wanted}, and was given {}", args.len()),
None,
));
}
let cfg = EvalCfg {
agreement: self.agreement,
fmt: self.fmt,
tol: self.rules.tol(),
rules: self.rules,
};
let mut env = Env::new(args.to_vec());
let mut inp = inp;
let inp = crate::verb::reborrow_input(&mut inp);
let mut ctx = Ctx { cfg, out, inp, env: &mut env, device };
let mut last = None;
for stmt in &self.stmts {
let (v, flow) = eval_stmt(stmt, &mut ctx, rec)?;
if flow != Flow::Normal {
return Err(Error::internal("a control signal escaped to the top level"));
}
last = if stmt.is_silent() { None } else { v };
}
Ok(last)
}
pub fn render_error(&self, e: &Error) -> String {
e.render(&self.display_src)
}
pub fn explain(&self, args: Option<&[Array]>) -> String {
crate::explain::explain(self, args, None)
}
pub fn explain_on(
&self,
device: &crate::device::Device,
args: Option<&[Array]>,
) -> String {
crate::explain::explain(self, args, Some(device))
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Flow {
Normal,
Return,
Break,
Continue,
Goto(usize),
}
pub(crate) fn run_block(
stmts: &[Expr],
last: Option<Array>,
ctx: &mut Ctx<'_>,
rec: &mut Option<Trace>,
) -> Result<(Option<Array>, Flow)> {
let mut last = last;
for stmt in stmts {
let (v, flow) = eval_stmt(stmt, ctx, rec)?;
if let Some(v) = v {
last = Some(v);
}
if flow != Flow::Normal {
return Ok((last, flow));
}
}
Ok((last, Flow::Normal))
}
fn eval_stmt(
e: &Expr,
ctx: &mut Ctx<'_>,
rec: &mut Option<Trace>,
) -> Result<(Option<Array>, Flow)> {
let Expr::Control(c, span) = e else {
return Ok((Some(eval(e, ctx, rec)?), Flow::Normal));
};
let (v, flow) = eval_control(c, *span, ctx, rec)?;
let v = match (v, flow) {
(Some(v), _) => Some(v),
(None, Flow::Normal) => ctx.env.current_def().and_then(|d| d.empty.clone()),
(None, _) => None,
};
if let (Some(t), Some(v)) = (rec.as_mut(), v.as_ref()) {
t.insert(
key(e),
Note {
shape: v.shape.clone(),
dtype: v.dtype(),
layout: v.layout(),
kernel_ran: None,
decline: None,
placement: crate::device::Placement::Default,
},
);
}
Ok((v, flow))
}
pub(crate) fn empty_result() -> Array {
Array::new(vec![0, 0], Data::I64(Vec::new().into()))
}
fn is_true(a: &Array, span: Span) -> Result<bool> {
if a.count() == 0 {
return Ok(true);
}
match &a.data {
Data::I64(v) => Ok(v.as_slice()[0] != 0),
Data::F64(v) => Ok(v.as_slice()[0] != 0.0),
Data::Bool(v) => Ok(v.as_slice()[0] != 0),
Data::Char(v) => Ok(v.as_slice()[0] as u32 != 0),
Data::Complex(v) => Ok(v.as_slice()[0] != crate::complex::ZERO),
Data::Ext(v) => Ok(v.as_slice()[0] != crate::exact::Ext::default()),
Data::Rat(v) => Ok(!v.as_slice()[0].is_zero()),
Data::Box(_) => Err(Error::domain("a condition must be numeric, not boxed", span)),
Data::Symbol(_) => {
Err(Error::domain("a condition must be numeric, not a symbol", span))
}
}
}
fn eval_control(
c: &Control,
span: Span,
ctx: &mut Ctx<'_>,
rec: &mut Option<Trace>,
) -> Result<(Option<Array>, Flow)> {
match c {
Control::Return => Ok((None, Flow::Return)),
Control::Branch(target) => {
let to = eval(target, ctx, rec)?;
if to.count() == 0 {
return Ok((None, Flow::Normal));
}
let line = to
.to_i64_vec()
.and_then(|v| v.first().copied())
.ok_or_else(|| Error::domain("a branch target is a line number", span))?;
let lines = ctx.env.current_def().map_or(0, |d| d.body.len() as i64);
if line >= 1 && line <= lines {
return Ok((None, Flow::Goto(line as usize - 1)));
}
Ok((None, Flow::Return))
}
Control::Break => Ok((None, Flow::Break)),
Control::Continue => Ok((None, Flow::Continue)),
Control::If { arms, otherwise } => {
for arm in arms {
let test = arm.test.as_deref().unwrap_or(&[]);
let (t, flow) = run_block(test, None, ctx, rec)?;
if flow != Flow::Normal {
return Ok((t, flow));
}
let taken = match &t {
Some(v) => is_true(v, span)?,
None => true,
};
if taken {
return run_block(&arm.body, None, ctx, rec);
}
}
match otherwise {
Some(body) => run_block(body, None, ctx, rec),
None => Ok((None, Flow::Normal)),
}
}
Control::While { test, body, body_first, until } => {
let mut last = None;
let mut first = *body_first;
loop {
if !first {
let (t, flow) = run_block(test, None, ctx, rec)?;
if flow != Flow::Normal {
return Ok((t, flow));
}
let mut go = match &t {
Some(v) => is_true(v, span)?,
None => false,
};
if *until {
go = !go;
}
if !go {
return Ok((last, Flow::Normal));
}
}
first = false;
let (v, flow) = run_block(body, last, ctx, rec)?;
last = v;
match flow {
Flow::Normal | Flow::Continue => {}
Flow::Break => return Ok((last, Flow::Normal)),
other => return Ok((last, other)),
}
}
}
Control::For { name, source, body } => {
let src = eval(source, ctx, rec)?;
let n = if src.rank() == 0 { 1 } else { src.shape[0] };
let mut last = None;
for i in 0..n {
if let Some(name) = name {
let item = if src.rank() == 0 { src.clone() } else { src.item(i) };
ctx.env.assign(name.clone(), item, Scope::Local);
ctx.env.assign(
format!("{name}_index"),
Array::scalar_i64(i as i64),
Scope::Local,
);
}
let (v, flow) = run_block(body, last, ctx, rec)?;
last = v;
match flow {
Flow::Normal | Flow::Continue => {}
Flow::Break => return Ok((last, Flow::Normal)),
other => return Ok((last, other)),
}
}
Ok((last, Flow::Normal))
}
Control::Select { subject, cases } => {
let subject = eval(subject, ctx, rec)?;
let tol = ctx.cfg.tol;
let mut running = false;
let mut last = None;
for case in cases {
if !running {
match &case.test {
None => running = true,
Some(test) => {
let (t, flow) = run_block(test, None, ctx, rec)?;
if flow != Flow::Normal {
return Ok((t, flow));
}
running = t.is_some_and(|v| arrays_match(&subject, &v, tol));
}
}
}
if running {
let (v, flow) = run_block(&case.body, last, ctx, rec)?;
last = v;
if flow != Flow::Normal {
return Ok((last, flow));
}
if !case.fall_through {
return Ok((last, Flow::Normal));
}
running = true;
}
}
Ok((last, Flow::Normal))
}
Control::Try { body, catch } => {
match run_block(body, None, ctx, rec) {
Ok(r) => Ok(r),
Err(e) if matches!(e.kind, ErrorKind::NotYet | ErrorKind::Internal) => Err(e),
Err(_) => run_block(catch, None, ctx, rec),
}
}
}
}
pub(crate) fn call_explicit(
def: &Arc<ExplicitDef>,
x: Option<&Array>,
y: &Array,
ctx: &mut Ctx<'_>,
span: Span,
) -> Result<Array> {
if x.is_some() && def.left.is_none() {
return Err(Error::new(
ErrorKind::Domain,
format!("{} has no dyadic definition", def.name),
Some(span),
));
}
if x.is_none() && def.dyad_only {
return Err(Error::new(
ErrorKind::Domain,
format!(
"{} has no monadic definition: it names {}",
def.name,
def.left.as_deref().unwrap_or("a left argument")
),
Some(span),
));
}
let mut frame: HashMap<String, Array> = HashMap::new();
frame.insert(def.right.clone(), y.clone());
if let (Some(name), Some(v)) = (&def.left, x) {
frame.insert(name.clone(), v.clone());
}
for (label, at) in &def.labels {
frame.insert(label.clone(), Array::scalar_i64(*at as i64 + 1));
}
ctx.env.enter(frame, Arc::clone(def), span)?;
let mut rec = None;
let out = run_body(&def.body, ctx, &mut rec);
let frame = ctx.env.leave();
let value = out?;
if let Some(name) = &def.result {
return frame.get(name).cloned().ok_or_else(|| {
Error::new(
ErrorKind::Value,
format!("{} did not set its result {name}", def.name),
Some(span),
)
});
}
match value {
Some(v) => Ok(v),
None => def.empty.clone().ok_or_else(|| {
Error::new(
ErrorKind::Value,
format!("{} produced no result", def.name),
Some(span),
)
}),
}
}
const BRANCH_LIMIT: usize = 1 << 22;
fn run_body(
stmts: &[Expr],
ctx: &mut Ctx<'_>,
rec: &mut Option<Trace>,
) -> Result<Option<Array>> {
let mut last = None;
let mut at = 0usize;
let mut steps = 0usize;
while at < stmts.len() {
steps += 1;
if steps > BRANCH_LIMIT {
return Err(Error::new(
ErrorKind::Domain,
format!("a definition branched more than {BRANCH_LIMIT} times"),
Some(stmts[at].span()),
)
.note("a loop written with → needs a branch that leaves it"));
}
let (v, flow) = eval_stmt(&stmts[at], ctx, rec)?;
if let Some(v) = v {
last = Some(v);
}
match flow {
Flow::Normal => at += 1,
Flow::Goto(to) => at = to,
_ => break,
}
}
Ok(last)
}
pub(crate) fn fold_const(e: &Expr, cfg: EvalCfg) -> Option<Array> {
fn closed(e: &Expr) -> bool {
match e {
Expr::Const(..) => true,
Expr::Monad { verb, y, .. } => verb.is_pure() && closed(y),
Expr::Dyad { verb, x, y, .. } => verb.is_pure() && closed(x) && closed(y),
_ => false,
}
}
if !closed(e) {
return None;
}
cfg.pure(|ctx| eval(e, ctx, &mut None).ok())
}
fn eval(e: &Expr, ctx: &mut Ctx<'_>, rec: &mut Option<Trace>) -> Result<Array> {
let _depth = crate::verb::Nesting::enter(e.span())?;
let v = eval_node(e, ctx, rec)?;
if let Some(t) = rec.as_mut() {
let (kernel_ran, decline, placement) = t.get(&key(e)).map_or(
(None, None, crate::device::Placement::Default),
|n| (n.kernel_ran, n.decline, n.placement.clone()),
);
t.insert(
key(e),
Note {
shape: v.shape.clone(),
dtype: v.dtype(),
layout: v.layout(),
kernel_ran,
decline,
placement,
},
);
}
Ok(v)
}
fn eval_node(e: &Expr, ctx: &mut Ctx<'_>, rec: &mut Option<Trace>) -> Result<Array> {
match e {
Expr::Const(a, _) => Ok(a.clone()),
Expr::Param(i, _) => ctx.env.arg(*i),
Expr::Name(n, span) => ctx.env.get(n).ok_or_else(|| {
Error::new(ErrorKind::Value, format!("undefined name: {n}"), Some(*span))
}),
Expr::Assign { name, value, scope, .. } => {
let v = eval(value, ctx, rec)?;
ctx.env.assign(name.clone(), v.clone(), *scope);
Ok(v)
}
Expr::AmendIndex { name, slots, value, origin, scope, span } => {
let base = ctx.env.get(name).ok_or_else(|| {
Error::new(ErrorKind::Value, format!("undefined name: {name}"), Some(*span))
})?;
let v = eval(value, ctx, rec)?;
let mut idx = Vec::with_capacity(slots.len());
for slot in slots {
idx.push(match slot {
Some(e) => Some(eval(e, ctx, rec)?),
None => None,
});
}
let out = crate::verb::amend_at(&base, &idx, &v, *origin, *span)?;
ctx.env.assign(name.clone(), out.clone(), *scope);
Ok(out)
}
Expr::Control(..) => {
Err(Error::internal("a control sentence appeared in expression position"))
}
Expr::Monad { verb, y, span } => {
let vy = eval(y, ctx, rec)?;
verb.monad(&vy, ctx, *span)
}
Expr::Dyad { verb, x, y, span } => {
let vy = eval(y, ctx, rec)?;
let vx = eval(x, ctx, rec)?;
verb.dyad(&vx, &vy, ctx, *span)
}
Expr::PrintPass { value, bare, .. } => {
let v = eval(value, ctx, rec)?;
let text = format_array(&v, &ctx.cfg.fmt);
(ctx.out)(&text);
if !bare {
(ctx.out)("\n");
}
Ok(v)
}
Expr::Input { eval: run_it, span } => {
let line = ctx.read_line(*span)?;
if !run_it {
return Ok(Array::from_chars(line.chars().collect()));
}
crate::verb::execute_source(&line, true, ctx, *span)
}
Expr::Fused { kernel, inputs, orig, .. } => {
let mut vals = Vec::with_capacity(inputs.len());
for e in inputs {
vals.push(eval(e, ctx, rec)?);
}
let (ran, placement) = crate::fuse::eval_on(ctx.device, kernel, &vals);
if let Some(t) = rec.as_mut() {
let decline =
if ran.is_none() { crate::fuse::decline_reason(kernel, &vals) } else { None };
t.insert(
key(e),
Note {
shape: Vec::new(),
dtype: crate::dtype::DType::I64,
layout: crate::array::Layout::RowMajor,
kernel_ran: Some(ran.is_some()),
decline,
placement,
},
);
}
match ran {
Some(a) => Ok(a),
None => {
let tree = crate::fuse::fallback_tree(kernel, orig, &vals);
let v = eval(&tree, ctx, &mut None)?;
Ok(crate::fuse::fallback_finish(kernel, v))
}
}
}
Expr::VerbDef { name, verb, .. } => {
ctx.env.define(name.clone(), verb.clone());
Ok(Array::scalar_i64(0))
}
Expr::Elided { .. } | Expr::ModDef { .. } => Ok(Array::scalar_i64(0)),
}
}