Fix set inplace operation bug

kth496 · kth496 · commit cb66cc3cbad6 · 2022-07-20T01:30:48.000+09:00
diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs
@@ -375,10 +375,11 @@ impl PySetInner {
 
     fn difference_update(
         &self,
-        others: impl std::iter::Iterator<Item = ArgIterable>,
+        others: impl std::iter::Iterator<Item = PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
         for iterable in others {
+            let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)?;
             for item in iterable.iter(vm)? {
                 self.content.delete_if_exists(vm, &*item?)?;
             }
@@ -388,11 +389,12 @@ impl PySetInner {
 
     fn symmetric_difference_update(
         &self,
-        others: impl std::iter::Iterator<Item = ArgIterable>,
+        others: impl std::iter::Iterator<Item = PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
         for iterable in others {
             // We want to remove duplicates in iterable
+            let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)?;
             let iterable_set = Self::from_iter(iterable.iter(vm)?, vm)?;
             for item in iterable_set.elements() {
                 self.content.delete_or_insert(vm, &item, ())?;
@@ -681,43 +683,66 @@ impl PySet {
     }
 
     #[pymethod]
-    fn difference_update(&self, others: PosArgs<ArgIterable>, vm: &VirtualMachine) -> PyResult<()> {
-        self.inner.difference_update(others.into_iter(), vm)?;
+    fn difference_update(zelf: PyRef<Self>, others: PosArgs<PyObjectRef>, vm: &VirtualMachine) -> PyResult<()> {
+        let arguments = others.into_vec();
+        if arguments.len() == 1 {
+            if let Some(iterable) = arguments.first() {
+                if zelf.is(iterable) {
+                    zelf.inner.clear();
+                    return Ok(());
+                }
+            }
+        } else if arguments.is_empty() {
+            return Ok(());
+        }
+        zelf.inner.difference_update(arguments.into_iter(), vm)?;
         Ok(())
     }
 
     #[pymethod(magic)]
     fn isub(zelf: PyRef<Self>, set: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        if !set.class().is(vm.ctx.types.set_type) {
+            return Err(vm.new_type_error("Type error".to_owned()));
+        }
         if zelf.is(&set) {
             zelf.inner.clear();
             return Ok(zelf);
         }
-        let set = set.try_into_value::<AnySet>(vm)?;
-        zelf.inner
-            .difference_update(set.into_iterable_iter(vm)?, vm)?;
+        zelf.inner.difference_update(std::iter::once(set), vm)?;
         Ok(zelf)
     }
 
     #[pymethod]
     fn symmetric_difference_update(
-        &self,
-        others: PosArgs<ArgIterable>,
+        zelf: PyRef<Self>,
+        others: PosArgs<PyObjectRef>,
         vm: &VirtualMachine,
     ) -> PyResult<()> {
-        self.inner
-            .symmetric_difference_update(others.into_iter(), vm)?;
+        let arguments = others.into_vec();
+        if arguments.len() == 1 {
+            if let Some(iterable) = arguments.first() {
+                if zelf.is(iterable) {
+                    zelf.inner.clear();
+                    return Ok(());
+                }
+            }
+        } else if arguments.is_empty() {
+            return Ok(());
+        }
+        zelf.inner.symmetric_difference_update(arguments.into_iter(), vm)?;
         Ok(())
     }
 
     #[pymethod(magic)]
     fn ixor(zelf: PyRef<Self>, set: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        if !set.class().is(vm.ctx.types.set_type) {
+            return Err(vm.new_type_error("Type error".to_owned()));
+        }
         if zelf.is(&set) {
             zelf.inner.clear();
             return Ok(zelf);
         }
-        let set = set.try_into_value::<AnySet>(vm)?;
-        zelf.inner
-            .symmetric_difference_update(set.into_iterable_iter(vm)?, vm)?;
+        zelf.inner.symmetric_difference_update(std::iter::once(set), vm)?;
         Ok(zelf)
     }
 
