Unify string and sequence slicing operations

Daniel Watkins · Daniel Watkins · commit d9cd236e7164 · 2018-08-08T12:44:51.000-04:00
diff --git a/vm/src/objsequence.rs b/vm/src/objsequence.rs
@@ -13,7 +13,28 @@ pub fn get_pos(sequence_length: usize, p: i32) -> usize {
     }
 }
 
-fn get_slice_items(l: &Vec<PyObjectRef>, slice: &PyObjectRef) -> Vec<PyObjectRef> {
+pub trait PySliceableSequence {
+    fn do_slice(&self, start: usize, stop: usize) -> Self;
+    fn do_stepped_slice(&self, start: usize, stop: usize, step: usize) -> Self;
+    fn len(&self) -> usize;
+}
+
+impl PySliceableSequence for Vec<PyObjectRef> {
+    fn do_slice(&self, start: usize, stop: usize) -> Self {
+        self[start..stop].to_vec()
+    }
+    fn do_stepped_slice(&self, start: usize, stop: usize, step: usize) -> Self {
+        self[start..stop].iter().step_by(step).cloned().collect()
+    }
+    fn len(&self) -> usize {
+        self.len()
+    }
+}
+
+pub fn get_slice_items<S>(l: &S, slice: &PyObjectRef) -> S
+where
+    S: PySliceableSequence,
+{
     // TODO: we could potentially avoid this copy and use slice
     match &(slice.borrow()).kind {
         PyObjectKind::Slice { start, stop, step } => {
@@ -26,16 +47,12 @@ fn get_slice_items(l: &Vec<PyObjectRef>, slice: &PyObjectRef) -> Vec<PyObjectRef
                 &None => l.len() as usize,
             };
             match step {
-                &None | &Some(1) => l[start..stop].to_vec(),
+                &None | &Some(1) => l.do_slice(start, stop),
                 &Some(num) => {
                     if num < 0 {
                         unimplemented!("negative step indexing not yet supported")
                     };
-                    l[start..stop]
-                        .iter()
-                        .step_by(num as usize)
-                        .cloned()
-                        .collect()
+                    l.do_stepped_slice(start, stop, num as usize)
                 }
             }
         }
diff --git a/vm/src/objstr.rs b/vm/src/objstr.rs
@@ -2,30 +2,15 @@ use super::objsequence;
 use super::pyobject::{PyObjectKind, PyObjectRef, PyResult};
 use super::vm::VirtualMachine;
 
-fn get_slice_items(value: &String, slice: &PyObjectRef) -> String {
-    match &(slice.borrow()).kind {
-        PyObjectKind::Slice { start, stop, step } => {
-            let start2: usize = match start {
-                // &Some(_) => panic!("Bad start index for string slicing {:?}", start),
-                &Some(start) => objsequence::get_pos(value.len(), start),
-                &None => 0,
-            };
-            let stop2: usize = match stop {
-                &Some(stop) => objsequence::get_pos(value.len(), stop),
-                // &Some(_) => panic!("Bad stop index for string slicing"),
-                &None => value.len() as usize,
-            };
-            match step {
-                &None | &Some(1) => value[start2..stop2].to_string(),
-                &Some(num) => {
-                    if num < 0 {
-                        unimplemented!("negative step indexing not yet supported")
-                    };
-                    value[start2..stop2].chars().step_by(num as usize).collect()
-                }
-            }
-        }
-        kind => panic!("get_slice_items called with non-slice: {:?}", kind),
+impl objsequence::PySliceableSequence for String {
+    fn do_slice(&self, start: usize, stop: usize) -> Self {
+        self[start..stop].to_string()
+    }
+    fn do_stepped_slice(&self, start: usize, stop: usize, step: usize) -> Self {
+        self[start..stop].chars().step_by(step).collect()
+    }
+    fn len(&self) -> usize {
+        self.len()
     }
 }
 
@@ -40,7 +25,7 @@ pub fn subscript(vm: &mut VirtualMachine, value: &String, b: PyObjectRef) -> PyR
             start: _,
             stop: _,
             step: _,
-        } => Ok(vm.new_str(get_slice_items(value, &b))),
+        } => Ok(vm.new_str(objsequence::get_slice_items(value, &b))),
         _ => panic!(
             "TypeError: indexing type {:?} with index {:?} is not supported (yet?)",
             value, b
