From eeacdf1c9d29e29ed985003379527b5cc9253ef2 Mon Sep 17 00:00:00 2001
From: stopachka <stepan.p@gmail.com>
Date: Tue, 31 Dec 2024 11:53:26 -0500
Subject: [PATCH] op 1

---
 client/packages/core/src/schemaTypes.ts       | 44 ++++++++++++-------
 .../strong-init-vite/instant.schema.v2.ts     |  1 +
 .../src/typescript_tests_experimental_v2.tsx  |  8 ++++
 3 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/client/packages/core/src/schemaTypes.ts b/client/packages/core/src/schemaTypes.ts
index b0351584b..ae002ecb4 100644
--- a/client/packages/core/src/schemaTypes.ts
+++ b/client/packages/core/src/schemaTypes.ts
@@ -45,13 +45,6 @@ export class DataAttrDef<ValueType, IsRequired extends boolean> {
   // }
 }
 
-type ExtractValueType<T> =
-  T extends DataAttrDef<infer ValueType, infer isRequired>
-    ? isRequired extends true
-      ? ValueType
-      : ValueType | undefined
-    : never;
-
 export class LinkAttrDef<
   Cardinality extends CardinalityKind,
   EntityName extends string,
@@ -90,9 +83,7 @@ export class EntityDef<
   ) {}
 
   asType<
-    _AsType extends Partial<{
-      [AttrName in keyof Attrs]: ExtractValueType<Attrs[AttrName]>;
-    }>,
+    _AsType extends Partial<MappedAttrs<Attrs>>,
   >() {
     return new EntityDef<Attrs, Links, _AsType>(this.attrs, this.links);
   }
@@ -231,13 +222,36 @@ type LinksIndexedByEntity<
   };
 };
 
+type RequiredKeys<Attrs extends AttrsDefs> = {
+  [K in keyof Attrs]: Attrs[K] extends DataAttrDef<any, infer R>
+    ? R extends true
+      ? K
+      : never
+    : never;
+}[keyof Attrs];
+
+type OptionalKeys<Attrs extends AttrsDefs> = {
+  [K in keyof Attrs]: Attrs[K] extends DataAttrDef<any, infer R>
+    ? R extends false
+      ? K
+      : never
+    : never;
+}[keyof Attrs];
+
+/**
+ * MappedAttrs:
+ *   - Required keys => `key: ValueType`
+ *   - Optional keys => `key?: ValueType`
+ */
+type MappedAttrs<Attrs extends AttrsDefs> = {
+  [K in RequiredKeys<Attrs>]: Attrs[K] extends DataAttrDef<infer V, any> ? V : never;
+} & {
+  [K in OptionalKeys<Attrs>]?: Attrs[K] extends DataAttrDef<infer V, any> ? V : never;
+};
+
 export type ResolveEntityAttrs<
   EDef extends EntityDef<any, any, any>,
-  ResolvedAttrs = {
-    [AttrName in keyof EDef["attrs"]]: ExtractValueType<
-      EDef["attrs"][AttrName]
-    >;
-  },
+  ResolvedAttrs = MappedAttrs<EDef["attrs"]>
 > =
   EDef extends EntityDef<any, any, infer AsType>
     ? AsType extends void
diff --git a/client/sandbox/strong-init-vite/instant.schema.v2.ts b/client/sandbox/strong-init-vite/instant.schema.v2.ts
index f584f3db0..3e710004c 100644
--- a/client/sandbox/strong-init-vite/instant.schema.v2.ts
+++ b/client/sandbox/strong-init-vite/instant.schema.v2.ts
@@ -4,6 +4,7 @@ const _schema = i.schema({
   entities: {
     messages: i.entity({
       content: i.string(),
+      createdAt: i.date().optional(),
     }),
     $users: i.entity({
       email: i.string().unique().indexed(),
diff --git a/client/sandbox/strong-init-vite/src/typescript_tests_experimental_v2.tsx b/client/sandbox/strong-init-vite/src/typescript_tests_experimental_v2.tsx
index 5dda1685b..8043550ce 100644
--- a/client/sandbox/strong-init-vite/src/typescript_tests_experimental_v2.tsx
+++ b/client/sandbox/strong-init-vite/src/typescript_tests_experimental_v2.tsx
@@ -151,6 +151,14 @@ const messagesQuery = {
 } satisfies InstaQLParams<AppSchema>;
 
 type CoreMessage = InstaQLEntity<AppSchema, "messages">;
+const mWithOptionalFieldWorks: CoreMessage = { 
+  id: '1', 
+  content: 'hello',
+};
+
+// to silence ts warnings
+mWithOptionalFieldWorks;
+
 let coreMessage: CoreMessage = 1 as any;
 coreMessage.content;