fix(react): infer jsx component returns

Author: chrisbbreuer

packages/dtsx/src/processor/type-inference.ts

@@ -110,6 +110,45 @@ function countOccurrences(str: string, sub: string): number {
   return count
 }
 
+function isJsxTagNameChar(code: number): boolean {
+  return (code >= 65 && code <= 90)
+    || (code >= 97 && code <= 122)
+    || (code >= 48 && code <= 57)
+    || code === 36 || code === 95 || code === 46 || code === 58 || code === 45
+}
+
+/**
+ * Detect a complete JSX element or fragment without assuming React as the JSX
+ * runtime. A matching outer close prevents generic arrows and type assertions
+ * from being classified as JSX.
+ */
+function isJsxExpression(value: string): boolean {
+  const expression = value.trim()
+  if (expression.length < 3 || expression.charCodeAt(0) !== 60) return false
+  if (expression.startsWith('<>')) return expression.endsWith('</>')
+
+  const firstTagCode = expression.charCodeAt(1)
+  const isIdentifierStart = (firstTagCode >= 65 && firstTagCode <= 90)
+    || (firstTagCode >= 97 && firstTagCode <= 122)
+    || firstTagCode === 36 || firstTagCode === 95
+  if (!isIdentifierStart) return false
+
+  let tagEnd = 2
+  while (tagEnd < expression.length && isJsxTagNameChar(expression.charCodeAt(tagEnd))) tagEnd++
+  const tagName = expression.slice(1, tagEnd)
+  const delimiter = expression.charCodeAt(tagEnd)
+  if (delimiter > 32 && delimiter !== 62 && delimiter !== 47) return false
+  if (expression.endsWith('/>')) return true
+
+  const closeStart = expression.lastIndexOf('</')
+  if (closeStart === -1) return false
+  let closeEnd = closeStart + 2
+  while (closeEnd < expression.length && isJsxTagNameChar(expression.charCodeAt(closeEnd))) closeEnd++
+  return expression.slice(closeStart + 2, closeEnd) === tagName
+    && expression.charCodeAt(closeEnd) === 62
+    && closeEnd === expression.length - 1
+}
+
 /** Collapse runs of whitespace to single spaces (no regex) */
 function collapseWhitespace(s: string): string {
   const len = s.length
@@ -169,6 +208,8 @@ export function inferNarrowType(value: unknown, isConst: boolean = false, inUnio
 
   const trimmed = value.trim()
 
+  if (isJsxExpression(trimmed)) return 'JSX.Element'
+
   // BigInt expressions (check early)
   if (trimmed.startsWith('BigInt(')) {
     return 'bigint'
@@ -536,6 +577,11 @@ export function inferFunctionBodyReturnType(body: string, isAsync: boolean = fal
           braceDepth--
         }
         else if (expressionChar === 59 && parenDepth === 0 && bracketDepth === 0 && braceDepth === 0) break
+        else if ((expressionChar === 10 || expressionChar === 13) && parenDepth === 0 && bracketDepth === 0 && braceDepth === 0) {
+          let next = i + 1
+          while (next < body.length && body.charCodeAt(next) <= 32) next++
+          if (body.startsWith('return', next) && !isWordChar(body.charCodeAt(next + 6))) break
+        }
         i++
       }
 
@@ -586,6 +632,7 @@ function collectParameterTypes(parameters: string): Map<string, string> {
 function inferBodyExpressionType(expression: string, parameterTypes: ReadonlyMap<string, string>): string {
   let value = expression.trim()
   while (hasBalancedOuterParentheses(value)) value = value.slice(1, -1).trim()
+  if (isJsxExpression(value)) return 'JSX.Element'
   if (value.startsWith('await ')) {
     const awaited = inferBodyExpressionType(value.slice(6), parameterTypes)
     return awaited.startsWith('Promise<') && awaited.endsWith('>') ? awaited.slice(8, -1) : awaited

packages/dtsx/test/react-components.test.ts

@@ -0,0 +1,53 @@
+import { describe, expect, test } from 'bun:test'
+import { processSource } from '../src/generator'
+
+describe('React and JSX component declarations', () => {
+  test('infers JSX returns for function components', () => {
+    const output = processSource(`
+      export interface ButtonProps { label: string }
+      export function Button({ label }: ButtonProps) {
+        return <button aria-label={label}>{label}</button>
+      }
+    `)
+
+    expect(output).toContain('export declare function Button({ label }: ButtonProps): JSX.Element;')
+  })
+
+  test('infers JSX returns for arrow components', () => {
+    const output = processSource(`
+      export interface BadgeProps { text: string }
+      export const Badge = ({ text }: BadgeProps) => <span>{text}</span>
+    `)
+
+    expect(output).toContain('export declare const Badge: ({ text }: BadgeProps) => JSX.Element;')
+  })
+
+  test('supports fragments and member component tags', () => {
+    const output = processSource(`
+      export const Fields = () => <><Form.Field name="email" /><Form.Field name="name" /></>
+    `)
+
+    expect(output).toContain('export declare const Fields: () => JSX.Element;')
+  })
+
+  test('unions nullable component returns', () => {
+    const output = processSource(`
+      export function OptionalPanel(hidden: boolean) {
+        if (hidden) return null
+        return <section>Visible</section>
+      }
+    `)
+
+    expect(output).toContain('export declare function OptionalPanel(hidden: boolean): null | JSX.Element;')
+  })
+
+  test('preserves explicit JSX return annotations', () => {
+    const output = processSource('export function Portal(): JSX.Element | null { return null }')
+    expect(output).toContain('export declare function Portal(): JSX.Element | null;')
+  })
+
+  test('does not mistake angle-bracket assertions for JSX', () => {
+    const output = processSource('export const value = <string>input')
+    expect(output).not.toContain('value: JSX.Element')
+  })
+})

packages/zig-dtsx/src/type_inference.zig

@@ -217,6 +217,34 @@ fn isRegexLiteral(value: []const u8) bool {
     return false;
 }
 
+fn isJsxTagNameChar(c: u8) bool {
+    return ch.isIdentChar(c) or c == '.' or c == ':' or c == '-';
+}
+
+/// Detect a complete JSX element or fragment without assuming a particular JSX
+/// runtime. This deliberately requires a matching outer close or a self-close,
+/// keeping angle-bracket assertions and generic arrow functions distinct.
+fn isJsxExpression(value: []const u8) bool {
+    const expression = trim(value);
+    if (expression.len < 3 or expression[0] != '<') return false;
+    if (ch.startsWith(expression, "<>")) return ch.endsWith(expression, "</>");
+    if (!ch.isIdentStart(expression[1])) return false;
+
+    var tag_end: usize = 2;
+    while (tag_end < expression.len and isJsxTagNameChar(expression[tag_end])) tag_end += 1;
+    const tag_name = expression[1..tag_end];
+    if (tag_name.len == 0 or tag_end >= expression.len) return false;
+    const delimiter = expression[tag_end];
+    if (!ch.isWhitespace(delimiter) and delimiter != '>' and delimiter != '/') return false;
+    if (ch.endsWith(expression, "/>")) return true;
+
+    const close_start = std.mem.lastIndexOf(u8, expression, "</") orelse return false;
+    var close_end = close_start + 2;
+    while (close_end < expression.len and isJsxTagNameChar(expression[close_end])) close_end += 1;
+    return std.mem.eql(u8, expression[close_start + 2 .. close_end], tag_name) and
+        close_end < expression.len and expression[close_end] == '>' and close_end == expression.len - 1;
+}
+
 fn inferAccessType(alloc: std.mem.Allocator, value: []const u8) InferError!?[]const u8 {
     if (value.len > 3 and value[value.len - 1] == ']') {
         const bracket = std.mem.lastIndexOfScalar(u8, value, '[') orelse return null;
@@ -990,6 +1018,11 @@ pub fn inferFunctionBodyReturnType(alloc: std.mem.Allocator, body: []const u8, p
                 if (expression_depth == 0) break;
                 expression_depth -= 1;
             } else if (c == ';' and expression_depth == 0) break;
+            if ((c == '\n' or c == '\r') and expression_depth == 0) {
+                var next = i + 1;
+                while (next < content.len and ch.isWhitespace(content[next])) next += 1;
+                if (ch.startsWith(content[next..], "return") and (next + 6 >= content.len or !ch.isIdentChar(content[next + 6]))) break;
+            }
         }
         const expression = trim(content[expression_start..i]);
         const inferred = try inferBodyExpressionType(alloc, expression, parameters, depth + 1);
@@ -1012,6 +1045,7 @@ fn inferBodyExpressionType(alloc: std.mem.Allocator, expression: []const u8, par
     while (value.len >= 2 and value[0] == '(' and value[value.len - 1] == ')' and findMatchingBracket(value, 0, '(', ')') == value.len - 1) {
         value = trim(value[1 .. value.len - 1]);
     }
+    if (isJsxExpression(value)) return "JSX.Element";
     if (findParameterType(parameters, value)) |parameter_type| return parameter_type;
     if (try inferBodyCallType(alloc, value, parameters)) |call_type| return call_type;
 
@@ -1114,6 +1148,8 @@ pub fn inferNarrowType(alloc: std.mem.Allocator, value: []const u8, is_const: bo
         return inferNarrowType(alloc, trim(trimmed[1 .. trimmed.len - 1]), is_const, in_union, depth + 1);
     }
 
+    if (isJsxExpression(trimmed)) return "JSX.Element";
+
     // The comma operator evaluates to its final operand. Reuse the balanced
     // element splitter so commas inside calls, arrays, and objects are ignored.
     if (findTopLevelComma(trimmed)) |comma| return inferNarrowType(alloc, trim(trimmed[comma + 1 ..]), is_const, in_union, depth + 1);
@@ -2562,6 +2598,18 @@ test "arithmetic and update expressions emit valid result types" {
     try std.testing.expectEqualStrings("typeof state.count", try inferNarrowType(alloc, "--state.count", false, false, 0));
 }
 
+test "JSX expressions infer portable element types" {
+    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
+    defer arena.deinit();
+    const alloc = arena.allocator();
+    try std.testing.expectEqualStrings("JSX.Element", try inferNarrowType(alloc, "<button>Save</button>", false, false, 0));
+    try std.testing.expectEqualStrings("JSX.Element", try inferNarrowType(alloc, "<Form.Field name='email' />", false, false, 0));
+    try std.testing.expectEqualStrings("JSX.Element", try inferNarrowType(alloc, "<><span>One</span><span>Two</span></>", false, false, 0));
+    try std.testing.expectEqualStrings("JSX.Element", try inferFunctionBodyReturnType(alloc, "{ return <Card title={title} /> }", "(title: string)", 0));
+    try std.testing.expectEqualStrings("null | JSX.Element", try inferFunctionBodyReturnType(alloc, "{ if (hidden) return null; return <Panel /> }", "(hidden: boolean)", 0));
+    try std.testing.expectEqualStrings("boolean", try inferNarrowType(alloc, "<Value>input", false, false, 0));
+}
+
 test "extractSatisfiesType" {
     try std.testing.expectEqualStrings("Config", extractSatisfiesType("{ port: 3000 } satisfies Config").?);
     try std.testing.expect(extractSatisfiesType("just a value without it") == null);