From dc6804dabc5cc373a1830589274ea47acb416bed Mon Sep 17 00:00:00 2001
From: Thomas Honeyman <hello@thomashoneyman.com>
Date: Sat, 17 Oct 2026 19:01:09 +0000
Subject: [PATCH] Fix balancing when joining maps with unequal heights

---
 src/Data/Map/Internal.purs | 57 +++++++++++++++++++++++---------------
 test/Test/Data/Map.purs    | 23 +++++++++------
 test/Test/Data/Set.purs    |  5 ++++
 3 files changed, 54 insertions(+), 31 deletions(-)

diff --git a/src/Data/Map/Internal.purs b/src/Data/Map/Internal.purs
index 4dc3da6..13d0ebb 100644
--- a/src/Data/Map/Internal.purs
+++ b/src/Data/Map/Internal.purs
@@ -717,31 +717,42 @@ unsafeNode = mkFn4 \k v l r -> case l of
         Node (1 + if h1 > h2 then h1 else h2) (1 + s1 + s2) k v l r
 
 -- | Low-level Node constructor which maintains the balance invariants.
--- | This is unsafe because it assumes the child Maps are ordered.
+-- | This is unsafe because it assumes the child Maps are ordered and balanced.
+-- | Their heights may differ arbitrarily: join along the taller tree's inner
+-- | spine before restoring balance on the way back up.
 unsafeBalancedNode :: forall k v. Fn4 k v (Map k v) (Map k v) (Map k v)
-unsafeBalancedNode = mkFn4 \k v l r -> case l of
-  Leaf ->
-    case r of
-      Leaf ->
-        singleton k v
-      Node rh _ rk rv rl rr
-        | rh > 1 ->
-            runFn7 rotateLeft k v l rk rv rl rr
-      _ ->
-        runFn4 unsafeNode k v l r
-  Node lh _ lk lv ll lr ->
-    case r of
-      Node rh _ rk rv rl rr
-        | rh > lh + 1 ->
-            runFn7 rotateLeft k v l rk rv rl rr
-        | lh > rh + 1 ->
-            runFn7 rotateRight k v lk lv ll lr r
-      Leaf
-        | lh > 1 ->
-            runFn7 rotateRight k v lk lv ll lr r
-      _ ->
-        runFn4 unsafeNode k v l r
+unsafeBalancedNode = mkFn4 \k v l r -> case l, r of
+  Node lh _ lk lv ll lr, _ | lh > height r + 1 ->
+    runFn4 balance lk lv ll (runFn4 unsafeBalancedNode k v lr r)
+  _, Node rh _ rk rv rl rr | rh > height l + 1 ->
+    runFn4 balance rk rv (runFn4 unsafeBalancedNode k v l rl) rr
+  _, _ ->
+    runFn4 unsafeNode k v l r
   where
+  balance :: Fn4 k v (Map k v) (Map k v) (Map k v)
+  balance = mkFn4 \k v l r -> case l of
+    Leaf ->
+      case r of
+        Leaf ->
+          singleton k v
+        Node rh _ rk rv rl rr
+          | rh > 1 ->
+              runFn7 rotateLeft k v l rk rv rl rr
+        _ ->
+          runFn4 unsafeNode k v l r
+    Node lh _ lk lv ll lr ->
+      case r of
+        Node rh _ rk rv rl rr
+          | rh > lh + 1 ->
+              runFn7 rotateLeft k v l rk rv rl rr
+          | lh > rh + 1 ->
+              runFn7 rotateRight k v lk lv ll lr r
+        Leaf
+          | lh > 1 ->
+              runFn7 rotateRight k v lk lv ll lr r
+        _ ->
+          runFn4 unsafeNode k v l r
+
   rotateLeft :: Fn7 k v (Map k v) k v (Map k v) (Map k v) (Map k v)
   rotateLeft = mkFn7 \k v l rk rv rl rr -> case rl of
     Node lh _ lk lv ll lr
diff --git a/test/Test/Data/Map.purs b/test/Test/Data/Map.purs
index 8e5d790..c3994e5 100644
--- a/test/Test/Data/Map.purs
+++ b/test/Test/Data/Map.purs
@@ -370,6 +370,12 @@ mapTests = do
   quickCheck \(TestMap m :: TestMap SmallKey Int) ->
     M.submap Nothing Nothing m === m
 
+  log "submap preserves balance across unequal subtree heights"
+  let
+    rangeInput = M.fromFoldable (map (\k -> Tuple k (k * 3)) (A.range 0 5))
+    bounded = M.submap (Just 0) (Just 5) rangeInput
+  quickCheck' 1 (M.checkValid bounded && bounded == rangeInput)
+
   log "submap with lower bound"
   quickCheck' 1 $
     M.submap (Just B) Nothing (M.fromFoldable [ Tuple A 0, Tuple B 0 ])
@@ -389,14 +395,15 @@ mapTests = do
   quickCheck' 1000 \(TestMap m :: TestMap SmallKey Int) mmin mmax key -> do
     let
       m' = M.submap mmin mmax m
-    ( if
-        ( maybe true (\min -> min <= key) mmin &&
-            maybe true (\max -> max >= key) mmax
-        ) then
-        M.lookup key m == M.lookup key m'
-      else
-        (not (M.member key m'))
-    )
+      inRange =
+        maybe true (\min -> min <= key) mmin &&
+          maybe true (\max -> max >= key) mmax
+      matchesExpected =
+        if inRange then
+          M.lookup key m == M.lookup key m'
+        else
+          not (M.member key m')
+    M.checkValid m' && matchesExpected
       <?> "m: " <> show m <> ", mmin: " <> show mmin <> ", mmax: " <> show mmax <> ", key: " <> show key
 
   log "foldrWithIndex maintains order"
diff --git a/test/Test/Data/Set.purs b/test/Test/Data/Set.purs
index 90479a0..d98b6ed 100644
--- a/test/Test/Data/Set.purs
+++ b/test/Test/Data/Set.purs
@@ -49,3 +49,8 @@ setTests = do
       s1 = S.fromFoldable [ 1, 2, 3, 4, 5 ]
       s2 = S.fromFoldable [ 2, 3 ]
     assert $ S.range (Just 2) (Just 3) s1 == s2
+    let
+      s3 = S.fromFoldable [ 0, 1, 2, 3, 4, 5 ]
+      ranged = S.range (Just 0) (Just 5) s3
+    assert $ ranged == s3
+    assert $ S.checkValid ranged