A MultiMap for Salesforce Apex

MultiMap.cls

/**
 * @author aidan@processity.ai
 * @date 29/08/2024
 * @description A map where the items are lists of items and you can provide a function for how values map to keys.
 * Oh, if we had generics....
 */

public class MultiMap {
    private Map<Object, List<Object>> theMap;
    private Type listType;
    private KeyReader keyReader;

    public MultiMap(Type listType, MultiMap.KeyReader keyReader) {
        theMap = new Map<Object, List<Object>>();
        this.listType = listType;
        this.keyReader = keyReader;
    }

    public MultiMap(MultiMap.KeyReader keyReader) {
        this(List<Object>.class, keyReader);
    }

    public void putAll(List<Object> items) {
        Integer itemsSize = items.size();
        for (Integer i = 0; i < itemsSize; i++) {
            put(items[i]);
        }
    }

    public void put(Object item) {
        Object keyValue = keyReader.getKey(item);
        List<Object> thisList = theMap.get(keyValue);

        if (thisList == null) {
            thisList = (List<Object>) listType.newInstance();
            thisList.add(item);
            theMap.put(keyValue, thisList);
        } else {
            thisList.add(item);
        }
    }

    public List<Object> getWithKey(Object key) {
        return theMap.get(key);
    }

    public List<Object> getWithItem(Object item) {
        return theMap.get(keyReader.getKey(item));
    }

    public Set<Object> keySet() {
        return theMap.keySet();
    }

    public interface KeyReader {
        Object getKey(Object item);
    }
}

MultiMapTest.cls

/**
 * @author aidan@processity.ai
 * @date 29/08/2024
 */

@IsTest
private class MultiMapTest {
    
    @IsTest 
    static void contactsByLastName() {
        MultiMap multiMap = new MultiMap(List<Contact>.class, new FieldFromSObject(Contact.LastName));
        
        List<Contact> contacts = new List<Contact> {
                new Contact(FirstName = 'John', LastName = 'Smith'),
                new Contact(FirstName = 'Jane', LastName = 'Smith'),
                new Contact(FirstName = 'John', LastName = 'Not-Smith')
        };
        
        multiMap.putAll(contacts);

        Assert.areEqual(new List<Contact> { contacts[0], contacts[1] }, multiMap.getWithKey('Smith'));
        Assert.areEqual(new List<Contact> { contacts[2] }, multiMap.getWithKey('Not-Smith'));
    }
    
    @IsTest
    static void singleItem() {
        MultiMap multiMap = new MultiMap(new IsEvenKey());
        multiMap.put(5);

        Assert.areEqual(5, multiMap.getWithItem(5)[0]);
        Assert.areEqual(5, multiMap.getWithKey(false)[0]);
        Assert.areEqual(new Set<Object>{ false }, multiMap.keySet());
    }

    @IsTest
    static void twoItemsSameKey() {
        List<Integer> data = new List<Integer>{ 5, 7 };
        MultiMap multiMap = new MultiMap(List<Integer>.class, new IsEvenKey());

        multiMap.putAll(data);

        List<Integer> retrievedByItem = (List<Integer>) multiMap.getWithItem(5);
        List<Integer> retrievedByKey = (List<Integer>) multiMap.getWithKey(false);

        for (Integer i = 0; i < data.size(); i++) {
            Assert.areEqual(data[i], retrievedByItem[i]);
            Assert.areEqual(data[i], retrievedByKey[i]);
        }
        Assert.areEqual(new Set<Object>{ false }, multiMap.keySet());
    }

    @IsTest
    static void twoItemsDifferentKeys() {
        MultiMap multiMap = new MultiMap(new IsEvenKey());

        multiMap.put(5);
        multiMap.put(6);

        Assert.areEqual(5, multiMap.getWithItem(5)[0]);
        Assert.areEqual(5, multiMap.getWithKey(false)[0]);
        Assert.areEqual(6, multiMap.getWithItem(6)[0]);
        Assert.areEqual(6, multiMap.getWithKey(true)[0]);
        Assert.areEqual(new Set<Object>{ false, true }, multiMap.keySet());
    }

    private class IsEvenKey implements MultiMap.KeyReader {
        public Object getKey(Object item) {
            return ((Integer) item & 1) == 0;
        }
    }
    
    private class FieldFromSObject implements MultiMap.KeyReader {
        private SObjectField field;

        private FieldFromSObject(SObjectField field) {
            this.field = field;
        }

        public Object getKey(Object item) {
            return ((SObject)item).get(field);
        }
    }
}