diff --git a/compiler/gen/tests/gen_primitives.rs b/compiler/gen/tests/gen_primitives.rs
index 58bec5323e..53c276484e 100644
--- a/compiler/gen/tests/gen_primitives.rs
+++ b/compiler/gen/tests/gen_primitives.rs
@@ -1278,6 +1278,109 @@ mod gen_primitives {
     }
 
     #[test]
+    #[ignore]
+    fn rbtree_balance_inc_dec() {
+        // TODO does not define a variable correctly, but all is well with the type signature
+        assert_non_opt_evals_to!(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                NodeColor : [ Red, Black ]
+
+                Dict k : [ Node NodeColor k (Dict k)  (Dict k), Empty ]
+
+                # balance : NodeColor, k, Dict k,  Dict k -> Dict k
+                balance = \color, key, left, right ->
+                  when right is
+                    Node Red rK rLeft rRight ->
+                      when left is
+                        Node Red _ _ _ ->
+                          Node
+                            Red
+                            key
+                            Empty
+                            Empty
+
+                        _ ->
+                          Node color rK (Node Red key left rLeft) rRight
+
+                    _ ->
+                        Empty
+
+                main : Dict Int
+                main =
+                    balance Red 0 Empty Empty
+                "#
+            ),
+            0,
+            i64
+        );
+    }
+
+    #[test]
+    fn rbtree_balance_3() {
+        assert_non_opt_evals_to!(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                Dict k : [ Node k (Dict k) (Dict k), Empty ]
+
+                balance : k, Dict k -> Dict k
+                balance = \key, left ->
+                    Node key left Empty
+
+                main : Dict Int
+                main =
+                    balance 0 Empty
+                "#
+            ),
+            1,
+            i64
+        );
+    }
+
+    #[test]
+    fn rbtree_balance_2() {
+        assert_non_opt_evals_to!(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                NodeColor : [ Red, Black ]
+
+                Dict k : [ Node NodeColor k (Dict k), Empty ]
+
+                balance : NodeColor, k, Dict k,  Dict k -> Dict k
+                balance = \color, key, left, right ->
+                  when right is
+                    Node Red rK _ ->
+                      when left is
+                        Node Red _ _  ->
+                          Node
+                            Red
+                            key
+                            Empty
+
+                        _ ->
+                          Node color rK (Node Red key left )
+
+                    _ ->
+                        Empty
+
+                main : Dict Int
+                main =
+                    balance Red 0 Empty Empty
+                "#
+            ),
+            0,
+            i64
+        );
+    }
+
+    #[test]
+    #[ignore]
     fn rbtree_balance() {
         assert_non_opt_evals_to!(
             indoc!(
diff --git a/compiler/solve/tests/solve_expr.rs b/compiler/solve/tests/solve_expr.rs
index 2a33175150..27d58428b5 100644
--- a/compiler/solve/tests/solve_expr.rs
+++ b/compiler/solve/tests/solve_expr.rs
@@ -3519,4 +3519,99 @@ mod solve_expr {
             "Int, Int, List (Num a), Int, Num a -> [ Pair Int (List (Num a)) ]",
         );
     }
+
+    #[test]
+    fn rbtree_old_balance_simplified() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                Dict k : [ Node k (Dict k) (Dict k), Empty ]
+
+                balance : k, Dict k -> Dict k
+                balance = \key, left ->
+                    Node key left Empty
+
+                main : Dict Int
+                main =
+                    balance 0 Empty
+                "#
+            ),
+            "Dict Int",
+        );
+    }
+
+    #[test]
+    fn rbtree_balance_simplified() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                Dict k : [ Node k (Dict k) (Dict k), Empty ]
+
+                node = \x,y,z -> Node x y z
+
+                balance : k, Dict k -> Dict k
+                balance = \key, left ->
+                    node key left Empty
+
+                main : Dict Int
+                main =
+                    balance 0 Empty
+                "#
+            ),
+            "Dict Int",
+        );
+    }
+
+    #[test]
+    fn rbtree_balance() {
+        infer_eq_without_problem(
+            indoc!(
+                r#"
+                app Test provides [ main ] imports []
+
+                NodeColor : [ Red, Black ]
+
+                Dict k v : [ Node NodeColor k v (Dict k v) (Dict k v), Empty ]
+
+                balance : NodeColor, k, v, Dict k v, Dict k v -> Dict k v
+                balance = \color, key, value, left, right ->
+                  when right is
+                    Node Red rK rV rLeft rRight ->
+                      when left is
+                        Node Red lK lV lLeft lRight ->
+                          Node
+                            Red
+                            key
+                            value
+                            (Node Black lK lV lLeft lRight)
+                            (Node Black rK rV rLeft rRight)
+
+                        _ ->
+                          Node color rK rV (Node Red key value left rLeft) rRight
+
+                    _ ->
+                      when left is
+                        Node Red lK lV (Node Red llK llV llLeft llRight) lRight ->
+                          Node
+                            Red
+                            lK
+                            lV
+                            (Node Black llK llV llLeft llRight)
+                            (Node Black key value lRight right)
+
+                        _ ->
+                          Node color key value left right
+
+                main : Dict Int Int
+                main =
+                    balance Red 0 0 Empty Empty
+                "#
+            ),
+            "Dict Int Int",
+        );
+    }
 }