Merge pull request #3234 from youknowone/further-repr-fix

Single pass repr for set

vm/src/builtins/bytearray.rs

@@ -129,7 +129,7 @@ impl PyByteArray {
     #[pymethod(magic)]
     fn repr(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult<String> {
         let class_name = zelf.class().name();
-        let s = zelf.inner().repr(&format!("{}(", class_name), ")");
+        let s = zelf.inner().repr(Some(&class_name));
         Ok(s)
     }
 

vm/src/builtins/bytes.rs

@@ -112,7 +112,7 @@ impl SlotConstructor for PyBytes {
 impl PyBytes {
     #[pymethod(magic)]
     pub(crate) fn repr(&self) -> String {
-        self.inner.repr("", "")
+        self.inner.repr(None)
     }
 
     #[pymethod(magic)]

vm/src/builtins/set.rs

@@ -199,16 +199,41 @@ impl PySetInner {
     }
 
     fn repr(&self, class_name: Option<&str>, vm: &VirtualMachine) -> PyResult<String> {
-        let mut str_parts = Vec::with_capacity(self.content.len());
+        let mut repr_len = class_name.map_or(0, |name| name.len() + 2);
+        let mut parts = Vec::with_capacity(self.content.len());
         for key in self.elements() {
             let part = vm.to_repr(&key)?;
-            str_parts.push(part.as_str().to_owned());
+            repr_len += part.as_str().len() + 2;
+            parts.push(part);
         }
-        let inner_repr = format!("{{{}}}", str_parts.join(", "));
+        let (parts, repr_len) = (parts, repr_len);
+
+        let mut repr = String::with_capacity(repr_len);
         if let Some(name) = class_name {
-            return Ok(format!("{}({})", name, inner_repr));
-        };
-        Ok(inner_repr)
+            repr.push_str(name);
+            repr.push('(');
+        }
+        repr.push('{');
+        {
+            let mut parts_iter = parts.into_iter();
+            repr.push_str(
+                parts_iter
+                    .next()
+                    .expect("this is not called for empty set")
+                    .as_str(),
+            );
+            for part in parts_iter {
+                repr.push_str(", ");
+                repr.push_str(part.as_str());
+            }
+        }
+        repr.push('}');
+        if class_name.is_some() {
+            repr.push(')');
+        }
+        debug_assert_eq!(repr.len(), repr_len);
+
+        Ok(repr)
     }
 
     fn add(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
@@ -480,11 +505,12 @@ impl PySet {
         let s = if zelf.inner.len() == 0 {
             format!("{}()", class_name)
         } else if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
-            if class_name != "set" {
-                zelf.inner.repr(Some(&class_name), vm)?
+            let name = if class_name != "set" {
+                Some(class_name.as_str())
             } else {
-                zelf.inner.repr(None, vm)?
-            }
+                None
+            };
+            zelf.inner.repr(name, vm)?
         } else {
             format!("{}(...)", class_name)
         };

vm/src/bytesinner.rs

@@ -235,40 +235,45 @@ impl ByteInnerTranslateOptions {
 pub type ByteInnerSplitOptions<'a> = anystr::SplitArgs<'a, PyBytesInner>;
 
 impl PyBytesInner {
-    pub fn repr(&self, prefix: &str, suffix: &str) -> String {
+    pub fn repr(&self, class_name: Option<&str>) -> String {
         use std::fmt::Write;
 
-        let mut out_len = 0usize;
-        let mut squote = 0;
-        let mut dquote = 0;
+        let (quote, out_len) = {
+            let mut out_len = 0usize;
+            let mut squote = 0;
+            let mut dquote = 0;
 
-        for &ch in self.elements.iter() {
-            let incr = match ch {
-                b'\'' => {
-                    squote += 1;
-                    1
-                }
-                b'"' => {
-                    dquote += 1;
-                    1
-                }
-                b'\\' | b'\t' | b'\r' | b'\n' => 2,
-                0x20..=0x7e => 1,
-                _ => 4, // \xHH
-            };
-            // TODO: OverflowError
-            out_len = out_len.checked_add(incr).unwrap();
-        }
+            for &ch in self.elements.iter() {
+                let incr = match ch {
+                    b'\'' => {
+                        squote += 1;
+                        1
+                    }
+                    b'"' => {
+                        dquote += 1;
+                        1
+                    }
+                    b'\\' | b'\t' | b'\r' | b'\n' => 2,
+                    0x20..=0x7e => 1,
+                    _ => 4, // \xHH
+                };
+                // TODO: OverflowError
+                out_len = out_len.checked_add(incr).unwrap();
+            }
 
-        let (quote, num_escaped_quotes) = anystr::choose_quotes_for_repr(squote, dquote);
-        // we'll be adding backslashes in front of the existing inner quotes
-        out_len += num_escaped_quotes;
-
-        // 3 is for b prefix + outer quotes
-        out_len += 3 + prefix.len() + suffix.len();
+            let (quote, num_escaped_quotes) = anystr::choose_quotes_for_repr(squote, dquote);
+            // we'll be adding backslashes in front of the existing inner quotes
+            out_len += num_escaped_quotes;
 
+            // 3 is for b prefix + outer quotes
+            out_len += 3 + class_name.map_or(0, |name| name.len() + 2);
+            (quote, out_len)
+        };
         let mut res = String::with_capacity(out_len);
-        res.push_str(prefix);
+        if let Some(name) = class_name {
+            res.push_str(name);
+            res.push('(');
+        }
         res.push('b');
         res.push(quote);
         for &ch in self.elements.iter() {
@@ -288,7 +293,10 @@ impl PyBytesInner {
             }
         }
         res.push(quote);
-        res.push_str(suffix);
+        if class_name.is_some() {
+            res.push(')');
+        }
+        debug_assert_eq!(res.len(), out_len);
 
         res
     }