From e487c4e504aeb72b28f759f5aad38adccc895ed9 Mon Sep 17 00:00:00 2001
From: Mark Sirek <sirek@cockroachlabs.com>
Date: Wed, 21 Dec 2022 15:13:33 -0800
Subject: [PATCH] main: update optgen Op alias to match current memo group
 member

Fixes #94042

This fixes a slightly confusing feature of optgen where an alias in
a match rule is assigned to the first expression in the group, and
may be of a different expression type than the allowed Ops in the
match rule.

Example:
```
[GenerateStreamingGroupByLimitOrderingHint, Explore]
(Limit
    $aggregation:(GroupBy | DistinctOn
```
Given this match rule, it looks like `$aggregation` can only be a
GroupBy or DistinctOn expression. But in #94042, it is sometimes a
Select expression. The variable is later used to generate a new
expression with the assumption it is one of the allowed Op types,
causing a panic.

This is fixed by updating the optgen rule gen code to keep the alias
up-to-date with the current memo group member being examined in the
match rule, as it cycles through all members in the group.

Release note: None
---
 pkg/sql/opt/optgen/cmd/optgen/rule_gen.go       | 3 +++
 pkg/sql/opt/optgen/cmd/optgen/testdata/explorer | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/pkg/sql/opt/optgen/cmd/optgen/rule_gen.go b/pkg/sql/opt/optgen/cmd/optgen/rule_gen.go
index c2870b58bec1..a9f01d25cb8d 100644
--- a/pkg/sql/opt/optgen/cmd/optgen/rule_gen.go
+++ b/pkg/sql/opt/optgen/cmd/optgen/rule_gen.go
@@ -492,6 +492,9 @@ func (g *newRuleGen) genMatchNameAndChildren(
 			g.w.nest("else {\n")
 			g.w.writeIndent("_member = _member.NextExpr()\n")
 			g.w.unnest("}\n")
+			if context.untypedAlias != "" {
+				g.w.write("%s = _member\n", context.untypedAlias)
+			}
 
 			if match == g.innerExploreMatch {
 				// This is the innermost match expression, so skip over already
diff --git a/pkg/sql/opt/optgen/cmd/optgen/testdata/explorer b/pkg/sql/opt/optgen/cmd/optgen/testdata/explorer
index 46323f5a11e9..1bf20fbda5e5 100644
--- a/pkg/sql/opt/optgen/cmd/optgen/testdata/explorer
+++ b/pkg/sql/opt/optgen/cmd/optgen/testdata/explorer
@@ -302,6 +302,7 @@ func (_e *explorer) exploreLimit(
 			} else {
 				_member = _member.NextExpr()
 			}
+			scan = _member
 			if !_partlyExplored || _ord >= _state.start {
 				_scan, _ := _member.(*memo.ScanExpr)
 				if _scan != nil {
@@ -420,6 +421,7 @@ func (_e *explorer) exploreLimit(
 			} else {
 				_member = _member.NextExpr()
 			}
+			input = _member
 			if !_partlyExplored || _ord >= _state.start {
 				_scan, _ := _member.(*memo.ScanExpr)
 				if _scan == nil {
@@ -652,6 +654,7 @@ func (_e *explorer) exploreUnion(
 			} else {
 				_member = _member.NextExpr()
 			}
+			left = _member
 			if !_partlyExplored || _ord >= _state.start {
 				_union, _ := _member.(*memo.UnionExpr)
 				if _union != nil {