-- | This module implements a least-recently-used structure. Conceptually,
--   it's a list of values where the head of the list is the most recently used
--   value. When a value is used, it's moved from its place in the list to the
--   head of the list. The last element in the list is thus the
--   least-recently-used value.
--
--   This structure is often used in caches to decide which values to evict when
--   the cache becomes full.
--
--   This module uses a Map to implement the LRU efficiently and thus there's the
--   requirement that the elements of the LRU be instances of Ord, which a more
--   general (but slower) LRU implementation could avoid.

module Data.LRU (LRU, empty, null, size, hit, delete, toList, member, pop, last) where

import Prelude hiding (last, null)
import qualified Data.Map as Map
import Data.Maybe (fromJust, isNothing)

-- Impementation notes:
--
-- The LRU is a Map mapping the element to an Item. An Item is a link in a
-- double-linked list. We also hold the values of the elements at the head and
-- tail of the list. So, to get the list you can take the head value and look
-- that up in the map. That will give you the prev and next pointers (the prev
-- pointer will be Nothing in the case of the head element). Taking the next
-- pointer value, you can lookup again to get the next element in the list
-- and so on.

-- | A double-linked list element. Contains the values of the elements before
-- and after this element (resp)
data (Ord a) => Item a = Item (Maybe a) (Maybe a)

-- | An LRU. Contains the head element, last element and the map from elements
-- to their Items
data (Ord a) => LRU a = LRU (Maybe a) (Maybe a) (Map.Map a (Item a))

instance (Ord a, Show a) => Show (LRU a) where
  show x = "LRU" ++ show (toList x)

-- | Returns an empty LRU. O(1)
empty :: (Ord a) => LRU a
empty = LRU Nothing Nothing Map.empty

-- | Returns True iff the LRU is empty. O(1)
null :: (Ord a) => LRU a -> Bool
null (LRU Nothing _ _) = True
null _ = False

-- | Returns the number of elements in the LRU. O(1)
size :: (Ord a) => LRU a -> Int
size (LRU _ _ map) = Map.size map

-- | Utility function. Updates a map entry so that the prev pointer points to
--   a new value
updatePrev :: (Ord a)
              => a  -- ^ the entry to update
              -> (Maybe a)  -- ^ the new prev pointer value
              -> Map.Map a (Item a)  -- ^ the map to update
              -> Map.Map a (Item a)
updatePrev toupdate newvalue =
  Map.adjust (\(Item _ next) -> Item newvalue next) toupdate

-- | Utility function. Same as updatePrev, but with the other pointer
updateNext :: (Ord a) => a -> (Maybe a) -> Map.Map a (Item a) -> Map.Map a (Item a)
updateNext toupdate newvalue =
  Map.adjust (\(Item prev _) -> Item prev newvalue) toupdate

-- | Inserts a new element into an LRU. O(log n) (not exported, see hit)
insert :: (Ord a) => a -> LRU a -> LRU a
insert value (LRU head tail map) = LRU newhead newtail newmap where
  (newmap, newhead, newtail) =
    case tail of
         -- We are inserting the first element
         Nothing -> (Map.singleton value (Item Nothing Nothing), Just value, Just value)
         -- We are inserting into a non-empty LRU. We change the prev pointer
         -- of the current head to point to the new head and insert the new
         -- head, with a next value pointing to the previous head.
         (Just x) -> (Map.insert value (Item Nothing head) $
                      updatePrev (fromJust head) (Just value) map,
                      Just value, tail)

-- | Insert a value into an LRU. If the value is already in the LRU, it's
--   moved to the head of the list. O(log n)
hit :: (Ord a) => a -> LRU a -> LRU a
hit value lru@(LRU head tail map) = r where
  r = case Map.lookup value map of
           Just (Item (Just prev) (Just next)) -> -- it's in the middle
             LRU (Just value) tail $ updateHit $ updateNext prev (Just next) $
             updatePrev next (Just prev) $ updatePrev (fromJust head) (Just value) map
           Just (Item (Just prev) Nothing) ->  -- it's the very last element
             LRU (Just value) (Just prev) $ updateHit $ updateNext prev Nothing $ updatePrev (fromJust head) (Just value) map
           Just (Item Nothing _) -> lru  -- it's already at the top
           Nothing -> insert value lru
  -- | Set the item for a given value so that it's the head of the list, pointing to the old head
  updateHit = Map.insert value (Item Nothing head)

-- | Returns a list of the members of the LRU in order, newest first. O(n(log n))
toList :: (Ord a) => LRU a -> [a]
toList (LRU head tail map) =
  if Map.null map
    then []
    else (fromJust head) : (inner (fromJust head)) where
         inner current = case next of
                              (Just value) -> value : (inner value)
                              Nothing -> []
                         where
                         Just (Item prev next) = Map.lookup current map

-- | Returns true iff the given element is in the LRU. O(log n)
member :: (Ord a) => a -> LRU a -> Bool
member value (LRU _ _ map) = Map.member value map

-- | Removes an element from the LRU, if it exists. O(log n)
delete :: (Ord a) => a -> LRU a -> LRU a
delete value lru@(LRU head tail map) =
  case head of
    Nothing -> lru  -- empty case
    (Just x)
      | x == value ->
        -- it's the head item
        let (Just (Item _ next), map') = Map.updateLookupWithKey (const $ const Nothing) value map in
        LRU next (if isNothing next then Nothing else tail) map'
      | otherwise ->
          case tail of
            (Just x)
              | x == value ->  -- removing the tail
                pop lru
              | otherwise ->
                  -- removing a middle element
                  LRU head tail $ updateNext (fromJust prev) next $ updatePrev (fromJust next) prev map' where
                  (Just (Item prev next), map') = Map.updateLookupWithKey (const $ const Nothing) value map

-- | Return the last element of the LRU. O(1)
last :: (Ord a, Monad m) => LRU a -> m a
last (LRU _ tail _) = if isNothing tail then fail "Empty LRU" else return $ fromJust tail

-- | Remove the last element of the LRU. Errors out if the LRU is empty. O(log n)
pop :: (Ord a) => LRU a -> LRU a
pop (LRU head (Just tailvalue) map) = newlru where
  newlru = if isOnlyElement
             then empty
             else LRU head newtail newmap
  Just (Item prev _) = Map.lookup tailvalue map
  isOnlyElement = isNothing prev
  newtail = prev
  newmap = updateNext (fromJust prev) Nothing $ Map.delete tailvalue map