diff --git a/src/main/java/io/usethesource/capsule/core/PersistentTrieSetMultimap.java b/src/main/java/io/usethesource/capsule/core/PersistentTrieSetMultimap.java index 30bbf72..ecea588 100644 --- a/src/main/java/io/usethesource/capsule/core/PersistentTrieSetMultimap.java +++ b/src/main/java/io/usethesource/capsule/core/PersistentTrieSetMultimap.java @@ -2223,6 +2223,8 @@ private static final class HashCollisionNode extends AbstractHashCollision final List>> collisionContent) { this.hash = hash; this.collisionContent = collisionContent; + + assert this.collisionContent.size() >= 2; } @Override @@ -2652,8 +2654,24 @@ public AbstractSetMultimapNode removed(AtomicReference mutator, K .filter(kImmutableSetEntry -> kImmutableSetEntry != optionalTuple.get()) .collect(Collectors.toList()); - details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE)); - return new HashCollisionNode(hash, updatedCollisionContent); + if (updatedCollisionContent.size() == 1) { + /* + * Create a root node that stores data for the single remaining key. This node will a) either become the + * new root returned, or b) unwrapped and inlined on the path upwards. + */ + Map.Entry> remainingEntry = updatedCollisionContent.get(0); + + K remainingKey = remainingEntry.getKey(); + // Passing the `remainingValues` set as an argument to `MultimapNode#inserted`. The method ensures to unbox + // singleton sets, hence we do not have to deal with the case distinction between 1:1 and 1:n mappings here. + io.usethesource.capsule.Set.Immutable remainingValues = remainingEntry.getValue(); + + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE)); + return CompactSetMultimapNode.nodeOf(null).inserted(null, remainingKey, remainingValues, keyHash, 0, MultimapResult.unchanged()); + } else { + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE)); + return new HashCollisionNode(hash, updatedCollisionContent); + } } else { Function>, Map.Entry>> substitutionMapper = (kImmutableSetEntry) -> { @@ -2698,12 +2716,30 @@ public AbstractSetMultimapNode removed(AtomicReference mutator, K .filter(kImmutableSetEntry -> kImmutableSetEntry != optionalTuple.get()) .collect(Collectors.toList()); - if (values.size() == 1) { - details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE), values); - return new HashCollisionNode(hash, updatedCollisionContent); + if (updatedCollisionContent.size() == 1) { + /* + * Create a root node that stores data for the single remaining key. This node will a) either become the + * new root returned, or b) unwrapped and inlined on the path upwards. + */ + Map.Entry> remainingEntry = updatedCollisionContent.get(0); + + K remainingKey = remainingEntry.getKey(); + // Passing the `remainingValues` set as an argument to `MultimapNode#inserted`. The method ensures to unbox + // singleton sets, hence we do not have to deal with the case distinction between 1:1 and 1:n mappings here. + io.usethesource.capsule.Set.Immutable remainingValues = remainingEntry.getValue(); + + if (values.size() == 1) { + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE), values); + } else { + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE_COLLECTION), values); + } + return CompactSetMultimapNode.nodeOf(null).inserted(null, remainingKey, remainingValues, keyHash, 0, MultimapResult.unchanged()); } else { - details - .modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE_COLLECTION), values); + if (values.size() == 1) { + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE), values); + } else { + details.modified(REMOVED_PAYLOAD, MultimapResult.Modification.flag(REMOVED_KEY, REMOVED_VALUE_COLLECTION), values); + } return new HashCollisionNode(hash, updatedCollisionContent); } }