From 8e0d144a1c1eb76122a902a443100e8b9d84845e Mon Sep 17 00:00:00 2001
From: Tilo Wiklund <tilo.wiklund@sensmetry.com>
Date: Mon, 8 Nov 2021 16:08:01 +0100
Subject: [PATCH] Implemented a lookupAround for IntervalMap

---
 src/Data/IntervalMap/Base.hs   | 14 ++++++++++++++
 src/Data/IntervalMap/Lazy.hs   |  1 +
 src/Data/IntervalMap/Strict.hs |  1 +
 test/TestIntervalMap.hs        |  9 +++++++++
 4 files changed, 25 insertions(+)

diff --git a/src/Data/IntervalMap/Base.hs b/src/Data/IntervalMap/Base.hs
index 7322cd7..033eaf1 100644
--- a/src/Data/IntervalMap/Base.hs
+++ b/src/Data/IntervalMap/Base.hs
@@ -32,6 +32,7 @@ module Data.IntervalMap.Base
   , member
   , notMember
   , lookup
+  , lookupAround
   , findWithDefault
   , span
 
@@ -246,6 +247,19 @@ lookup k (IntervalMap m) =
     Just (_, (i, a)) | k `Interval.member` i -> Just a
     _ -> Nothing
 
+-- | Lookup the value and interval at a key in the map.
+--
+-- The function will return the corresponding value and a
+-- surrounding interval as @('Just' (interval, value))@,
+-- or 'Nothing' if the key isn't in the map.  The surrounding interval
+-- @interval@ is such that it contains @value@ and such that @k@ is
+-- constantly equal to @value@ on this interval.
+lookupAround :: Ord k => k -> IntervalMap k a -> Maybe (Interval k, a)
+lookupAround k (IntervalMap m) =
+  case Map.lookupLE (LB (Finite k, Interval.Closed)) m of
+    Just (_, (i, a)) | k `Interval.member` i -> Just (i, a)
+    _ -> Nothing
+
 -- | The expression @('findWithDefault' def k map)@ returns
 -- the value at key @k@ or returns default value @def@
 -- when the key is not in the map.
diff --git a/src/Data/IntervalMap/Lazy.hs b/src/Data/IntervalMap/Lazy.hs
index fc4a55a..da77d9a 100644
--- a/src/Data/IntervalMap/Lazy.hs
+++ b/src/Data/IntervalMap/Lazy.hs
@@ -43,6 +43,7 @@ module Data.IntervalMap.Lazy
   , member
   , notMember
   , lookup
+  , lookupAround
   , findWithDefault
   , span
 
diff --git a/src/Data/IntervalMap/Strict.hs b/src/Data/IntervalMap/Strict.hs
index 045b8d8..9f06246 100644
--- a/src/Data/IntervalMap/Strict.hs
+++ b/src/Data/IntervalMap/Strict.hs
@@ -44,6 +44,7 @@ module Data.IntervalMap.Strict
   , member
   , notMember
   , lookup
+  , lookupAround
   , findWithDefault
   , span
 
diff --git a/test/TestIntervalMap.hs b/test/TestIntervalMap.hs
index f4004e0..dcbedf7 100644
--- a/test/TestIntervalMap.hs
+++ b/test/TestIntervalMap.hs
@@ -148,6 +148,15 @@ prop_insert_lookup =
           Just k -> IML.lookup k (IML.insert i a m) == Just a
           Nothing -> True
 
+prop_insert_lookupAround_just_inserted :: Property
+prop_insert_lookupAround_just_inserted =
+  forAll arbitrary $ \(m :: IntervalMap Rational Integer) ->
+    forAll arbitrary $ \i ->
+      forAll arbitrary $ \a ->
+        case Interval.pickup i of
+          Just k -> IML.lookupAround k (IML.insert i a m) == Just (i, a)
+          Nothing -> True
+
 prop_insert_bang :: Property
 prop_insert_bang =
   forAll arbitrary $ \(m :: IntervalMap Rational Integer) ->