Working TraitExpressionBuilder

src/TraitExpressionBuilder.ts

@@ -1,6 +1,6 @@
 import { Call, Fn, Pipe, Tuples } from "hotscript"
 import { AbstractClass } from "type-fest"
-import { Trait, TraitExpression, TraitOwnAbstractFn, TraitOwnImplClassFn, TraitOwnImplInstanceFn, TraitSupertraitsFn, emptyTraitExpression } from "."
+import { Trait, TraitExpression, TraitOwnAbstractFn, TraitOwnImplClassFn, TraitOwnImplInstanceFn, emptyTraitExpression } from "."
 import { ExtendableFn, StaticMembersFn } from "./util"
 
 
@@ -27,17 +27,18 @@ class TraitExpressionBuilder<
         Traits extends Trait<any, any, any>[]
     >(
         ...traits: Traits
-    ) {
+    ): TraitExpressionBuilder<
+        Super,
+        [...OwnTraits, ...Traits],
+        [...AllTraits, SpreadSupertraits<Traits>]
+    > {
         return new TraitExpressionBuilder(
             new TraitExpression(
                 this.expression.superclass,
                 [...this.expression.ownTraits, ...traits] as const,
-                [
-                    ...this.expression.allTraits,
-                    ...this.spreadSupertraits(traits) as SpreadTraitsHierarchy<Traits>,
-                ] as const,
+                [...this.expression.allTraits, ...this.spreadSupertraits(traits)] as const,
             )
-        )
+        ) as any
     }
 
     private spreadSupertraits<
@@ -50,9 +51,9 @@ class TraitExpressionBuilder<
         traits: Traits
     ) {
         return traits.flatMap(trait => [
-            ...trait.superExpression.allTraits,
+            ...trait.supertraits.allTraits,
             trait,
-        ])
+        ]) as SpreadSupertraits<Traits>
     }
 
     get() {
@@ -65,18 +66,17 @@ class TraitExpressionBuilder<
 }
 
 
-type SpreadTraitsHierarchy<Traits extends Trait<any, any, any>[]> = (
+type SpreadSupertraits<Traits extends Trait<any, any, any>[]> = (
     Call<
         Tuples.FlatMap<PrependTraitSupertraitsFn>,
         Traits
     >
 )
 interface PrependTraitSupertraitsFn extends Fn {
-    return: this["arg0"] extends Trait<any, any, any>
-        ? [
-            ...Call<TraitSupertraitsFn, this["arg0"]>,
-            this["arg0"],
-        ]
+    return: this["arg0"] extends Trait<infer Super, any, any>
+        ? Super extends TraitExpression<any, any, infer AllTraits>
+            ? [...AllTraits, this["arg0"]]
+            : never
         : never
 }