From d9124166a1203e43f1d9093db82365cbaf9b70bc Mon Sep 17 00:00:00 2001
From: MECHERI Akram <mecheri.akram@gmail.com>
Date: Fri, 27 Jan 2023 13:54:40 +0100
Subject: [PATCH] fix(cdk/collections): SelectionModel does not always respect
 the compareWith function (#26447)

Fixed bug in SelectionModel where compareWith function was not consistently respected in deselect method.

Fixes #25878

(cherry picked from commit 50ec808eaef5b554df382591a6a58f88b429237f)
---
 src/cdk/collections/selection-model.ts | 26 +++++++++++++++++---------
 src/cdk/collections/selection.spec.ts  | 20 ++++++++++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/src/cdk/collections/selection-model.ts b/src/cdk/collections/selection-model.ts
index c5ac4daff494..7a1a36b02b7b 100644
--- a/src/cdk/collections/selection-model.ts
+++ b/src/cdk/collections/selection-model.ts
@@ -131,15 +131,7 @@ export class SelectionModel<T> {
    * Determines whether a value is selected.
    */
   isSelected(value: T): boolean {
-    if (this.compareWith) {
-      for (const otherValue of this._selection) {
-        if (this.compareWith(otherValue, value)) {
-          return true;
-        }
-      }
-      return false;
-    }
-    return this._selection.has(value);
+    return this._selection.has(this._getConcreteValue(value));
   }
 
   /**
@@ -191,6 +183,7 @@ export class SelectionModel<T> {
 
   /** Selects a value. */
   private _markSelected(value: T) {
+    value = this._getConcreteValue(value);
     if (!this.isSelected(value)) {
       if (!this._multiple) {
         this._unmarkAll();
@@ -208,6 +201,7 @@ export class SelectionModel<T> {
 
   /** Deselects a value. */
   private _unmarkSelected(value: T) {
+    value = this._getConcreteValue(value);
     if (this.isSelected(value)) {
       this._selection.delete(value);
 
@@ -238,6 +232,20 @@ export class SelectionModel<T> {
   private _hasQueuedChanges() {
     return !!(this._deselectedToEmit.length || this._selectedToEmit.length);
   }
+
+  /** Returns a value that is comparable to inputValue by applying compareWith function, returns the same inputValue otherwise. */
+  private _getConcreteValue(inputValue: T): T {
+    if (!this.compareWith) {
+      return inputValue;
+    } else {
+      for (let selectedValue of this._selection) {
+        if (this.compareWith!(inputValue, selectedValue)) {
+          return selectedValue;
+        }
+      }
+      return inputValue;
+    }
+  }
 }
 
 /**
diff --git a/src/cdk/collections/selection.spec.ts b/src/cdk/collections/selection.spec.ts
index ac430f962f23..41d0c7f7784f 100644
--- a/src/cdk/collections/selection.spec.ts
+++ b/src/cdk/collections/selection.spec.ts
@@ -279,4 +279,24 @@ describe('SelectionModel', () => {
     let singleSelectionModel = new SelectionModel();
     expect(singleSelectionModel.isMultipleSelection()).toBe(false);
   });
+
+  it('should deselect value if comparable to another one', () => {
+    type Item = {key: number; value: string};
+    const v1: Item = {key: 1, value: 'blue'};
+    const v2: Item = {key: 1, value: 'green'};
+    const compareFun = (x: Item, y: Item) => x.key === y.key;
+    const model = new SelectionModel<Item>(false, [v1], false, compareFun);
+    model.deselect(v2);
+    expect(model.selected.length).toBe(0);
+  });
+
+  it('should not deselect value if not comparable to another one', () => {
+    type Item = {key: number; value: string};
+    const v1: Item = {key: 1, value: 'blue'};
+    const v2: Item = {key: 2, value: 'apple'};
+    const compareFun = (x: Item, y: Item) => x.key === y.key;
+    const model = new SelectionModel<Item>(false, [v1], false, compareFun);
+    model.deselect(v2);
+    expect(model.selected.length).toBe(1);
+  });
 });