From 5d4bc2ed62475062a09d76a037ca7a72ad3cbd45 Mon Sep 17 00:00:00 2001
From: j6carey <jcarey@arista.com>
Date: Wed, 5 Jun 2024 12:44:38 -0700
Subject: [PATCH] Fix Issue 1004: effect order in filterAMissing (#1005)

The order of Applicative effects in filterAMissing
was incorrect, causing the order of effects to
differ from key order and be influenced by how
the binary tree was balanced.

The fix is to arrange that effects arising from
the key and value at an internal node come after
those in its left branch instead of before.
(Regardless of this fix such effects come before
those effects arising from the right branch.)

This change also expands test coverage
to detect a regression of this fix.
---
 containers-tests/tests/map-properties.hs | 46 +++++++++++++++++++++---
 containers/src/Data/Map/Internal.hs      |  6 ++--
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/containers-tests/tests/map-properties.hs b/containers-tests/tests/map-properties.hs
index 1a47bfa57..df03dab4f 100644
--- a/containers-tests/tests/map-properties.hs
+++ b/containers-tests/tests/map-properties.hs
@@ -1212,13 +1212,49 @@ prop_mergeWithKeyModel xs ys
 -- This uses the instance
 --     Monoid a => Applicative ((,) a)
 -- to test that effects are sequenced in ascending key order.
-prop_mergeA_effects :: UMap -> UMap -> Property
-prop_mergeA_effects xs ys
+prop_mergeA_effects :: WhenMissingSpec -> WhenMissingSpec -> WhenMatchedSpec -> UMap -> UMap -> Property
+prop_mergeA_effects onlyLeft onlyRight both xs ys
   = effects === sort effects
   where
-    (effects, _m) = mergeA whenMissing whenMissing whenMatched xs ys
-    whenMissing = traverseMissing (\k _ -> ([k], ()))
-    whenMatched = zipWithAMatched (\k _ _ -> ([k], ()))
+    (effects, _m) = mergeA (whenMissing onlyLeft) (whenMissing onlyRight) (whenMatched both) xs ys
+    whenMissing spec = case spec of
+      DropMissing -> dropMissing
+      PreserveMissing -> preserveMissing
+      FilterMissing -> filterMissing (\_ _ -> False)
+      FilterAMissing -> filterAMissing (\k _ -> ([k], False))
+      MapMissing -> mapMissing (\_ _ -> ())
+      TraverseMissing -> traverseMissing (\k _ -> ([k], ()))
+      MapMaybeMissing -> mapMaybeMissing (\_ _ -> Nothing)
+      TraverseMaybeMissing -> traverseMaybeMissing (\k _ -> ([k], Nothing))
+    whenMatched spec = case spec of
+      ZipWithMatched -> zipWithMatched (\_ _ _ -> ())
+      ZipWithAMatched -> zipWithAMatched (\k _ _ -> ([k], ()))
+      ZipWithMaybeMatched -> zipWithMaybeMatched (\_ _ _ -> Nothing)
+      ZipWithMaybeAMatched -> zipWithMaybeAMatched (\k _ _ -> ([k], Nothing))
+
+data WhenMissingSpec
+  = DropMissing
+  | PreserveMissing
+  | FilterMissing
+  | FilterAMissing
+  | MapMissing
+  | TraverseMissing
+  | MapMaybeMissing
+  | TraverseMaybeMissing
+  deriving (Bounded, Enum, Show)
+
+instance Arbitrary WhenMissingSpec where
+  arbitrary = arbitraryBoundedEnum
+
+data WhenMatchedSpec
+  = ZipWithMatched
+  | ZipWithAMatched
+  | ZipWithMaybeMatched
+  | ZipWithMaybeAMatched
+  deriving (Bounded, Enum, Show)
+
+instance Arbitrary WhenMatchedSpec where
+  arbitrary = arbitraryBoundedEnum
 
 ----------------------------------------------------------------
 
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index 99531e3a1..199204fe0 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/containers/src/Data/Map/Internal.hs
@@ -2891,12 +2891,12 @@ filterWithKey p t@(Bin _ kx x l r)
 filterWithKeyA :: Applicative f => (k -> a -> f Bool) -> Map k a -> f (Map k a)
 filterWithKeyA _ Tip = pure Tip
 filterWithKeyA p t@(Bin _ kx x l r) =
-  liftA3 combine (p kx x) (filterWithKeyA p l) (filterWithKeyA p r)
+  liftA3 combine (filterWithKeyA p l) (p kx x) (filterWithKeyA p r)
   where
-    combine True pl pr
+    combine pl True pr
       | pl `ptrEq` l && pr `ptrEq` r = t
       | otherwise = link kx x pl pr
-    combine False pl pr = link2 pl pr
+    combine pl False pr = link2 pl pr
 
 -- | \(O(\log n)\). Take while a predicate on the keys holds.
 -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map,