Array and dictionary indexing

2 files changed, 48 insertions, 9 deletions

diff --git a/src/util/eco.rs b/src/util/eco.rs
index 00f87872..2e326f32 100644
--- a/src/util/eco.rs
+++ b/src/util/eco.rs
@@ -2,11 +2,14 @@
 
 use std::borrow::Borrow;
 use std::cmp::Ordering;
+use std::convert::TryFrom;
 use std::fmt::{self, Debug, Display, Formatter, Write};
 use std::hash::{Hash, Hasher};
 use std::ops::{Add, AddAssign, Deref};
 use std::rc::Rc;
 
+use crate::diag::StrResult;
+
 /// An economical string with inline storage and clone-on-write value semantics.
 #[derive(Clone)]
 pub struct EcoString(Repr);
@@ -143,21 +146,25 @@ impl EcoString {
     }
 
     /// Repeat this string `n` times.
-    pub fn repeat(&self, n: usize) -> Self {
+    pub fn repeat(&self, n: i64) -> StrResult<Self> {
+        let (n, new) = usize::try_from(n)
+            .ok()
+            .and_then(|n| self.len().checked_mul(n).map(|new| (n, new)))
+            .ok_or_else(|| format!("cannot repeat this string {} times", n))?;
+
         if let Repr::Small { buf, len } = &self.0 {
             let prev = usize::from(*len);
-            let new = prev.saturating_mul(n);
             if new <= LIMIT {
                 let src = &buf[.. prev];
                 let mut buf = [0; LIMIT];
                 for i in 0 .. n {
                     buf[prev * i .. prev * (i + 1)].copy_from_slice(src);
                 }
-                return Self(Repr::Small { buf, len: new as u8 });
+                return Ok(Self(Repr::Small { buf, len: new as u8 }));
             }
         }
 
-        self.as_str().repeat(n).into()
+        Ok(self.as_str().repeat(n).into())
     }
 }
 
@@ -216,6 +223,13 @@ impl Deref for EcoString {
 
     fn deref(&self) -> &str {
         match &self.0 {
+            // Safety:
+            // The buffer contents stem from correct UTF-8 sources:
+            // - Valid ASCII characters
+            // - Other string slices
+            // - Chars that were encoded with char::encode_utf8
+            // Furthermore, we still do the bounds-check on the len in case
+            // it gets corrupted somehow.
             Repr::Small { buf, len } => unsafe {
                 std::str::from_utf8_unchecked(&buf[.. usize::from(*len)])
             },
@@ -424,13 +438,17 @@ mod tests {
     #[test]
     fn test_str_repeat() {
         // Test with empty string.
-        assert_eq!(EcoString::new().repeat(0), "");
-        assert_eq!(EcoString::new().repeat(100), "");
+        assert_eq!(EcoString::new().repeat(0).unwrap(), "");
+        assert_eq!(EcoString::new().repeat(100).unwrap(), "");
 
         // Test non-spilling and spilling case.
         let v = EcoString::from("abc");
-        assert_eq!(v.repeat(0), "");
-        assert_eq!(v.repeat(3), "abcabcabc");
-        assert_eq!(v.repeat(5), "abcabcabcabcabc");
+        assert_eq!(v.repeat(0).unwrap(), "");
+        assert_eq!(v.repeat(3).unwrap(), "abcabcabc");
+        assert_eq!(v.repeat(5).unwrap(), "abcabcabcabcabc");
+        assert_eq!(
+            v.repeat(-1).unwrap_err(),
+            "cannot repeat this string -1 times",
+        );
     }
 }
diff --git a/src/util/mod.rs b/src/util/mod.rs
index dc400af8..309c1241 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -4,6 +4,7 @@ mod eco;
 
 pub use eco::EcoString;
 
+use std::cell::RefMut;
 use std::cmp::Ordering;
 use std::ops::Range;
 use std::path::{Component, Path, PathBuf};
@@ -153,3 +154,23 @@ impl PathExt for Path {
         out
     }
 }
+
+/// Additional methods for [`RefMut`].
+pub trait RefMutExt<'a, T> {
+    fn try_map<U, F, E>(orig: Self, f: F) -> Result<RefMut<'a, U>, E>
+    where
+        F: FnOnce(&mut T) -> Result<&mut U, E>;
+}
+
+impl<'a, T> RefMutExt<'a, T> for RefMut<'a, T> {
+    fn try_map<U, F, E>(mut orig: Self, f: F) -> Result<RefMut<'a, U>, E>
+    where
+        F: FnOnce(&mut T) -> Result<&mut U, E>,
+    {
+        // Taken from here:
+        // https://github.com/rust-lang/rust/issues/27746#issuecomment-172899746
+        f(&mut orig)
+            .map(|new| new as *mut U)
+            .map(|raw| RefMut::map(orig, |_| unsafe { &mut *raw }))
+    }
+}