schemachanger: Fixed a crash when creating UDFs whose body casts to UDT

Fixes cockroachdb#104242

Release note (bug fix): Fixed a bug when creating UDFs whose body
contains a ShortCast operator. E.g.
`CREATE FUNCTION f() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 FROM t WHERE NULL::typ IN () $$;`

pkg/sql/logictest/testdata/logic_test/udf

@@ -3566,3 +3566,19 @@ query T
 SELECT f104927()
 ----
 [{"i": 1, "s": "foo"}]
+
+subtest 104242
+
+statement ok
+CREATE TABLE t_104242 (i INT PRIMARY KEY);
+
+statement ok
+CREATE TYPE typ_104242 AS ENUM ('a', 'b', 'c');
+
+# We've only fixed issue with declarative schema changer.
+skipif config local-legacy-schema-changer
+skipif config local-mixed-22.2-23.1
+statement ok
+CREATE FUNCTION f_104242() RETURNS INT LANGUAGE SQL AS $$ SELECT 1 FROM t_104242 WHERE NULL::typ_104242 IN () $$;
+
+subtest end

pkg/sql/schemachanger/scbuild/BUILD.bazel

@@ -33,6 +33,7 @@ go_library(
         "//pkg/sql/catalog/typedesc",
         "//pkg/sql/faketreeeval",
         "//pkg/sql/parser",
+        "//pkg/sql/parser/statements",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/privilege",

pkg/sql/schemachanger/scbuild/ast_annotator.go

@@ -12,6 +12,9 @@ package scbuild
 
 import (
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scbuild/internal/scbuildstmt"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/errors"
@@ -35,6 +38,13 @@ func newAstAnnotator(original tree.Statement) (*astAnnotator, error) {
 	if err != nil {
 		return nil, err
 	}
+	// For `CREATE FUNCTION` stmt, overwrite `NumAnnotations` to be the largest of
+	// parsing the node and parsing all its body statements, to avoid annotation
+	// index-out-of-bound issue during formatting (https://github.com/cockroachdb/cockroach/issues/104242).
+	if _, ok := statement.AST.(*tree.CreateFunction); ok {
+		statement.NumAnnotations = numberAnnotationsOfCreateFunction(statement)
+	}
+
 	return &astAnnotator{
 		nonExistentNames: map[*tree.TableName]struct{}{},
 		statement:        statement.AST,
@@ -88,3 +98,39 @@ func (ann *astAnnotator) ValidateAnnotations() {
 		}))
 	f.FormatNode(ann.statement)
 }
+
+// numberAnnotationsOfCreateFunction parse each function body statement
+// and return the largest NumAnnotations in `cfStmt` and in each of those
+// parsed function body statement.
+func numberAnnotationsOfCreateFunction(
+	cfStmt statements.Statement[tree.Statement],
+) tree.AnnotationIdx {
+	if cf, ok := cfStmt.AST.(*tree.CreateFunction); !ok {
+		panic(errors.AssertionFailedf("statement is not CREATE FUNCTION"))
+	} else {
+		var funcBodyFound bool
+		var funcBodyStr string
+		for _, option := range cf.Options {
+			switch opt := option.(type) {
+			case tree.FunctionBodyStr:
+				funcBodyFound = true
+				funcBodyStr = string(opt)
+			}
+		}
+		if !funcBodyFound {
+			panic(pgerror.New(pgcode.InvalidFunctionDefinition, "no function body specified"))
+		}
+		// ret = max(cfSmt.NumAnnotations, max(funcBodyStmt.NumAnnotations within funcBodyStmts))
+		ret := cfStmt.NumAnnotations
+		funcBodyStmts, err := parser.Parse(funcBodyStr)
+		if err != nil {
+			panic(err)
+		}
+		for _, funcBodyStmt := range funcBodyStmts {
+			if funcBodyStmt.NumAnnotations > ret {
+				ret = funcBodyStmt.NumAnnotations
+			}
+		}
+		return ret
+	}
+}