From 26d132faa9ec7cfea5e7dd0bf2585bd90bb56d59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Ad=C3=A1mek?= <banan23@gmail.com>
Date: Tue, 26 May 2020 22:39:30 +0200
Subject: [PATCH] feat(core): accept references in collection `add/remove/set`
 methods

---
 packages/core/src/entity/ArrayCollection.ts | 32 ++++++++++++++-------
 packages/core/src/entity/Collection.ts      | 30 ++++++++++---------
 2 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/packages/core/src/entity/ArrayCollection.ts b/packages/core/src/entity/ArrayCollection.ts
index 188b9328a837..78367dd7cd20 100644
--- a/packages/core/src/entity/ArrayCollection.ts
+++ b/packages/core/src/entity/ArrayCollection.ts
@@ -1,7 +1,9 @@
 import { AnyEntity, Dictionary, EntityProperty, IPrimaryKey, Primary } from '../typings';
 import { ReferenceType } from './enums';
 import { Collection } from './Collection';
+import { Reference } from './Reference';
 import { wrap } from './wrap';
+import { Utils } from '../utils';
 
 export class ArrayCollection<T extends AnyEntity<T>, O extends AnyEntity<O>> {
 
@@ -46,31 +48,37 @@ export class ArrayCollection<T extends AnyEntity<T>, O extends AnyEntity<O>> {
     return this.getItems().map(i => i[field as keyof T]) as unknown as U[];
   }
 
-  add(...items: T[]): void {
+  add(...items: (T | Reference<T>)[]): void {
     for (const item of items) {
-      if (!this.contains(item, false)) {
-        this.items.push(item);
-        this.propagate(item, 'add');
+      const entity = Utils.unwrapReference(item);
+
+      if (!this.contains(entity, false)) {
+        this.items.push(entity);
+        this.propagate(entity, 'add');
       }
     }
 
     Object.assign(this, this.items);
   }
 
-  set(items: T[]): void {
+  set(items: (T | Reference<T>)[]): void {
     this.removeAll();
     this.add(...items);
   }
 
+  /**
+   * @internal
+   */
   hydrate(items: T[]): void {
     this.items.length = 0;
     this.items.push(...items);
     Object.assign(this, this.items);
   }
 
-  remove(...items: T[]): void {
+  remove(...items: (T | Reference<T>)[]): void {
     for (const item of items) {
-      const idx = this.items.findIndex(i => wrap(i, true).__serializedPrimaryKey === wrap(item, true).__serializedPrimaryKey);
+      const entity = Utils.unwrapReference(item);
+      const idx = this.items.findIndex(i => wrap(i, true).__serializedPrimaryKey === wrap(entity, true).__serializedPrimaryKey);
 
       if (idx !== -1) {
         delete this[this.items.length - 1]; // remove last item
@@ -78,7 +86,7 @@ export class ArrayCollection<T extends AnyEntity<T>, O extends AnyEntity<O>> {
         Object.assign(this, this.items); // reassign array access
       }
 
-      this.propagate(item, 'remove');
+      this.propagate(entity, 'remove');
     }
   }
 
@@ -86,10 +94,12 @@ export class ArrayCollection<T extends AnyEntity<T>, O extends AnyEntity<O>> {
     this.remove(...this.items);
   }
 
-  contains(item: T, check?: boolean): boolean {
+  contains(item: T | Reference<T>, check?: boolean): boolean {
+    const entity = Utils.unwrapReference(item);
+
     return !!this.items.find(i => {
-      const objectIdentity = i === item;
-      const primaryKeyIdentity = !!wrap(i, true).__primaryKey && !!wrap(item, true).__primaryKey && wrap(i, true).__serializedPrimaryKey === wrap(item, true).__serializedPrimaryKey;
+      const objectIdentity = i === entity;
+      const primaryKeyIdentity = !!wrap(i, true).__primaryKey && !!wrap(entity, true).__primaryKey && wrap(i, true).__serializedPrimaryKey === wrap(entity, true).__serializedPrimaryKey;
 
       return objectIdentity || primaryKeyIdentity;
     });
diff --git a/packages/core/src/entity/Collection.ts b/packages/core/src/entity/Collection.ts
index 72111b826e92..de74f9df1693 100644
--- a/packages/core/src/entity/Collection.ts
+++ b/packages/core/src/entity/Collection.ts
@@ -3,6 +3,7 @@ import { ArrayCollection } from './index';
 import { ReferenceType } from './enums';
 import { Utils, ValidationError } from '../utils';
 import { QueryOrder, QueryOrderMap } from '../enums';
+import { Reference } from './Reference';
 import { wrap } from './wrap';
 
 export class Collection<T extends AnyEntity<T>, O extends AnyEntity<O> = AnyEntity> extends ArrayCollection<T, O> {
@@ -41,18 +42,20 @@ export class Collection<T extends AnyEntity<T>, O extends AnyEntity<O> = AnyEnti
     return super.getItems();
   }
 
-  add(...items: T[]): void {
-    items.map(item => this.validateItemType(item));
-    this.modify('add', items);
-    this.cancelOrphanRemoval(items);
+  add(...items: (T | Reference<T>)[]): void {
+    const unwrapped = items.map(i => Utils.unwrapReference(i));
+    unwrapped.map(item => this.validateItemType(item));
+    this.modify('add', unwrapped);
+    this.cancelOrphanRemoval(unwrapped);
   }
 
-  set(items: T[]): void {
-    items.map(item => this.validateItemType(item));
-    this.validateModification(items);
-    super.set(items);
+  set(items: (T | Reference<T>)[]): void {
+    const unwrapped = items.map(i => Utils.unwrapReference(i));
+    unwrapped.map(item => this.validateItemType(item));
+    this.validateModification(unwrapped);
+    super.set(unwrapped);
     this.setDirty();
-    this.cancelOrphanRemoval(items);
+    this.cancelOrphanRemoval(unwrapped);
   }
 
   /**
@@ -71,18 +74,19 @@ export class Collection<T extends AnyEntity<T>, O extends AnyEntity<O> = AnyEnti
     }
   }
 
-  remove(...items: T[]): void {
-    this.modify('remove', items);
+  remove(...items: (T | Reference<T>)[]): void {
+    const unwrapped = items.map(i => Utils.unwrapReference(i));
+    this.modify('remove', unwrapped);
     const em = wrap(this.owner, true).__em;
 
     if (this.property.orphanRemoval && em) {
-      for (const item of items) {
+      for (const item of unwrapped) {
         em.getUnitOfWork().scheduleOrphanRemoval(item);
       }
     }
   }
 
-  contains(item: T, check = true): boolean {
+  contains(item: (T | Reference<T>), check = true): boolean {
     if (check) {
       this.checkInitialized();
     }