Array and dictionary indexing

6 files changed, 261 insertions, 137 deletions

diff --git a/src/eval/array.rs b/src/eval/array.rs
index 38876c96..356aa0ca 100644
--- a/src/eval/array.rs
+++ b/src/eval/array.rs
@@ -1,12 +1,14 @@
+use std::convert::TryFrom;
 use std::fmt::{self, Debug, Formatter};
 use std::iter::FromIterator;
 use std::ops::{Add, AddAssign};
 use std::rc::Rc;
 
 use super::Value;
+use crate::diag::StrResult;
 
 /// Create a new [`Array`] from values.
-#[macro_export]
+#[allow(unused_macros)]
 macro_rules! array {
     ($value:expr; $count:expr) => {
         $crate::eval::Array::from_vec(vec![$crate::eval::Value::from($value); $count])
@@ -47,25 +49,25 @@ impl Array {
     }
 
     /// The length of the array.
-    pub fn len(&self) -> usize {
-        self.vec.len()
+    pub fn len(&self) -> i64 {
+        self.vec.len() as i64
     }
 
     /// Borrow the value at the given index.
-    pub fn get(&self, index: usize) -> Option<&Value> {
-        self.vec.get(index)
+    pub fn get(&self, index: i64) -> StrResult<&Value> {
+        usize::try_from(index)
+            .ok()
+            .and_then(|i| self.vec.get(i))
+            .ok_or_else(|| out_of_bounds(index, self.len()))
     }
 
     /// Mutably borrow the value at the given index.
-    pub fn get_mut(&mut self, index: usize) -> Option<&mut Value> {
-        Rc::make_mut(&mut self.vec).get_mut(index)
-    }
-
-    /// Set the value at the given index.
-    ///
-    /// This panics the `index` is out of range.
-    pub fn set(&mut self, index: usize, value: Value) {
-        Rc::make_mut(&mut self.vec)[index] = value;
+    pub fn get_mut(&mut self, index: i64) -> StrResult<&mut Value> {
+        let len = self.len();
+        usize::try_from(index)
+            .ok()
+            .and_then(move |i| Rc::make_mut(&mut self.vec).get_mut(i))
+            .ok_or_else(|| out_of_bounds(index, len))
     }
 
     /// Push a value to the end of the array.
@@ -82,16 +84,26 @@ impl Array {
         }
     }
 
-    /// Repeat this array `n` times.
-    pub fn repeat(&self, n: usize) -> Self {
-        let len = self.len().checked_mul(n).expect("capacity overflow");
-        self.iter().cloned().cycle().take(len).collect()
-    }
-
     /// Iterate over references to the contained values.
     pub fn iter(&self) -> std::slice::Iter<Value> {
         self.vec.iter()
     }
+
+    /// Repeat this array `n` times.
+    pub fn repeat(&self, n: i64) -> StrResult<Self> {
+        let count = usize::try_from(n)
+            .ok()
+            .and_then(|n| self.vec.len().checked_mul(n))
+            .ok_or_else(|| format!("cannot repeat this array {} times", n))?;
+
+        Ok(self.iter().cloned().cycle().take(count).collect())
+    }
+}
+
+/// The out of bounds access error message.
+#[cold]
+fn out_of_bounds(index: i64, len: i64) -> String {
+    format!("array index out of bounds (index: {}, len: {})", index, len)
 }
 
 impl Default for Array {
diff --git a/src/eval/dict.rs b/src/eval/dict.rs
index bc7df2b8..730392dc 100644
--- a/src/eval/dict.rs
+++ b/src/eval/dict.rs
@@ -5,10 +5,12 @@ use std::ops::{Add, AddAssign};
 use std::rc::Rc;
 
 use super::Value;
+use crate::diag::StrResult;
+use crate::pretty::pretty;
 use crate::util::EcoString;
 
 /// Create a new [`Dict`] from key-value pairs.
-#[macro_export]
+#[allow(unused_macros)]
 macro_rules! dict {
     ($($key:expr => $value:expr),* $(,)?) => {{
         #[allow(unused_mut)]
@@ -41,18 +43,21 @@ impl Dict {
     }
 
     /// The number of pairs in the dictionary.
-    pub fn len(&self) -> usize {
-        self.map.len()
+    pub fn len(&self) -> i64 {
+        self.map.len() as i64
     }
 
     /// Borrow the value the given `key` maps to.
-    pub fn get(&self, key: &str) -> Option<&Value> {
-        self.map.get(key)
+    pub fn get(&self, key: &str) -> StrResult<&Value> {
+        self.map.get(key).ok_or_else(|| missing_key(key))
     }
 
     /// Mutably borrow the value the given `key` maps to.
-    pub fn get_mut(&mut self, key: &str) -> Option<&mut Value> {
-        Rc::make_mut(&mut self.map).get_mut(key)
+    ///
+    /// This inserts the key with [`None`](Value::None) as the value if not
+    /// present so far.
+    pub fn get_mut(&mut self, key: EcoString) -> &mut Value {
+        Rc::make_mut(&mut self.map).entry(key).or_default()
     }
 
     /// Insert a mapping from the given `key` to the given `value`.
@@ -75,6 +80,12 @@ impl Dict {
     }
 }
 
+/// The missing key access error message.
+#[cold]
+fn missing_key(key: &str) -> String {
+    format!("dictionary does not contain key: {}", pretty(key))
+}
+
 impl Default for Dict {
     fn default() -> Self {
         Self::new()
diff --git a/src/eval/function.rs b/src/eval/function.rs
index 8b1e883f..be7894f4 100644
--- a/src/eval/function.rs
+++ b/src/eval/function.rs
@@ -134,7 +134,7 @@ impl FuncArgs {
         let value = self.items.remove(index).value;
         let span = value.span;
 
-        T::cast(value).map(Some).map_err(Error::partial(self.source, span))
+        T::cast(value).map(Some).map_err(Error::at(self.source, span))
     }
 
     /// Return an "unexpected argument" error if there is any remaining
@@ -146,3 +146,40 @@ impl FuncArgs {
         Ok(())
     }
 }
+
+impl FuncArgs {
+    /// Reinterpret these arguments as actually being an array index.
+    pub fn into_index(self) -> TypResult<i64> {
+        self.into_castable("index")
+    }
+
+    /// Reinterpret these arguments as actually being a dictionary key.
+    pub fn into_key(self) -> TypResult<EcoString> {
+        self.into_castable("key")
+    }
+
+    /// Reinterpret these arguments as actually being a single castable thing.
+    fn into_castable<T>(self, what: &str) -> TypResult<T>
+    where
+        T: Cast<Value>,
+    {
+        let mut iter = self.items.into_iter();
+        let value = match iter.next() {
+            None => {
+                bail!(self.source, self.span, "missing {}", what);
+            }
+            Some(FuncArg { name: Some(_), span, .. }) => {
+                bail!(self.source, span, "named pair is not allowed here");
+            }
+            Some(FuncArg { name: None, value, .. }) => {
+                value.v.cast().map_err(Error::at(self.source, value.span))?
+            }
+        };
+
+        if let Some(arg) = iter.next() {
+            bail!(self.source, arg.span, "only one {} is allowed", what);
+        }
+
+        Ok(value)
+    }
+}
diff --git a/src/eval/mod.rs b/src/eval/mod.rs
index e911596d..abb7cab7 100644
--- a/src/eval/mod.rs
+++ b/src/eval/mod.rs
@@ -20,6 +20,7 @@ pub use scope::*;
 pub use template::*;
 pub use value::*;
 
+use std::cell::RefMut;
 use std::collections::HashMap;
 use std::io;
 use std::mem;
@@ -34,7 +35,7 @@ use crate::parse::parse;
 use crate::source::{SourceId, SourceStore};
 use crate::syntax::visit::Visit;
 use crate::syntax::*;
-use crate::util::EcoString;
+use crate::util::{EcoString, RefMutExt};
 use crate::Context;
 
 /// Evaluate a parsed source file into a module.
@@ -60,15 +61,6 @@ pub struct Module {
     pub template: Template,
 }
 
-/// Evaluate an expression.
-pub trait Eval {
-    /// The output of evaluating the expression.
-    type Output;
-
-    /// Evaluate the expression to the output value.
-    fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output>;
-}
-
 /// The context for evaluation.
 pub struct EvalContext<'a> {
     /// The loader from which resources (files and images) are loaded.
@@ -168,45 +160,19 @@ impl<'a> EvalContext<'a> {
     }
 }
 
+/// Evaluate an expression.
+pub trait Eval {
+    /// The output of evaluating the expression.
+    type Output;
+
+    /// Evaluate the expression to the output value.
+    fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output>;
+}
+
 impl Eval for Rc<SyntaxTree> {
     type Output = Template;
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
-        trait Walk {
-            fn walk(&self, ctx: &mut EvalContext) -> TypResult<()>;
-        }
-
-        impl Walk for SyntaxTree {
-            fn walk(&self, ctx: &mut EvalContext) -> TypResult<()> {
-                for node in self.iter() {
-                    node.walk(ctx)?;
-                }
-                Ok(())
-            }
-        }
-
-        impl Walk for SyntaxNode {
-            fn walk(&self, ctx: &mut EvalContext) -> TypResult<()> {
-                match self {
-                    Self::Text(_) => {}
-                    Self::Space => {}
-                    Self::Linebreak(_) => {}
-                    Self::Parbreak(_) => {}
-                    Self::Strong(_) => {}
-                    Self::Emph(_) => {}
-                    Self::Raw(_) => {}
-                    Self::Heading(n) => n.body.walk(ctx)?,
-                    Self::List(n) => n.body.walk(ctx)?,
-                    Self::Enum(n) => n.body.walk(ctx)?,
-                    Self::Expr(n) => {
-                        let value = n.eval(ctx)?;
-                        ctx.map.insert(n as *const _, value);
-                    }
-                }
-                Ok(())
-            }
-        }
-
         let map = {
             let prev = mem::take(&mut ctx.map);
             self.walk(ctx)?;
@@ -302,8 +268,8 @@ impl Eval for BlockExpr {
         let mut output = Value::None;
         for expr in &self.exprs {
             let value = expr.eval(ctx)?;
-            output = ops::join(output, value)
-                .map_err(Error::partial(ctx.source, expr.span()))?;
+            output =
+                ops::join(output, value).map_err(Error::at(ctx.source, expr.span()))?;
         }
 
         if self.scoping {
@@ -324,7 +290,7 @@ impl Eval for UnaryExpr {
             UnOp::Neg => ops::neg(value),
             UnOp::Not => ops::not(value),
         };
-        result.map_err(Error::partial(ctx.source, self.span))
+        result.map_err(Error::at(ctx.source, self.span))
     }
 }
 
@@ -371,7 +337,7 @@ impl BinaryExpr {
         }
 
         let rhs = self.rhs.eval(ctx)?;
-        op(lhs, rhs).map_err(Error::partial(ctx.source, self.span))
+        op(lhs, rhs).map_err(Error::at(ctx.source, self.span))
     }
 
     /// Apply an assignment operation.
@@ -379,27 +345,11 @@ impl BinaryExpr {
     where
         F: FnOnce(Value, Value) -> StrResult<Value>,
     {
-        let lspan = self.lhs.span();
-        let slot = if let Expr::Ident(id) = self.lhs.as_ref() {
-            match ctx.scopes.get(id) {
-                Some(slot) => Rc::clone(slot),
-                None => bail!(ctx.source, lspan, "unknown variable"),
-            }
-        } else {
-            bail!(ctx.source, lspan, "cannot assign to this expression",);
-        };
-
+        let source = ctx.source;
         let rhs = self.rhs.eval(ctx)?;
-        let mut mutable = match slot.try_borrow_mut() {
-            Ok(mutable) => mutable,
-            Err(_) => {
-                bail!(ctx.source, lspan, "cannot assign to a constant",);
-            }
-        };
-
-        let lhs = mem::take(&mut *mutable);
-        *mutable = op(lhs, rhs).map_err(Error::partial(ctx.source, self.span))?;
-
+        let mut target = self.lhs.access(ctx)?;
+        let lhs = mem::take(&mut *target);
+        *target = op(lhs, rhs).map_err(Error::at(source, self.span))?;
         Ok(Value::None)
     }
 }
@@ -408,32 +358,45 @@ impl Eval for CallExpr {
     type Output = Value;
 
     fn eval(&self, ctx: &mut EvalContext) -> TypResult<Self::Output> {
-        let callee = self
-            .callee
-            .eval(ctx)?
-            .cast::<Function>()
-            .map_err(Error::partial(ctx.source, self.callee.span()))?;
-
+        let callee = self.callee.eval(ctx)?;
         let mut args = self.args.eval(ctx)?;
-        let returned = callee(ctx, &mut args).map_err(|mut errors| {
-            for error in errors.iter_mut() {
-                // Skip errors directly related to arguments.
-                if error.source == ctx.source && self.span.contains(error.span) {
-                    continue;
-                }
 
-                error.trace.push((
-                    ctx.source,
-                    self.span,
-                    Tracepoint::Call(callee.name().map(Into::into)),
-                ));
+        match callee {
+            Value::Array(array) => array
+                .get(args.into_index()?)
+                .map(Value::clone)
+                .map_err(Error::at(ctx.source, self.span)),
+
+            Value::Dict(dict) => dict
+                .get(&args.into_key()?)
+                .map(Value::clone)
+                .map_err(Error::at(ctx.source, self.span)),
+
+            Value::Func(func) => {
+                let returned = func(ctx, &mut args).map_err(|mut errors| {
+                    for error in errors.iter_mut() {
+                        // Skip errors directly related to arguments.
+                        if error.source != ctx.source || !self.span.contains(error.span) {
+                            error.trace.push((
+                                ctx.source,
+                                self.span,
+                                Tracepoint::Call(func.name().map(Into::into)),
+                            ));
+                        }
+                    }
+                    errors
+                })?;
+                args.finish()?;
+                Ok(returned)
             }
-            errors
-        })?;
 
-        args.finish()?;
-
-        Ok(returned)
+            v => bail!(
+                ctx.source,
+                self.callee.span(),
+                "expected function or collection, found {}",
+                v.type_name(),
+            ),
+        }
     }
 }
 
@@ -518,7 +481,7 @@ impl Eval for WithExpr {
             .callee
             .eval(ctx)?
             .cast::<Function>()
-            .map_err(Error::partial(ctx.source, self.callee.span()))?;
+            .map_err(Error::at(ctx.source, self.callee.span()))?;
 
         let applied = self.args.eval(ctx)?;
 
@@ -568,7 +531,7 @@ impl Eval for IfExpr {
             .condition
             .eval(ctx)?
             .cast::<bool>()
-            .map_err(Error::partial(ctx.source, self.condition.span()))?;
+            .map_err(Error::at(ctx.source, self.condition.span()))?;
 
         if condition {
             self.if_body.eval(ctx)
@@ -590,11 +553,11 @@ impl Eval for WhileExpr {
             .condition
             .eval(ctx)?
             .cast::<bool>()
-            .map_err(Error::partial(ctx.source, self.condition.span()))?
+            .map_err(Error::at(ctx.source, self.condition.span()))?
         {
             let value = self.body.eval(ctx)?;
             output = ops::join(output, value)
-                .map_err(Error::partial(ctx.source, self.body.span()))?;
+                .map_err(Error::at(ctx.source, self.body.span()))?;
         }
 
         Ok(output)
@@ -616,7 +579,7 @@ impl Eval for ForExpr {
 
                     let value = self.body.eval(ctx)?;
                     output = ops::join(output, value)
-                        .map_err(Error::partial(ctx.source, self.body.span()))?;
+                        .map_err(Error::at(ctx.source, self.body.span()))?;
                 }
 
                 ctx.scopes.exit();
@@ -662,7 +625,7 @@ impl Eval for ImportExpr {
             .path
             .eval(ctx)?
             .cast::<EcoString>()
-            .map_err(Error::partial(ctx.source, self.path.span()))?;
+            .map_err(Error::at(ctx.source, self.path.span()))?;
 
         let file = ctx.import(&path, self.path.span())?;
         let module = &ctx.modules[&file];
@@ -696,7 +659,7 @@ impl Eval for IncludeExpr {
             .path
             .eval(ctx)?
             .cast::<EcoString>()
-            .map_err(Error::partial(ctx.source, self.path.span()))?;
+            .map_err(Error::at(ctx.source, self.path.span()))?;
 
         let file = ctx.import(&path, self.path.span())?;
         let module = &ctx.modules[&file];
@@ -704,3 +667,97 @@ impl Eval for IncludeExpr {
         Ok(Value::Template(module.template.clone()))
     }
 }
+
+/// Walk a node in a template, filling the context's expression map.
+pub trait Walk {
+    /// Walk the node.
+    fn walk(&self, ctx: &mut EvalContext) -> TypResult<()>;
+}
+
+impl Walk for SyntaxTree {
+    fn walk(&self, ctx: &mut EvalContext) -> TypResult<()> {
+        for node in self.iter() {
+            node.walk(ctx)?;
+        }
+        Ok(())
+    }
+}
+
+impl Walk for SyntaxNode {
+    fn walk(&self, ctx: &mut EvalContext) -> TypResult<()> {
+        match self {
+            Self::Text(_) => {}
+            Self::Space => {}
+            Self::Linebreak(_) => {}
+            Self::Parbreak(_) => {}
+            Self::Strong(_) => {}
+            Self::Emph(_) => {}
+            Self::Raw(_) => {}
+            Self::Heading(n) => n.body.walk(ctx)?,
+            Self::List(n) => n.body.walk(ctx)?,
+            Self::Enum(n) => n.body.walk(ctx)?,
+            Self::Expr(n) => {
+                let value = n.eval(ctx)?;
+                ctx.map.insert(n as *const _, value);
+            }
+        }
+        Ok(())
+    }
+}
+
+/// Try to mutably access the value an expression points to.
+///
+/// This only works if the expression is a valid lvalue.
+pub trait Access {
+    /// Try to access the value.
+    fn access<'a>(&self, ctx: &'a mut EvalContext) -> TypResult<RefMut<'a, Value>>;
+}
+
+impl Access for Expr {
+    fn access<'a>(&self, ctx: &'a mut EvalContext) -> TypResult<RefMut<'a, Value>> {
+        match self {
+            Expr::Ident(ident) => ident.access(ctx),
+            Expr::Call(call) => call.access(ctx),
+            _ => bail!(
+                ctx.source,
+                self.span(),
+                "cannot access this expression mutably",
+            ),
+        }
+    }
+}
+
+impl Access for Ident {
+    fn access<'a>(&self, ctx: &'a mut EvalContext) -> TypResult<RefMut<'a, Value>> {
+        ctx.scopes
+            .get(self)
+            .ok_or("unknown variable")
+            .and_then(|slot| {
+                slot.try_borrow_mut().map_err(|_| "cannot mutate a constant")
+            })
+            .map_err(Error::at(ctx.source, self.span))
+    }
+}
+
+impl Access for CallExpr {
+    fn access<'a>(&self, ctx: &'a mut EvalContext) -> TypResult<RefMut<'a, Value>> {
+        let source = ctx.source;
+        let args = self.args.eval(ctx)?;
+        let guard = self.callee.access(ctx)?;
+
+        RefMut::try_map(guard, |value| match value {
+            Value::Array(array) => array
+                .get_mut(args.into_index()?)
+                .map_err(Error::at(source, self.span)),
+
+            Value::Dict(dict) => Ok(dict.get_mut(args.into_key()?)),
+
+            v => bail!(
+                source,
+                self.callee.span(),
+                "expected collection, found {}",
+                v.type_name(),
+            ),
+        })
+    }
+}
diff --git a/src/eval/ops.rs b/src/eval/ops.rs
index 2bf1c189..f79b2bf6 100644
--- a/src/eval/ops.rs
+++ b/src/eval/ops.rs
@@ -150,12 +150,12 @@ pub fn mul(lhs: Value, rhs: Value) -> StrResult<Value> {
         (Fractional(a), Float(b)) => Fractional(a * b),
         (Int(a), Fractional(b)) => Fractional(a as f64 * b),
 
-        (Str(a), Int(b)) => Str(a.repeat(b.max(0) as usize)),
-        (Int(a), Str(b)) => Str(b.repeat(a.max(0) as usize)),
-        (Array(a), Int(b)) => Array(a.repeat(b.max(0) as usize)),
-        (Int(a), Array(b)) => Array(b.repeat(a.max(0) as usize)),
-        (Template(a), Int(b)) => Template(a.repeat(b.max(0) as usize)),
-        (Int(a), Template(b)) => Template(b.repeat(a.max(0) as usize)),
+        (Str(a), Int(b)) => Str(a.repeat(b)?),
+        (Int(a), Str(b)) => Str(b.repeat(a)?),
+        (Array(a), Int(b)) => Array(a.repeat(b)?),
+        (Int(a), Array(b)) => Array(b.repeat(a)?),
+        (Template(a), Int(b)) => Template(a.repeat(b)?),
+        (Int(a), Template(b)) => Template(b.repeat(a)?),
 
         (a, b) => mismatch!("cannot multiply {} with {}", a, b),
     })
diff --git a/src/eval/template.rs b/src/eval/template.rs
index 29b5663d..4e20b8f8 100644
--- a/src/eval/template.rs
+++ b/src/eval/template.rs
@@ -1,9 +1,11 @@
 use std::collections::HashMap;
+use std::convert::TryFrom;
 use std::fmt::{self, Debug, Formatter};
 use std::ops::{Add, AddAssign, Deref};
 use std::rc::Rc;
 
 use super::Value;
+use crate::diag::StrResult;
 use crate::exec::ExecContext;
 use crate::syntax::{Expr, SyntaxTree};
 use crate::util::EcoString;
@@ -26,10 +28,15 @@ impl Template {
     }
 
     /// Repeat this template `n` times.
-    pub fn repeat(&self, n: usize) -> Self {
-        let len = self.nodes.len().checked_mul(n).expect("capacity overflow");
-        let nodes = self.iter().cloned().cycle().take(len).collect();
-        Self { nodes: Rc::new(nodes) }
+    pub fn repeat(&self, n: i64) -> StrResult<Self> {
+        let count = usize::try_from(n)
+            .ok()
+            .and_then(|n| self.nodes.len().checked_mul(n))
+            .ok_or_else(|| format!("cannot repeat this template {} times", n))?;
+
+        Ok(Self {
+            nodes: Rc::new(self.iter().cloned().cycle().take(count).collect()),
+        })
     }
 }
 
@@ -123,7 +130,7 @@ pub struct TemplateTree {
 /// The raw pointers point into the expressions contained in some
 /// [`SyntaxTree`]. Since the lifetime is erased, the tree could go out of scope
 /// while the hash map still lives. Although this could lead to lookup panics,
-/// it is not unsafe since the pointers are never dereferenced.
+/// it is safe since the pointers are never dereferenced.
 pub type ExprMap = HashMap<*const Expr, Value>;
 
 /// A reference-counted dynamic template node that can implement custom