From d0d6fdfe7faf265a40de8dc4208e29962028935b Mon Sep 17 00:00:00 2001
From: Eric Kidd <git@randomhacks.net>
Date: Mon, 30 Oct 2023 14:37:38 -0400
Subject: [PATCH] Infer set operations

---
 src/infer.rs | 17 ++++++++++++++++-
 src/types.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/infer.rs b/src/infer.rs
index f862003..f8da877 100644
--- a/src/infer.rs
+++ b/src/infer.rs
@@ -231,7 +231,22 @@ impl InferTypes for ast::QueryExpression {
                 }
                 query.infer_types(&scope)
             }
-            ast::QueryExpression::SetOperation { .. } => Err(nyi(self, "set operation")),
+            ast::QueryExpression::SetOperation {
+                left,
+                set_operator,
+                right,
+            } => {
+                let left_ty = left.infer_types(scope)?.0;
+                let right_ty = right.infer_types(scope)?.0;
+                let result_ty = left_ty.common_supertype(&right_ty).ok_or_else(|| {
+                    Error::annotated(
+                        format!("cannot combine {} and {}", left_ty, right_ty),
+                        set_operator.span(),
+                        "incompatible types",
+                    )
+                })?;
+                Ok((result_ty, scope.clone()))
+            }
         }
     }
 }
diff --git a/src/types.rs b/src/types.rs
index 10e6ad5..1f54e41 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -778,6 +778,30 @@ impl TableType {
         Ok(())
     }
 
+    /// Compute the least common supertype of two types. Returns `None` if the
+    /// only common super type would be top (⊤), which isn't part of our type
+    /// system.
+    ///
+    /// For some nice theoretical terminology, see [this
+    /// page](https://orc.csres.utexas.edu/documentation/html/refmanual/ref.types.subtyping.html).
+    pub fn common_supertype<'a>(&'a self, other: &'a Self) -> Option<Self> {
+        // Make sure we have the same number of columns.
+        if self.columns.len() != other.columns.len() {
+            return None;
+        }
+
+        // For each column, see if we can merge the column definitions.
+        let mut columns = Vec::new();
+        for (a, b) in self.columns.iter().zip(&other.columns) {
+            if let Some(column) = a.common_supertype(b) {
+                columns.push(column);
+            } else {
+                return None;
+            }
+        }
+        Some(TableType { columns })
+    }
+
     /// Expect this table to be creatable, i.e., that it contains no aggregate
     /// columns or uninhabited types.
     pub fn expect_creatable(&self, spanned: &dyn Spanned) -> Result<()> {
@@ -860,6 +884,26 @@ impl ColumnType {
                 _ => true,
             }
     }
+
+    /// Merge two column types, if possible. Returns `None` if the only common
+    /// super type would be top (⊤), which isn't part of our type system.
+    pub fn common_supertype<'a>(&'a self, other: &'a Self) -> Option<Self> {
+        // If we have two names, they must match. If we have one or zero names,
+        // do the best we can.
+        let name = match (&self.name, &other.name) {
+            (Some(a), Some(b)) if a == b => Some(a.clone()),
+            (Some(_), Some(_)) => None,
+            (Some(a), None) => Some(a.clone()),
+            (None, Some(b)) => Some(b.clone()),
+            (None, None) => None,
+        };
+        let ty = self.ty.common_supertype(&other.ty)?;
+        Some(ColumnType {
+            name,
+            ty,
+            not_null: self.not_null && other.not_null,
+        })
+    }
 }
 
 impl fmt::Display for ColumnType {