MessageController: return all messages by sent at, not just 1

ts/models/messages.ts

@@ -10,6 +10,7 @@ import {
   QuotedMessageType,
   WhatIsThis,
 } from '../model-types.d';
+import { find } from '../util/iterables';
 import { DataMessageClass } from '../textsecure.d';
 import { ConversationModel } from './conversations';
 import { ConversationType } from '../state/ducks/conversations';
@@ -970,10 +971,13 @@ export class MessageModel extends window.Backbone.Model<MessageAttributesType> {
       window.log.info(
         `doubleCheckMissingQuoteReference/${logId}: Verifying reference to ${sentAt}`
       );
-      const inMemoryMessage = window.MessageController.findBySentAt(
+      const inMemoryMessages = window.MessageController.filterBySentAt(
         Number(sentAt)
       );
-      if (!isQuoteAMatch(inMemoryMessage, this.get('conversationId'), quote)) {
+      const matchingMessage = find(inMemoryMessages, message =>
+        isQuoteAMatch(message, this.get('conversationId'), quote)
+      );
+      if (!matchingMessage) {
         window.log.info(
           `doubleCheckMissingQuoteReference/${logId}: No match for ${sentAt}.`
         );
@@ -992,7 +996,7 @@ export class MessageModel extends window.Backbone.Model<MessageAttributesType> {
         `doubleCheckMissingQuoteReference/${logId}: Found match for ${sentAt}, updating.`
       );
 
-      await this.copyQuoteContentFromOriginal(inMemoryMessage, quote);
+      await this.copyQuoteContentFromOriginal(matchingMessage, quote);
       this.set({
         quote: {
           ...quote,
@@ -2287,12 +2291,15 @@ export class MessageModel extends window.Backbone.Model<MessageAttributesType> {
     }
 
     const { id } = quote;
-    const inMemoryMessage = window.MessageController.findBySentAt(id);
+    const inMemoryMessages = window.MessageController.filterBySentAt(id);
+    const matchingMessage = find(inMemoryMessages, item =>
+      isQuoteAMatch(item, conversationId, quote)
+    );
 
-    let queryMessage;
+    let queryMessage: undefined | MessageModel;
 
-    if (isQuoteAMatch(inMemoryMessage, conversationId, quote)) {
-      queryMessage = inMemoryMessage;
+    if (matchingMessage) {
+      queryMessage = matchingMessage;
     } else {
       window.log.info('copyFromQuotedMessage: db lookup needed', id);
       const collection = await window.Signal.Data.getMessagesBySentAt(id, {

ts/test-both/util/iterables_test.ts

@@ -7,6 +7,7 @@ import * as sinon from 'sinon';
 import {
   concat,
   filter,
+  find,
   groupBy,
   isIterable,
   map,
@@ -211,6 +212,29 @@ describe('iterable utilities', () => {
     });
   });
 
+  describe('find', () => {
+    const isOdd = (n: number) => Boolean(n % 2);
+
+    it('returns undefined if the value is not found', () => {
+      assert.isUndefined(find([], isOdd));
+      assert.isUndefined(find([2, 4], isOdd));
+    });
+
+    it('returns the first matching value', () => {
+      assert.strictEqual(find([0, 1, 2, 3], isOdd), 1);
+    });
+
+    it('only iterates until a value is found', () => {
+      function* numbers() {
+        yield 2;
+        yield 3;
+        throw new Error('this should never happen');
+      }
+
+      find(numbers(), isOdd);
+    });
+  });
+
   describe('groupBy', () => {
     it('returns an empty object if passed an empty iterable', () => {
       const fn = sinon.fake();

ts/test-electron/util/MessageController_test.ts

@@ -0,0 +1,56 @@
+// Copyright 2021 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import { assert } from 'chai';
+import { MessageModel } from '../../models/messages';
+
+import { MessageController } from '../../util/MessageController';
+
+describe('MessageController', () => {
+  describe('filterBySentAt', () => {
+    it('returns an empty iterable if no messages match', () => {
+      const mc = new MessageController();
+
+      assert.isEmpty([...mc.filterBySentAt(123)]);
+    });
+
+    it('returns all messages that match the timestamp', () => {
+      const mc = new MessageController();
+      const message1 = new MessageModel({
+        conversationId: 'xyz',
+        id: 'abc',
+        received_at: 1,
+        sent_at: 1234,
+        timestamp: 9999,
+        type: 'incoming',
+      });
+      const message2 = new MessageModel({
+        conversationId: 'xyz',
+        id: 'def',
+        received_at: 2,
+        sent_at: 1234,
+        timestamp: 9999,
+        type: 'outgoing',
+      });
+      const message3 = new MessageModel({
+        conversationId: 'xyz',
+        id: 'ignored',
+        received_at: 3,
+        sent_at: 5678,
+        timestamp: 9999,
+        type: 'outgoing',
+      });
+      mc.register(message1.id, message1);
+      mc.register(message2.id, message2);
+      // We deliberately register this message twice for testing.
+      mc.register(message2.id, message2);
+      mc.register(message3.id, message3);
+
+      assert.sameMembers([...mc.filterBySentAt(1234)], [message1, message2]);
+
+      mc.unregister(message2.id);
+
+      assert.sameMembers([...mc.filterBySentAt(1234)], [message1]);
+    });
+  });
+});

ts/util/MessageController.ts

@@ -2,6 +2,8 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 
 import { MessageModel } from '../models/messages';
+import { map, filter } from './iterables';
+import { isNotNil } from './isNotNil';
 
 const SECOND = 1000;
 const MINUTE = SECOND * 60;
@@ -19,7 +21,7 @@ export class MessageController {
 
   private msgIDsBySender = new Map<string, string>();
 
-  private msgIDsBySentAt = new Map<number, string>();
+  private msgIDsBySentAt = new Map<number, Set<string>>();
 
   static install(): MessageController {
     const instance = new MessageController();
@@ -48,7 +50,14 @@ export class MessageController {
       timestamp: Date.now(),
     };
 
-    this.msgIDsBySentAt.set(message.get('sent_at'), id);
+    const sentAt = message.get('sent_at');
+    const previousIdsBySentAt = this.msgIDsBySentAt.get(sentAt);
+    if (previousIdsBySentAt) {
+      previousIdsBySentAt.add(id);
+    } else {
+      this.msgIDsBySentAt.set(sentAt, new Set([id]));
+    }
+
     this.msgIDsBySender.set(message.getSenderIdentifier(), id);
 
     return message;
@@ -58,7 +67,13 @@ export class MessageController {
     const { message } = this.messageLookup[id] || {};
     if (message) {
       this.msgIDsBySender.delete(message.getSenderIdentifier());
-      this.msgIDsBySentAt.delete(message.get('sent_at'));
+
+      const sentAt = message.get('sent_at');
+      const idsBySentAt = this.msgIDsBySentAt.get(sentAt) || new Set();
+      idsBySentAt.delete(id);
+      if (!idsBySentAt.size) {
+        this.msgIDsBySentAt.delete(sentAt);
+      }
     }
     delete this.messageLookup[id];
   }
@@ -87,12 +102,10 @@ export class MessageController {
     return existing && existing.message ? existing.message : undefined;
   }
 
-  findBySentAt(sentAt: number): MessageModel | undefined {
-    const id = this.msgIDsBySentAt.get(sentAt);
-    if (!id) {
-      return undefined;
-    }
-    return this.getById(id);
+  filterBySentAt(sentAt: number): Iterable<MessageModel> {
+    const ids = this.msgIDsBySentAt.get(sentAt) || [];
+    const maybeMessages = map(ids, id => this.getById(id));
+    return filter(maybeMessages, isNotNil);
   }
 
   findBySender(sender: string): MessageModel | undefined {

ts/util/iterables.ts

@@ -90,6 +90,18 @@ class FilterIterator<T> implements Iterator<T> {
   }
 }
 
+export function find<T>(
+  iterable: Iterable<T>,
+  predicate: (value: T) => unknown
+): undefined | T {
+  for (const value of iterable) {
+    if (predicate(value)) {
+      return value;
+    }
+  }
+  return undefined;
+}
+
 export function groupBy<T>(
   iterable: Iterable<T>,
   fn: (value: T) => string