fix stackoverflow in comparing null booleans of unit length

pola-rs · Jan 30, 2022 · 49e5cc1 · 49e5cc1
1 parent d7da65b
commit 49e5cc1
Showing 1 changed file with 85 additions and 48 deletions.
diff --git a/polars/polars-core/src/chunked_array/bitwise.rs b/polars/polars-core/src/chunked_array/bitwise.rs
@@ -107,22 +107,29 @@ impl BitOr for &BooleanChunked {
     type Output = BooleanChunked;
 
     fn bitor(self, rhs: Self) -> Self::Output {
-        if self.len() == 1 {
-            return match self.get(0) {
-                Some(true) => BooleanChunked::full(self.name(), true, rhs.len()),
-                Some(false) => {
-                    let mut rhs = rhs.clone();
-                    rhs.rename(self.name());
-                    rhs
-                }
-                None => &self.expand_at_index(0, rhs.len()) | rhs,
-            };
-        } else if rhs.len() == 1 {
-            return match rhs.get(0) {
-                Some(true) => BooleanChunked::full(self.name(), true, self.len()),
-                Some(false) => self.clone(),
-                None => &rhs.expand_at_index(0, self.len()) | self,
-            };
+        match (self.len(), rhs.len()) {
+            // make sure that we fall through if both are equal unit lengths
+            // otherwise we stackoverflow
+            (1, 1) => {}
+            (1, _) => {
+                return match self.get(0) {
+                    Some(true) => BooleanChunked::full(self.name(), true, rhs.len()),
+                    Some(false) => {
+                        let mut rhs = rhs.clone();
+                        rhs.rename(self.name());
+                        rhs
+                    }
+                    None => &self.expand_at_index(0, rhs.len()) | rhs,
+                };
+            }
+            (_, 1) => {
+                return match rhs.get(0) {
+                    Some(true) => BooleanChunked::full(self.name(), true, self.len()),
+                    Some(false) => self.clone(),
+                    None => &rhs.expand_at_index(0, self.len()) | self,
+                };
+            }
+            _ => {}
         }
 
         let (lhs, rhs) = align_chunks_binary(self, rhs);
@@ -150,26 +157,33 @@ impl BitXor for &BooleanChunked {
     type Output = BooleanChunked;
 
     fn bitxor(self, rhs: Self) -> Self::Output {
-        if self.len() == 1 {
-            return match self.get(0) {
-                Some(true) => {
-                    let mut rhs = rhs.not();
-                    rhs.rename(self.name());
-                    rhs
-                }
-                Some(false) => {
-                    let mut rhs = rhs.clone();
-                    rhs.rename(self.name());
-                    rhs
-                }
-                None => &self.expand_at_index(0, rhs.len()) | rhs,
-            };
-        } else if rhs.len() == 1 {
-            return match rhs.get(0) {
-                Some(true) => self.not(),
-                Some(false) => self.clone(),
-                None => &rhs.expand_at_index(0, self.len()) | self,
-            };
+        match (self.len(), rhs.len()) {
+            // make sure that we fall through if both are equal unit lengths
+            // otherwise we stackoverflow
+            (1, 1) => {}
+            (1, _) => {
+                return match self.get(0) {
+                    Some(true) => {
+                        let mut rhs = rhs.not();
+                        rhs.rename(self.name());
+                        rhs
+                    }
+                    Some(false) => {
+                        let mut rhs = rhs.clone();
+                        rhs.rename(self.name());
+                        rhs
+                    }
+                    None => &self.expand_at_index(0, rhs.len()) | rhs,
+                };
+            }
+            (_, 1) => {
+                return match rhs.get(0) {
+                    Some(true) => self.not(),
+                    Some(false) => self.clone(),
+                    None => &rhs.expand_at_index(0, self.len()) | self,
+                };
+            }
+            _ => {}
         }
 
         let (l, r) = align_chunks_binary(self, rhs);
@@ -211,18 +225,25 @@ impl BitAnd for &BooleanChunked {
     type Output = BooleanChunked;
 
     fn bitand(self, rhs: Self) -> Self::Output {
-        if self.len() == 1 {
-            return match self.get(0) {
-                Some(true) => rhs.clone(),
-                Some(false) => BooleanChunked::full(self.name(), false, rhs.len()),
-                None => &self.expand_at_index(0, rhs.len()) & rhs,
-            };
-        } else if rhs.len() == 1 {
-            return match rhs.get(0) {
-                Some(true) => self.clone(),
-                Some(false) => BooleanChunked::full(self.name(), false, self.len()),
-                None => self & &rhs.expand_at_index(0, self.len()),
-            };
+        match (self.len(), rhs.len()) {
+            // make sure that we fall through if both are equal unit lengths
+            // otherwise we stackoverflow
+            (1, 1) => {}
+            (1, _) => {
+                return match self.get(0) {
+                    Some(true) => rhs.clone(),
+                    Some(false) => BooleanChunked::full(self.name(), false, rhs.len()),
+                    None => &self.expand_at_index(0, rhs.len()) & rhs,
+                };
+            }
+            (_, 1) => {
+                return match rhs.get(0) {
+                    Some(true) => self.clone(),
+                    Some(false) => BooleanChunked::full(self.name(), false, self.len()),
+                    None => self & &rhs.expand_at_index(0, self.len()),
+                };
+            }
+            _ => {}
         }
 
         let (lhs, rhs) = align_chunks_binary(self, rhs);
@@ -274,3 +295,19 @@ macro_rules! impl_floats {
 
 impl_floats!(Float64Chunked);
 impl_floats!(Float32Chunked);
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn guard_so_issue_2494() {
+        // this cause a stack overflow
+        let a = BooleanChunked::new("a", [None]);
+        let b = BooleanChunked::new("b", [None]);
+
+        assert_eq!((&a).bitand(&b).null_count(), 1);
+        assert_eq!((&a).bitor(&b).null_count(), 1);
+        assert_eq!((&a).bitxor(&b).null_count(), 1);
+    }
+}