Fix optimization-dependent invalid Map decoding

changelog.md

@@ -2,6 +2,10 @@ For the latest version of this document, please see [https://github.com/haskell/
 
 ### Unreleased - Future date
 
+* Fix optimized decoding through `fmap` on `FromJSONKeyFunction` producing an
+  invalid `Map` when the target key type has a different `Ord` instance. This
+  is a patch-level behavioral correction with no public API change.
+
 ### 2.3.1.0 - 2026-07-05
 
 * Add `FromJSONKey` instance for `Data.Fixed`.

src/Data/Aeson/Internal/Functions.hs

@@ -35,8 +35,23 @@ mapKey :: (Eq k2, Hashable k2) => (k1 -> k2) -> H.HashMap k1 v -> H.HashMap k2 v
 mapKey fk = mapKeyVal fk id
 {-# INLINE mapKey #-}
 
--- | Transform the keys of a 'M.Map'.
+-- | Transform the keys of a 'M.Map', preserving its tree shape only when the
+-- projected keys prove strictly ascending.
 mapKeyO :: (Ord k2) => (k1 -> k2) -> M.Map k1 v -> M.Map k2 v
-mapKeyO fk = mapKeyValO fk id
-{-# INLINE mapKeyO #-}
+mapKeyO fk sourceMap
+    | keysRemainAscending = M.mapKeysMonotonic fk sourceMap
+    | otherwise = M.mapKeysWith (\_ earlierValue -> earlierValue) fk sourceMap
+  where
+    keysRemainAscending = case M.foldrWithKey checkKeyOrder (Just Nothing) sourceMap of
+        Nothing -> False
+        Just _  -> True
 
+    -- The outer 'Nothing' is an ordering obstruction; the inner one is the
+    -- empty suffix before its greatest projected key is seen.
+    checkKeyOrder sourceKey _ maybeNextKey = do
+        nextKey <- maybeNextKey
+        let currentKey = fk sourceKey
+        if maybe True (currentKey <) nextKey
+        then Just (Just currentKey)
+        else Nothing
+{-# INLINE mapKeyO #-}

src/Data/Aeson/Types/FromJSON.hs

@@ -454,7 +454,7 @@ class FromJSONKey a where
 --   For performance reasons, these exist as three options instead of one.
 data FromJSONKeyFunction a where
     FromJSONKeyCoerce :: Coercible Text a => FromJSONKeyFunction a
-      -- ^ uses 'coerce', we expect that 'Hashable' and 'Ord' instance are compatible.
+      -- ^ uses 'coerce', we expect that 'Eq', 'Hashable', and 'Ord' instances are compatible.
     FromJSONKeyText :: !(Text -> a) -> FromJSONKeyFunction a
       -- ^ conversion from 'Text' that always succeeds
     FromJSONKeyTextParser :: !(Text -> Parser a) -> FromJSONKeyFunction a
@@ -470,31 +470,24 @@ instance Functor FromJSONKeyFunction where
     fmap h (FromJSONKeyValue f)      = FromJSONKeyValue (fmap h . f)
 
 -- | Construct 'FromJSONKeyFunction' for types coercible from 'Text'. This
--- conversion is still unsafe, as 'Hashable' and 'Eq' instances of @a@ should be
--- compatible with 'Text' i.e. hash values should be equal for wrapped values as well.
--- This property will always be maintained if the 'Hashable' and 'Eq' instances
--- are derived with generalized newtype deriving.
--- compatible with 'Text' i.e. hash values be equal for wrapped values as well.
+-- conversion is still unsafe, as 'Hashable', 'Eq', and 'Ord' instances of @a@
+-- should be compatible with 'Text'. This property will always be maintained if
+-- these instances are derived with generalized newtype deriving.
 --
 -- On pre GHC 7.8 this is unconstrained function.
 fromJSONKeyCoerce ::
     Coercible Text a =>
     FromJSONKeyFunction a
 fromJSONKeyCoerce = FromJSONKeyCoerce
 
--- | Semantically the same as @coerceFromJSONKeyFunction = fmap coerce = coerce@.
+-- | Coerce the result type without changing the parsing strategy.
 --
 -- See note on 'fromJSONKeyCoerce'.
 coerceFromJSONKeyFunction ::
     Coercible a b =>
     FromJSONKeyFunction a -> FromJSONKeyFunction b
 coerceFromJSONKeyFunction = coerce
 
-{-# RULES
-  "FromJSONKeyCoerce: fmap coerce" forall x .
-                                   fmap coerce x = coerceFromJSONKeyFunction x
-  #-}
-
 -- | Same as 'fmap'. Provided for the consistency with 'ToJSONKeyFunction'.
 mapFromJSONKeyFunction :: (a -> b) -> FromJSONKeyFunction a -> FromJSONKeyFunction b
 mapFromJSONKeyFunction = fmap
@@ -2491,7 +2484,7 @@ instance FromJSON b => FromJSON (Tagged a b) where
 
 instance FromJSONKey b => FromJSONKey (Tagged a b) where
     fromJSONKey = coerceFromJSONKeyFunction (fromJSONKey :: FromJSONKeyFunction b)
-    fromJSONKeyList = (fmap . fmap) Tagged fromJSONKeyList
+    fromJSONKeyList = coerceFromJSONKeyFunction (fromJSONKeyList :: FromJSONKeyFunction [b])
 
 -------------------------------------------------------------------------------
 -- these

tests/UnitTests/FromJSONKey.hs

@@ -1,49 +1,74 @@
+{-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedStrings #-}
 module UnitTests.FromJSONKey (fromJSONKeyTests) where
 
-import Test.Tasty (TestTree, testGroup)
-import Test.Tasty.HUnit (testCase, Assertion, assertFailure)
-import Data.Text (Text)
-import Data.Tagged (Tagged)
 import Control.Applicative (Const)
-
 import Data.Aeson
+import Data.Map.Strict (Map)
+import Data.Ord (Down(Down))
+import Data.Tagged (Tagged)
+import Data.Text (Text)
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit (Assertion, assertFailure, testCase, (@?=))
+import qualified Data.Map.Strict as Map
+import qualified Data.Text as Text
 
 newtype MyText = MyText Text
     deriving (FromJSONKey)
 
 newtype MyText' = MyText' Text
+    deriving FromJSON via Text
 
 instance FromJSONKey MyText' where
     fromJSONKey = fmap MyText' fromJSONKey
-    fromJSONKeyList = error "not used"
+
+newtype ReverseText = ReverseText Text
+    deriving (Eq, Ord) via Down Text
+    deriving FromJSON via Text
+
+instance FromJSONKey ReverseText where
+    fromJSONKey = fmap ReverseText fromJSONKey
+
+newtype FoldedText = FoldedText Text
+    deriving (Eq, Ord, Show) via Text
+    deriving FromJSON via Text
+
+instance FromJSONKey FoldedText where
+    fromJSONKey = fmap (FoldedText . Text.toCaseFold) fromJSONKey
 
 fromJSONKeyTests :: TestTree
-fromJSONKeyTests = testGroup "FromJSONKey" $ fmap (testCase "-") fromJSONKeyAssertions
+fromJSONKeyTests = testGroup "FromJSONKey"
+    [ testGroup "strategies" $ fmap (testCase "-") fromJSONKeyAssertions
+    , testCase "fmap decoding preserves Map ordering" assertDecodedMapIsValid
+    , testCase "fmap decoding rebuilds colliding Map keys" assertDecodedCollidingMap
+    ]
 
 fromJSONKeyAssertions :: [Assertion]
 fromJSONKeyAssertions =
     [ assertIsCoerce  "Text"            (fromJSONKey :: FromJSONKeyFunction Text)
     , assertIsCoerce  "Tagged Int Text" (fromJSONKey :: FromJSONKeyFunction (Tagged Int Text))
     , assertIsCoerce  "MyText"          (fromJSONKey :: FromJSONKeyFunction MyText)
 
-    , assertIsCoerce' "MyText'"         (fromJSONKey :: FromJSONKeyFunction MyText')
+    , assertIsText    "MyText'"         (fromJSONKey :: FromJSONKeyFunction MyText')
     , assertIsCoerce  "Const Text"      (fromJSONKey :: FromJSONKeyFunction (Const Text ()))
     ]
   where
     assertIsCoerce :: String -> FromJSONKeyFunction a -> Assertion
     assertIsCoerce _ FromJSONKeyCoerce = pure ()
     assertIsCoerce n _                 = assertFailure n
 
-    assertIsCoerce' :: String -> FromJSONKeyFunction a -> Assertion
-    assertIsCoerce' _ FromJSONKeyCoerce = pure ()
-    assertIsCoerce' n _                 = pickWithRules (assertFailure n) (pure ())
-
--- | Pick the first when RULES are enabled, e.g. optimisations are on
-pickWithRules
-    :: a -- ^ Pick this when RULES are on
-    -> a -- ^ use this otherwise
-    -> a
-pickWithRules _ = id
-{-# NOINLINE pickWithRules #-}
-{-# RULES "pickWithRules/rule" [0] forall x. pickWithRules x = const x #-}
+    assertIsText :: String -> FromJSONKeyFunction a -> Assertion
+    assertIsText _ (FromJSONKeyText _) = pure ()
+    assertIsText n _                   = assertFailure n
+
+assertDecodedMapIsValid :: Assertion
+assertDecodedMapIsValid = fmap Map.valid decodedMap @?= Just True
+  where
+    decodedMap = decode "{\"a\":\"a\",\"b\":\"b\"}" :: Maybe (Map ReverseText Text)
+
+assertDecodedCollidingMap :: Assertion
+assertDecodedCollidingMap = decodedMap @?= Just (Map.singleton (FoldedText "a") "upper")
+  where
+    decodedMap = decode "{\"A\":\"upper\",\"a\":\"lower\"}" :: Maybe (Map FoldedText Text)