From e1ce20358f4dc63c178d24bb5e52405e840d9725 Mon Sep 17 00:00:00 2001
From: "Marc J. Schmidt" <marc@marcjschmidt.de>
Date: Wed, 8 May 2024 22:54:59 +0200
Subject: [PATCH] fix(type-compiler): support ReceiveType in arrow function
 with body expression

This makes sure the following works correctly:

```typescript
(type?: ReceiveType) => (v: any) => {
    type = resolveReceiveType(type);
    // ...
};
```
---
 packages/type-compiler/src/compiler.ts         | 16 +++++++++++++---
 packages/type-compiler/tests/transpile.spec.ts | 15 +++++++++++++++
 packages/type/tests/receive-type.spec.ts       | 13 ++++++++++++-
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/packages/type-compiler/src/compiler.ts b/packages/type-compiler/src/compiler.ts
index d43d8a898..30c68bcb4 100644
--- a/packages/type-compiler/src/compiler.ts
+++ b/packages/type-compiler/src/compiler.ts
@@ -8,7 +8,7 @@
  * You should have received a copy of the MIT License along with this program.
  */
 
-import type {
+import ts, {
     __String,
     ArrayTypeNode,
     ArrowFunction,
@@ -43,6 +43,8 @@ import type {
     InferTypeNode,
     InterfaceDeclaration,
     IntersectionTypeNode,
+    isBlock,
+    isExpression,
     LiteralTypeNode,
     MappedTypeNode,
     MethodDeclaration,
@@ -72,7 +74,6 @@ import type {
     TypeReferenceNode,
     UnionTypeNode,
 } from 'typescript';
-import ts from 'typescript';
 
 import {
     ensureImportIsEmitted,
@@ -1105,7 +1106,16 @@ export class ReflectionTransformer implements CustomTransformer {
             this.f.createToken(ts.SyntaxKind.EqualsToken),
             this.f.createIdentifier('undefined'),
         ));
-        const body = node.body ? this.f.updateBlock(node.body as Block, [reset, ...(node.body as Block).statements]) : undefined;
+
+        // convert expression into statements array
+        let body = node.body && isBlock(node.body) ? node.body : undefined;
+        let bodyStatements: Statement[] = node.body && isBlock(node.body) ? [...node.body.statements] : [];
+        if (node.body) {
+            if (isExpression(node.body)) {
+                bodyStatements = [this.f.createReturnStatement(node.body)];
+            }
+            body = this.f.updateBlock(node.body as Block, [reset, ...bodyStatements]);
+        }
 
         if (isArrowFunction(node)) {
             return this.f.updateArrowFunction(node, node.modifiers, node.typeParameters, node.parameters, node.type, node.equalsGreaterThanToken, body as ConciseBody) as T;
diff --git a/packages/type-compiler/tests/transpile.spec.ts b/packages/type-compiler/tests/transpile.spec.ts
index ed44d10e6..c4c487af2 100644
--- a/packages/type-compiler/tests/transpile.spec.ts
+++ b/packages/type-compiler/tests/transpile.spec.ts
@@ -536,3 +536,18 @@ test('resolve type ref2', () => {
     console.log(res.app);
     expect(res.app).toContain(`() => Guest, 'Guest'`);
 });
+
+test('ReceiveType arrow function', () => {
+    const res = transpile({
+        'app': `
+export const typeValidation = <T>(type?: ReceiveType<T>): ValidatorFn => (control: AbstractControl) => {
+        type = resolveReceiveType(type);
+        const errors = validate<T>(control.value)
+        console.log(errors)
+        return errors.length ? {validation: errors[0]} : null
+}
+        `
+    });
+    console.log(res.app);
+    expect(res.app).toContain(`exports.typeValidation.Ω = undefined; return __assignType((control) =>`);
+});
diff --git a/packages/type/tests/receive-type.spec.ts b/packages/type/tests/receive-type.spec.ts
index e9634aa61..2d9348356 100644
--- a/packages/type/tests/receive-type.spec.ts
+++ b/packages/type/tests/receive-type.spec.ts
@@ -1,7 +1,7 @@
 import { expect, test } from '@jest/globals';
 import { ReceiveType, resolveReceiveType, typeOf } from '../src/reflection/reflection.js';
 import { ReflectionKind, Type } from '../src/reflection/type.js';
-import { ReflectionOp } from '@deepkit/type-spec';
+import { validates } from '../src/validator.js';
 
 test('typeOf', () => {
     const type = typeOf<string>();
@@ -97,3 +97,14 @@ test('class constructor multiple', () => {
     expect(aString.type1).toMatchObject({ kind: ReflectionKind.string });
     expect(aString.type2).toMatchObject({ kind: ReflectionKind.number });
 });
+
+test('function with ReceiveType return expression', () => {
+    const typeValidation = <T>(type?: ReceiveType<T>) => (value: any) => {
+        type = resolveReceiveType(type);
+        return validates(value, type);
+    }
+
+    const validateString = typeValidation<string>();
+    expect(validateString('hello')).toBe(true);
+    expect(validateString(2)).toBe(false);
+});