fix: ensure coalesce handles optional/unknown/null properly

Author: sgulseth

src/typeEvaluator/functions.ts

@@ -1,5 +1,6 @@
 /* eslint-disable max-statements */
 import type {FuncCallNode} from '../nodeTypes'
+import {optimizeUnions} from './optimizations'
 import {Scope} from './scope'
 import {walk} from './typeEvaluate'
 import {mapNode, nullUnion} from './typeHelpers'
@@ -175,12 +176,32 @@ export function handleFuncCallNode(node: FuncCallNode, scope: Scope): TypeNode {
       const typeNodes: TypeNode[] = []
       let canBeNull = true
       for (const arg of node.args) {
-        const type = walk({node: arg, scope})
-        typeNodes.push(unionWithoutNull(type))
+        const argNode = optimizeUnions(walk({node: arg, scope}))
+
+        // Check if all types are null
+        const allNull =
+          argNode.type === 'null' ||
+          (argNode.type === 'union' && argNode.of.every((t) => t.type === 'null'))
+
+        // Can the argument be null, if all is null, unknown, or if its a union with at least one null or unknown
         canBeNull =
-          type.type === 'null' || (type.type === 'union' && type.of.some((t) => t.type === 'null'))
+          allNull ||
+          argNode.type === 'unknown' ||
+          (argNode.type === 'union' &&
+            argNode.of.some((t) => t.type === 'null' || t.type === 'unknown'))
+
+        // As long as some type is not null or unknown, we add it to the union, but skip nulls
+        if (!allNull) {
+          typeNodes.push(unionWithoutNull(argNode))
+        }
+
+        // If we have a type that can't be null, we can break.
+        if (!canBeNull) {
+          break
+        }
       }
 
+      // If the last argument can be null, we add null to the union
       if (canBeNull) {
         typeNodes.push({type: 'null'} satisfies NullTypeNode)
       }

tap-snapshots/test/typeEvaluate.test.ts.test.cjs

@@ -257,54 +257,6 @@ Object {
 }
 `
 
-exports[`test/typeEvaluate.test.ts TAP coalesce only > must match snapshot 1`] = `
-Object {
-  "of": Object {
-    "attributes": Object {
-      "maybe": Object {
-        "type": "objectAttribute",
-        "value": Object {
-          "of": Array [
-            Object {
-              "attributes": Object {
-                "subfield": Object {
-                  "type": "objectAttribute",
-                  "value": Object {
-                    "type": "string",
-                  },
-                },
-              },
-              "type": "object",
-            },
-            Object {
-              "type": "null",
-            },
-          ],
-          "type": "union",
-        },
-      },
-      "name": Object {
-        "type": "objectAttribute",
-        "value": Object {
-          "of": Array [
-            Object {
-              "type": "string",
-            },
-            Object {
-              "type": "string",
-              "value": "unknown",
-            },
-          ],
-          "type": "union",
-        },
-      },
-    },
-    "type": "object",
-  },
-  "type": "array",
-}
-`
-
 exports[`test/typeEvaluate.test.ts TAP coalesce with projection > must match snapshot 1`] = `
 Object {
   "of": Array [

test/typeEvaluate.test.ts

@@ -1623,12 +1623,68 @@ t.test('with conditional splat', (t) => {
 
 t.test('coalesce only', async (t) => {
   const query = `*[_type == "author"]{
-          "name": coalesce(name, "unknown"),
-          "maybe": coalesce(optionalObject, dontExists)
+          "name": coalesce(name, "unknown"), // should be string, since name is not optional
+          "maybe": coalesce(optionalObject, dontExists), // should be object or null
+          "multiple": coalesce(optionalAge, missing, "foo"), // should be either number or "foo"
+          "allMissing": coalesce(missing, missing2, missing3), // should be null, since all are missing
+          "fallback": coalesce(missing, missing2, missing3, "fallback"), // should be "fallback"
         }`
   const ast = parse(query)
   const res = typeEvaluate(ast, schemas)
-  t.matchSnapshot(res)
+  t.strictSame(res, {
+    type: 'array',
+    of: {
+      type: 'object',
+      attributes: {
+        name: {
+          type: 'objectAttribute',
+          value: {
+            type: 'string',
+          },
+        },
+        maybe: {
+          type: 'objectAttribute',
+          value: {
+            type: 'union',
+            of: [
+              {
+                type: 'object',
+                attributes: {
+                  subfield: {
+                    type: 'objectAttribute',
+                    value: {
+                      type: 'string',
+                    },
+                  },
+                },
+              },
+              {type: 'null'},
+            ],
+          },
+        },
+        multiple: {
+          type: 'objectAttribute',
+          value: {
+            type: 'union',
+            of: [{type: 'number'}, {type: 'string', value: 'foo'}],
+          },
+        },
+        allMissing: {
+          type: 'objectAttribute',
+          value: {
+            type: 'null',
+          },
+        },
+        fallback: {
+          type: 'objectAttribute',
+          value: {
+            type: 'string',
+            value: 'fallback',
+          },
+        },
+      },
+    },
+  } satisfies TypeNode)
   t.end()
 })