From 2685b5337ecd5e8f2d120718e8acd88c5423c2b1 Mon Sep 17 00:00:00 2001
From: rei <rei.miyasaka@protonmail.com>
Date: Mon, 17 Apr 2023 17:15:01 -0700
Subject: [PATCH] Preliminary Map.binarySearch function with tests

---
 src/FSharp.Core/map.fs                        | 33 +++++++++-
 src/FSharp.Core/map.fsi                       | 17 ++++++
 .../Microsoft.FSharp.Collections/MapModule.fs | 60 +++++++++++++++++++
 3 files changed, 109 insertions(+), 1 deletion(-)

diff --git a/src/FSharp.Core/map.fs b/src/FSharp.Core/map.fs
index 2535dff7522..cb65776fdba 100644
--- a/src/FSharp.Core/map.fs
+++ b/src/FSharp.Core/map.fs
@@ -679,7 +679,31 @@ module MapTree =
                 (m.Key, m.Value)
             else
                 rightmost nd.Right
-
+                
+    let rec binarySearch (comparer: IComparer<'Key>) k lower higher (m: MapTree<'Key, 'Value>) =
+        if isEmpty m then
+            lower, None, higher
+        else
+            let c = comparer.Compare(k, m.Key)
+            if m.Height = 1 then
+                if c = 0 then
+                    lower, Some (m.Key, m.Value), higher
+                elif c < 0 then
+                    lower, None, Some (m.Key, m.Value)
+                else
+                    Some (m.Key, m.Value), None, higher
+            elif c = 0 then
+                let nd = asNode m
+                (if isEmpty nd.Left then None else Some (rightmost nd.Left)),
+                Some (m.Key, m.Value),
+                (if isEmpty nd.Right then None else Some (leftmost nd.Right))
+            elif c > 0 then
+                let nd = asNode m
+                binarySearch comparer k (Some (nd.Key, nd.Value)) higher nd.Right
+            else
+                let nd = asNode m
+                binarySearch comparer k lower (Some (nd.Key, nd.Value)) nd.Left
+                
 [<System.Diagnostics.DebuggerTypeProxy(typedefof<MapDebugView<_, _>>)>]
 [<System.Diagnostics.DebuggerDisplay("Count = {Count}")>]
 [<Sealed>]
@@ -843,6 +867,9 @@ type Map<[<EqualityConditionalOn>] 'Key, [<EqualityConditionalOn; ComparisonCond
 
     member m.MinKeyValue = MapTree.leftmost tree
     member m.MaxKeyValue = MapTree.rightmost tree
+    
+    member m.BinarySearch key =
+        MapTree.binarySearch comparer key None None tree
 
     static member ofList l : Map<'Key, 'Value> =
         let comparer = LanguagePrimitives.FastGenericComparer<'Key>
@@ -1305,3 +1332,7 @@ module Map =
     [<CompiledName("MaxKeyValue")>]
     let maxKeyValue (table: Map<_, _>) =
         table.MaxKeyValue
+        
+    [<CompiledName("BinarySearch")>]
+    let binarySearch (key: 'Key) (table: Map<'Key, 'T>) =
+        table.BinarySearch key
\ No newline at end of file
diff --git a/src/FSharp.Core/map.fsi b/src/FSharp.Core/map.fsi
index abf960909a2..743d7921725 100644
--- a/src/FSharp.Core/map.fsi
+++ b/src/FSharp.Core/map.fsi
@@ -781,3 +781,20 @@ module Map =
     /// </example>
     [<CompiledName("MaxKeyValue")>]
     val maxKeyValue: table: Map<'Key, 'T> -> 'Key * 'T
+
+    /// <summary>Returns a 3-tuple of the item with the closest key below the map, if any;
+    /// the matching item, if any; and the closest key above the map, if any.</summary>
+    ///
+    /// <param name="key">The key to search.</param>
+    /// <param name="table">The input map.</param>
+    ///
+    /// <example>
+    /// <code lang="fsharp">
+    /// let sample = Map [ (10, "a"); (12, "b"); (20, "c"); (22, "d"); (25, "e"); (28, "f")
+    ///                    (30, "g"); (36, "h"); (38, "i"); (40, "j"); (48, "k") ]
+    ///
+    /// sample |> Map.binarySearch 20 // evaluates to Some (12, "b"), Some (20, "c"), Some (25, "e")
+    /// </code>
+    /// </example>
+    [<CompiledName("BinarySearch")>]
+    val binarySearch: key: 'Key -> table: Map<'Key, 'T> -> ('Key * 'T) option * ('Key * 'T) option * ('Key * 'T) option
\ No newline at end of file
diff --git a/tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Collections/MapModule.fs b/tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Collections/MapModule.fs
index 90c14d92532..b517e626655 100644
--- a/tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Collections/MapModule.fs
+++ b/tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Collections/MapModule.fs
@@ -725,4 +725,64 @@ type MapModule() =
         let eptMap = Map.empty
         CheckThrowsKeyNotFoundException (fun () -> Map.maxKeyValue eptMap |> ignore)
                
+        ()
+        
+    [<Fact>]
+    member _.BinarySearch() =
+        let map = Map [ (10, "a"); (12, "b"); (20, "c"); (22, "d"); (25, "e"); (28, "f")
+                        (30, "g"); (36, "h"); (38, "i"); (40, "j"); (48, "k") ]
+        // Produces the following tree at tolerance 2:
+        // 22
+        // |- 12
+        //    |- 10
+        //    |- 20
+        // |- 28
+        //    |- 25
+        //    |- 36
+        //       |- 30
+        //       |- 40
+        //          |- 38
+        //          |- 48
+        
+        // Matches exact and either side at middle of tree
+        let result = Map.binarySearch 12 map
+        let expected = Some (10, "a"), Some (12, "b"), Some (20, "c")
+        Assert.Equal(expected, result)
+        
+        // Matches exact and either side at bottom of tree
+        let result = Map.binarySearch 20 map
+        let expected = Some (12, "b"), Some (20, "c"), Some (22, "d")
+        Assert.Equal(expected, result)
+        
+        // Matches exact and either side at top of tree
+        let result = Map.binarySearch 22 map
+        let expected = Some (20, "c"), Some (22, "d"), Some (25, "e")
+        Assert.Equal(expected, result)
+        
+        // Matches on either side
+        let result = Map.binarySearch 11 map
+        let expected = Some (10, "a"), None, Some (12, "b")
+        Assert.Equal(expected, result)
+        
+        // Only matches to left
+        let result = Map.binarySearch 50 map
+        let expected = Some (48, "k"), None, None
+        Assert.Equal(expected, result)
+        
+        // Only matches to right
+        let result = Map.binarySearch 1 map
+        let expected = None, None, Some (10, "a")
+        Assert.Equal(expected, result)
+        
+        // One-element Map
+        let map = Map [ (1, "a") ]
+        let result = Map.binarySearch 1 map
+        let expected = None, Some (1, "a"), None
+        Assert.Equal(expected, result)
+        
+        // Empty Map
+        let result = Map.binarySearch 1 (Map [])
+        let expected = None, None, None
+        Assert.Equal(expected, result)
+        
         ()
\ No newline at end of file