From a74c2c6792e8828ab811a908b237bf735b832eb6 Mon Sep 17 00:00:00 2001
From: Andrew Gallant <andrew@astral.sh>
Date: Thu, 23 May 2024 13:25:42 -0400
Subject: [PATCH] pep508: add `and` and `or` mutators to `MarkerTree`

While this could be done by callers since the representation
of `MarkerTree` is public, they are just annoying enough to do
that I think it makes sense to provide them on `MarkerTree`
itself.

These could also be improved in the future to do even more
flattening of conjunctions/disjunctions (or perhaps even
more robust simplification). But for now, some basic flattening
is good enough.

These routines will be used to combine marker expressions when
merging forked resolutions.
---
 crates/pep508-rs/src/marker.rs | 36 ++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/crates/pep508-rs/src/marker.rs b/crates/pep508-rs/src/marker.rs
index c576ef222..788387e4b 100644
--- a/crates/pep508-rs/src/marker.rs
+++ b/crates/pep508-rs/src/marker.rs
@@ -1801,6 +1801,42 @@ impl MarkerTree {
             }
         }
     }
+
+    /// Combine this marker tree with the one given via a conjunction.
+    ///
+    /// This does some shallow flattening. That is, if `self` is a conjunction
+    /// already, then `tree` is added to it instead of creating a new
+    /// conjunction.
+    pub fn and(&mut self, tree: MarkerTree) {
+        match *self {
+            MarkerTree::Expression(_) | MarkerTree::Or(_) => {
+                let this = std::mem::replace(self, MarkerTree::And(vec![]));
+                *self = MarkerTree::And(vec![this]);
+            }
+            _ => {}
+        }
+        if let MarkerTree::And(ref mut exprs) = *self {
+            exprs.push(tree);
+        }
+    }
+
+    /// Combine this marker tree with the one given via a disjunction.
+    ///
+    /// This does some shallow flattening. That is, if `self` is a disjunction
+    /// already, then `tree` is added to it instead of creating a new
+    /// disjunction.
+    pub fn or(&mut self, tree: MarkerTree) {
+        match *self {
+            MarkerTree::Expression(_) | MarkerTree::And(_) => {
+                let this = std::mem::replace(self, MarkerTree::And(vec![]));
+                *self = MarkerTree::Or(vec![this]);
+            }
+            _ => {}
+        }
+        if let MarkerTree::Or(ref mut exprs) = *self {
+            exprs.push(tree);
+        }
+    }
 }
 
 impl Display for MarkerTree {