opt: make scalar list custom func accessor safer

The custom functions `FirstScalarListExpr` and `SecondScalarListExpr`
have been replaced by `ScalarExprAt`. This function is safer because it
returns ok=false if the given index is out-of-bounds instead of
panicking. This function has been added to `general_funcs.go` instead of
`comp_funcs.go` in preparation for usage outside of `comp.opt`.

Release note: None

Author: mgartner

pkg/sql/opt/norm/comp_funcs.go

@@ -76,16 +76,6 @@ func (c *CustomFuncs) NormalizeTupleEquality(left, right memo.ScalarListExpr) op
 	return result
 }
 
-// FirstScalarListExpr returns the first ScalarExpr in the given list.
-func (c *CustomFuncs) FirstScalarListExpr(list memo.ScalarListExpr) opt.ScalarExpr {
-	return list[0]
-}
-
-// SecondScalarListExpr returns the second ScalarExpr in the given list.
-func (c *CustomFuncs) SecondScalarListExpr(list memo.ScalarListExpr) opt.ScalarExpr {
-	return list[1]
-}
-
 // MakeTimeZoneFunction constructs a new timezone() function with the given zone
 // and timestamp as arguments. The type of the function result is TIMESTAMPTZ if
 // ts is of type TIMESTAMP, or TIMESTAMP if is of type TIMESTAMPTZ.

pkg/sql/opt/norm/general_funcs.go

@@ -658,6 +658,22 @@ func (c *CustomFuncs) IsLeakproof(expr memo.RelExpr) bool {
 	return expr.Relational().VolatilitySet.IsLeakproof()
 }
 
+// ----------------------------------------------------------------------
+//
+// ScalarListExpr functions
+//   General functions related to ScalarListExpr.
+//
+// ----------------------------------------------------------------------
+
+// ScalarExprAt returns the ScalarExpr in the i-th position in the given list.
+// Returns ok=false if i is out of bounds.
+func (c *CustomFuncs) ScalarExprAt(list memo.ScalarListExpr, i int) (_ opt.ScalarExpr, ok bool) {
+	if i >= 0 && i < len(list) {
+		return list[i], true
+	}
+	return nil, false
+}
+
 // ----------------------------------------------------------------------
 //
 // Ordering functions

pkg/sql/opt/norm/rules/comp.opt

@@ -220,14 +220,13 @@
     (Function $args:* $private:(FunctionPrivate "timezone"))
     $right:(ConstValue) &
         (IsTimestampTZ $right) &
-        (IsTimestamp $ts:(SecondScalarListExpr $args)) &
-        ^(IsConstValueOrGroupOfConstValues $ts)
+        (Let ($ts $tsOk):(ScalarExprAt $args 1) $tsOk) &
+        (IsTimestamp $ts) &
+        ^(IsConstValueOrGroupOfConstValues $ts) &
+        (Let ($zone $zoneOk):(ScalarExprAt $args 0) $zoneOk)
 )
 =>
-((OpName)
-    $ts
-    (MakeTimeZoneFunction (FirstScalarListExpr $args) $right)
-)
+((OpName) $ts (MakeTimeZoneFunction $zone $right))
 
 # NormalizeCmpTimeZoneFunctionTZ normalizes timezone functions within
 # comparison operators. It only matches expressions when:
@@ -249,14 +248,13 @@
     (Function $args:* $private:(FunctionPrivate "timezone"))
     $right:(ConstValue) &
         (IsTimestamp $right) &
-        (IsTimestampTZ $tz:(SecondScalarListExpr $args)) &
-        ^(IsConstValueOrGroupOfConstValues $tz)
+        (Let ($tz $tzOk):(ScalarExprAt $args 1) $tzOk) &
+        (IsTimestampTZ $tz) &
+        ^(IsConstValueOrGroupOfConstValues $tz) &
+        (Let ($zone $zoneOk):(ScalarExprAt $args 0) $zoneOk)
 )
 =>
-((OpName)
-    $tz
-    (MakeTimeZoneFunction (FirstScalarListExpr $args) $right)
-)
+((OpName) $tz (MakeTimeZoneFunction $zone $right))
 
 # FoldEqTrue replaces x = True with x.
 [FoldEqTrue, Normalize]