From de6ba02cdaec6578f8175098bfe1f20589fb4be4 Mon Sep 17 00:00:00 2001
From: Karol Sobczak <napewnotrafi@gmail.com>
Date: Fri, 14 Jun 2024 17:05:35 +0200
Subject: [PATCH] Make sure expression type does not change after processing
 with rule

---
 .../io/trino/sql/ir/optimizer/IrExpressionOptimizer.java  | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/core/trino-main/src/main/java/io/trino/sql/ir/optimizer/IrExpressionOptimizer.java b/core/trino-main/src/main/java/io/trino/sql/ir/optimizer/IrExpressionOptimizer.java
index c37fc08684938..5dde2a5e14ca5 100644
--- a/core/trino-main/src/main/java/io/trino/sql/ir/optimizer/IrExpressionOptimizer.java
+++ b/core/trino-main/src/main/java/io/trino/sql/ir/optimizer/IrExpressionOptimizer.java
@@ -73,6 +73,8 @@
 import java.util.Map;
 import java.util.Optional;
 
+import static com.google.common.base.Preconditions.checkState;
+
 public class IrExpressionOptimizer
 {
     private final List<IrOptimizerRule> rules;
@@ -286,6 +288,12 @@ private Optional<Expression> applyRules(Expression expression, Session session,
         for (IrOptimizerRule rule : rules) {
             Optional<Expression> optimized = rule.apply(expression, session, bindings);
             if (optimized.isPresent()) {
+                checkState(
+                        expression.type().equals(optimized.get().type()),
+                        "Rule %s changed expression type from %s to %s",
+                        rule.getClass().getSimpleName(),
+                        expression.type(),
+                        optimized.get().type());
                 expression = optimized.get();
                 changed = true;
             }