Update

README.md

@@ -0,0 +1,9 @@
+# Red-black tree
+
+This is a basic implementation of a red-black tree in F#, which is pretty close to being correct by construction.
+
+## Things to fix
+
+* Can we prove that the `failwith` cases indeed will never be hit? (Is it even true?)
+* Can we do better than the awful recursion to `elevateBlack` and `elevateRed`? That is, can we make the type parameter truly phantom, erased at runtime?
+* Can we get some constraints at the type level to help us get the values assigned to the various nodes in the right order?

RedBlackTree.Test/Permutations.fs

@@ -0,0 +1,21 @@
+namespace RedBlackTree.Test
+
+/// A pretty naive implementation of "find all permutations of the given list".
+[<RequireQualifiedAccess>]
+module Permutations =
+    let rec private zipper (acc : 'a list) (results) (xs : 'a list) : ('a list * 'a list) seq =
+        match xs with
+        | [] -> seq { yield (acc, []) ; yield! results }
+        | x :: xs ->
+           zipper (x :: acc) ((acc, x :: xs) :: results) xs
+
+    let private insertions (y : 'a) (xs : 'a list) =
+        zipper [] [] xs
+        |> Seq.map (fun (pre, post) -> List.rev pre @ (y :: post))
+
+    let rec all (xs : 'a list) : 'a list seq =
+        match xs with
+        | [] -> Seq.singleton []
+        | x1 :: xs ->
+            all xs |> Seq.collect (insertions x1)
+

RedBlackTree.Test/RedBlackTree.Test.fsproj

@@ -13,7 +13,9 @@
     </ItemGroup>
 
     <ItemGroup>
+        <Compile Include="Permutations.fs" />
         <Compile Include="TestRedBlackTree.fs" />
+        <Compile Include="TestPermutations.fs" />
     </ItemGroup>
 
     <ItemGroup>

RedBlackTree.Test/TestPermutations.fs

@@ -0,0 +1,28 @@
+namespace RedBlackTree.Test
+
+open NUnit.Framework
+open FsUnitTyped
+
+[<TestFixture>]
+module TestPermutations =
+
+    let private factorial (n : int) =
+        let rec go acc i = if i <= 1 then acc else go (acc * i) (i - 1)
+        go 1 n
+
+    [<Test>]
+    let ``Test factorial`` () =
+        [1..5]
+        |> List.map factorial
+        |> shouldEqual [1 ; 2 ; 6 ; 24 ; 120]
+
+    [<TestCase 5>]
+    let ``Test permutations`` (n : int) =
+        let allPerms =
+            Permutations.all [1..n]
+            |> Seq.toList
+        let fact = factorial n
+        Set.ofList allPerms |> Set.count |> shouldEqual fact
+        allPerms |> List.length |> shouldEqual fact
+        allPerms |> List.iter (fun i -> Set.ofList i |> shouldEqual (Set.ofList [1..n]))
+

RedBlackTree.Test/TestRedBlackTree.fs

@@ -8,26 +8,31 @@ open FsCheck
 [<TestFixture>]
 module TestRedBlackTree =
 
-    [<Test>]
-    let ``First`` () =
-        [0 ; 1]
+    let propertyCases =
+        [
+            [0 ; 1]
+            [2 ; -1 ; 1 ; 0]
+            [3 ; 0 ; -1 ; 2 ; 1]
+        ]
+        |> List.map TestCaseData
+
+    [<TestCaseSource "propertyCases">]
+    let ``Examples found by property-based testing`` (l : int list) =
+        l
         |> List.fold RedBlackTree.add RedBlackTree.empty
         |> RedBlackTree.toList
-        |> shouldEqual (Set.ofList [0 ; 1] |> Set.toList)
+        |> shouldEqual (Set.ofList l |> Set.toList)
 
-    [<Test>]
-    let ``Second`` () =
-        [2 ; -1 ; 1 ; 0]
-        |> List.fold RedBlackTree.add RedBlackTree.empty
-        |> RedBlackTree.toList
-        |> shouldEqual (Set.ofList [2 ; -1 ; 1 ; 0] |> Set.toList)
-
-    [<Test>]
-    let ``Third`` () =
-        [3 ; 0 ; -1 ; 2 ; 1]
-        |> List.fold (RedBlackTree.add) RedBlackTree.empty
-        |> RedBlackTree.toList
-        |> shouldEqual (Set.ofList [-1 ; 0 ; 1 ; 2 ; 3] |> Set.toList)
+    [<TestCase 11>]
+    let ``Exhaustive test`` (n : int) =
+        for perm in Permutations.all [1..n] do
+            let rbt =
+                perm
+                |> List.fold RedBlackTree.add RedBlackTree.empty
+            if rbt |> RedBlackTree.toList <> [1..n] then failwithf "Correctness error: %+A produced %+A" perm rbt
+            let balance = RedBlackTree.balanceFactor rbt
+            if balance.Longest >= balance.Shortest * 2 then
+                failwithf "Unbalanced! %+A produced %+A (balance: %+A)"  perm rbt balance
 
     [<Test>]
     let ``Property-based test`` () =

RedBlackTree.sln.DotSettings

@@ -0,0 +1,2 @@
+<wpf:ResourceDictionary xml:space="preserve" xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml" xmlns:s="clr-namespace:System;assembly=mscorlib" xmlns:ss="urn:shemas-jetbrains-com:settings-storage-xaml" xmlns:wpf="http://schemas.microsoft.com/winfx/2006/xaml/presentation">
+	<s:Boolean x:Key="/Default/UserDictionary/Words/=Rebalance/@EntryIndexedValue">True</s:Boolean></wpf:ResourceDictionary>

RedBlackTree/RedBlackTree.fs

@@ -60,6 +60,7 @@ module BlackNodeCrate =
 [<RequireQualifiedAccess>]
 module RedBlackTree =
 
+    /// A red-black tree holding no data.
     let empty<'a when 'a : comparison> : RedBlackTree<'a> = RedBlackTree.BlackRoot (BlackNode.Leaf)
 
     let rec private elevateBlack<'a, 'depth when 'a : comparison>
@@ -253,11 +254,11 @@ module RedBlackTree =
         | AdditionResult.Black node -> RedBlackTree.BlackRoot node
         | AdditionResult.Red node -> RedBlackTree.RedRoot node
         | AdditionResult.NeedsRebalance info ->
-            RedNode
-                (BlackNode.BlackBlackNode (elevateBlack info.Left, elevateBlack info.Middle, info.Lower),
-                 info.Right,
+            BlackNode.RedBlackNode
+                (RedNode.RedNode (elevateBlack info.Left, elevateBlack info.Middle, ValueAtDepth.elevate info.Lower),
+                 elevateBlack info.Right,
                  info.Upper)
-            |> RedBlackTree.RedRoot
+            |> RedBlackTree.BlackRoot
 
     let rec private findBlack<'a, 'depth when 'a : comparison>
         (tree : BlackNode<'a, 'depth>)
@@ -363,3 +364,39 @@ module RedBlackTree =
 
     let toList<'a when 'a : comparison> (tree : RedBlackTree<'a>) : 'a list =
         fold (fun ls a -> a :: ls) [] tree |> List.rev
+
+    let rec private balanceFactorBlack<'a, 'depth when 'a : comparison> (node : BlackNode<'a, 'depth>) : int * int =
+        match node with
+        | BlackNode.Leaf -> 0, 0
+        | BlackNode.BlackBlackNode (left, right, _) ->
+            let (min1, max1) = balanceFactorBlack left
+            let (min2, max2) = balanceFactorBlack right
+            (min min1 min2, max max1 max2)
+        | BlackNode.BlackRedNode (left, right, _) ->
+            let (min1, max1) = balanceFactorBlack left
+            let (min2, max2) = balanceFactorRed right
+            (min min1 min2, max max1 max2)
+        | BlackNode.RedBlackNode (left, right, _) ->
+            let (min1, max1) = balanceFactorRed left
+            let (min2, max2) = balanceFactorBlack right
+            (min min1 min2, max max1 max2)
+        | BlackNode.RedRedNode (left, right, _) ->
+            let (min1, max1) = balanceFactorRed left
+            let (min2, max2) = balanceFactorRed right
+            (min min1 min2, max max1 max2)
+        |> fun (a, b) -> (a + 1, b + 1)
+
+    and private balanceFactorRed<'a, 'depth when 'a : comparison> (node : RedNode<'a, 'depth>) : int * int =
+        match node with
+        | RedNode (left, right, _) ->
+            let (min1, max1) = balanceFactorBlack left
+            let (min2, max2) = balanceFactorBlack right
+            (min min1 min2, max max1 max2)
+        |> fun (a, b) -> (a + 1, b + 1)
+
+    /// Answer the question: how long is the longest path through the graph, and the shortest?
+    let balanceFactor<'a when 'a : comparison> (tree : RedBlackTree<'a>) : {| Longest : int ; Shortest : int |} =
+        match tree with
+        | RedBlackTree.RedRoot red -> balanceFactorRed red
+        | RedBlackTree.BlackRoot black -> balanceFactorBlack black
+        |> fun (short, long) -> {| Shortest = short ; Longest = long |}