Add attributes when reconstituting, not after

src/GraphRewriting.ts

@@ -5,13 +5,13 @@ import { Map, OrderedMap, OrderedSet, Set, Collection, isCollection } from "immu
 import { PrimitiveTypeKind, Type, ClassProperty } from "./Type";
 import { combineTypeAttributesOfTypes } from "./TypeUtils";
 import { TypeGraph } from "./TypeGraph";
-import { TypeAttributes } from "./TypeAttributes";
+import { TypeAttributes, emptyTypeAttributes, combineTypeAttributes } from "./TypeAttributes";
 import { assert, panic } from "./Support";
 import { TypeRef, TypeBuilder, StringTypeMapping } from "./TypeBuilder";
 
 export interface TypeLookerUp {
     lookupTypeRefs(typeRefs: TypeRef[], forwardingRef?: TypeRef): TypeRef | undefined;
-    reconstituteTypeRef(typeRef: TypeRef, forwardingRef?: TypeRef): TypeRef;
+    reconstituteTypeRef(typeRef: TypeRef, attributes?: TypeAttributes, forwardingRef?: TypeRef): TypeRef;
 }
 
 export class TypeReconstituter<TBuilder extends BaseGraphRewriteBuilder> {
@@ -198,18 +198,25 @@ export abstract class BaseGraphRewriteBuilder extends TypeBuilder implements Typ
     }
 
     reconstituteType(t: Type, forwardingRef?: TypeRef): TypeRef {
-        return this.reconstituteTypeRef(t.typeRef, forwardingRef);
+        return this.reconstituteTypeRef(t.typeRef, undefined, forwardingRef);
     }
 
     abstract lookupTypeRefs(typeRefs: TypeRef[], forwardingRef?: TypeRef, replaceSet?: boolean): TypeRef | undefined;
-    protected abstract forceReconstituteTypeRef(originalRef: TypeRef, maybeForwardingRef?: TypeRef): TypeRef;
+    protected abstract forceReconstituteTypeRef(
+        originalRef: TypeRef,
+        attributes?: TypeAttributes,
+        maybeForwardingRef?: TypeRef
+    ): TypeRef;
 
-    reconstituteTypeRef(originalRef: TypeRef, maybeForwardingRef?: TypeRef): TypeRef {
+    reconstituteTypeRef(originalRef: TypeRef, attributes?: TypeAttributes, maybeForwardingRef?: TypeRef): TypeRef {
         const maybeRef = this.lookupTypeRefs([originalRef], maybeForwardingRef);
         if (maybeRef !== undefined) {
+            if (attributes !== undefined) {
+                this.addAttributes(maybeRef, attributes);
+            }
             return maybeRef;
         }
-        return this.forceReconstituteTypeRef(originalRef, maybeForwardingRef);
+        return this.forceReconstituteTypeRef(originalRef, attributes, maybeForwardingRef);
     }
 
     protected assertTypeRefsToReconstitute(typeRefs: TypeRef[], forwardingRef?: TypeRef): void {
@@ -302,18 +309,24 @@ export class GraphRemapBuilder extends BaseGraphRewriteBuilder {
         return first;
     }
 
-    protected forceReconstituteTypeRef(originalRef: TypeRef, maybeForwardingRef?: TypeRef): TypeRef {
+    protected forceReconstituteTypeRef(
+        originalRef: TypeRef,
+        attributes?: TypeAttributes,
+        maybeForwardingRef?: TypeRef
+    ): TypeRef {
         assert(maybeForwardingRef === undefined, "We can't have a forwarding ref when we remap");
 
         originalRef = this.getMapTarget(originalRef);
         const [originalType, originalAttributes] = originalRef.deref();
 
         const attributeSources = this._attributeSources.get(originalType);
-        let attributes: TypeAttributes;
+        if (attributes === undefined) {
+            attributes = emptyTypeAttributes;
+        }
         if (attributeSources === undefined) {
-            attributes = originalAttributes;
+            attributes = combineTypeAttributes(attributes, originalAttributes);
         } else {
-            attributes = combineTypeAttributesOfTypes(attributeSources);
+            attributes = combineTypeAttributes(attributes, combineTypeAttributesOfTypes(attributeSources));
         }
 
         const index = originalRef.index;
@@ -429,7 +442,11 @@ export class GraphRewriteBuilder<T extends Type> extends BaseGraphRewriteBuilder
         });
     }
 
-    protected forceReconstituteTypeRef(originalRef: TypeRef, maybeForwardingRef?: TypeRef): TypeRef {
+    protected forceReconstituteTypeRef(
+        originalRef: TypeRef,
+        attributes?: TypeAttributes,
+        maybeForwardingRef?: TypeRef
+    ): TypeRef {
         const [originalType, originalAttributes] = originalRef.deref();
         const index = originalRef.index;
 
@@ -438,10 +455,16 @@ export class GraphRewriteBuilder<T extends Type> extends BaseGraphRewriteBuilder
             this.changeDebugPrintIndent(1);
         }
 
+        if (attributes === undefined) {
+            attributes = originalAttributes;
+        } else {
+            attributes = combineTypeAttributes(attributes, originalAttributes);
+        }
+
         const reconstituter = new TypeReconstituter(
             this,
             this.alphabetizeProperties,
-            originalAttributes,
+            attributes,
             maybeForwardingRef,
             tref => {
                 if (this.debugPrint) {

src/UnifyClasses.ts

@@ -217,9 +217,7 @@ export function unifyTypes<T extends Type>(
     } else if (types.count() === 1) {
         const first = defined(types.first());
         if (!(first instanceof UnionType)) {
-            const tref = typeBuilder.reconstituteTypeRef(first.typeRef, maybeForwardingRef);
-            typeBuilder.addAttributes(tref, typeAttributes);
-            return tref;
+            return typeBuilder.reconstituteTypeRef(first.typeRef, typeAttributes, maybeForwardingRef);
         }
     }