-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

{-# LANGUAGE BangPatterns #-}

module Utils.ARC
  ( ARCTick
  , ARC, initARC
  , arcAlterF, lookupARC, insertARC
  , ARCHandle, initARCHandle, cachedARC, cachedARC'
  , lookupARCHandle
  , readARCHandle
  , arcRecentSize, arcFrequentSize, arcGhostRecentSize, arcGhostFrequentSize
  , getARCRecentWeight, getARCFrequentWeight
  , describeARC
  , NFDynamic(..), _NFDynamic, DynARC, DynARCHandle
  ) where

import ClassyPrelude

import Data.HashPSQ (HashPSQ)
import qualified Data.HashPSQ as HashPSQ

import Control.Lens

import Type.Reflection
import Text.Show (showString, shows)

import Data.Hashable (Hashed, hashed)

-- https://web.archive.org/web/20210115184012/https://dbs.uni-leipzig.de/file/ARC.pdf
-- https://jaspervdj.be/posts/2015-02-24-lru-cache.html


data NFDynamic where
  NFDynamic :: forall a. NFData a => TypeRep a -> a -> NFDynamic

_NFDynamic :: forall a. (Typeable a, NFData a) => Prism' NFDynamic a
_NFDynamic = prism' toNFDyn fromNFDynamic
  where
    toNFDyn v = NFDynamic typeRep v
    fromNFDynamic (NFDynamic t v)
      | Just HRefl <- t `eqTypeRep` rep = Just v
      | otherwise                       = Nothing
      where rep = typeRep :: TypeRep a

instance NFData NFDynamic where
  rnf (NFDynamic t v) = rnfTypeRep t `seq` rnf v

instance Show NFDynamic where
  showsPrec _ (NFDynamic t _) = showString "<<" . shows t . showString ">>"


newtype ARCTick = ARCTick { _getARCTick :: Word64 }
  deriving (Eq, Ord, Show)
  deriving newtype (NFData)

makeLenses ''ARCTick

data ARC k w v = ARC
  { arcRecent, arcFrequent :: !(HashPSQ (Hashed k) ARCTick (v, w))
  , arcGhostRecent, arcGhostFrequent :: !(HashPSQ (Hashed k) ARCTick ())
  , arcRecentWeight, arcFrequentWeight :: !w
  , arcTargetRecent, arcMaximumWeight :: !w
  , arcMaximumGhost :: !Int
  }

type DynARC k w = ARC (SomeTypeRep, k) w NFDynamic

instance (NFData k, NFData w, NFData v) => NFData (ARC k w v) where
  rnf ARC{..} = rnf arcRecent
          `seq` rnf arcFrequent
          `seq` rnf arcGhostRecent
          `seq` rnf arcGhostFrequent
          `seq` rnf arcRecentWeight
          `seq` rnf arcFrequentWeight
          `seq` rnf arcTargetRecent
          `seq` rnf arcMaximumWeight
          `seq` rnf arcMaximumGhost

describeARC :: Show w
            => ARC k w v
            -> String
describeARC ARC{..} = intercalate ", "
  [ "arcRecent: "         <> show (HashPSQ.size arcRecent)
  , "arcFrequent: "       <> show (HashPSQ.size arcFrequent)
  , "arcGhostRecent: "    <> show (HashPSQ.size arcGhostRecent)
  , "arcGhostFrequent: "  <> show (HashPSQ.size arcGhostFrequent)
  , "arcRecentWeight: "   <> show arcRecentWeight
  , "arcFrequentWeight: " <> show arcFrequentWeight
  , "arcTargetRecent: "   <> show arcTargetRecent
  , "arcMaximumWeight: "  <> show arcMaximumWeight
  , "arcMaximumGhost: "   <> show arcMaximumGhost
  ]

arcRecentSize, arcFrequentSize, arcGhostRecentSize, arcGhostFrequentSize :: ARC k w v -> Int
arcRecentSize = HashPSQ.size . arcRecent
arcFrequentSize = HashPSQ.size . arcFrequent
arcGhostRecentSize = HashPSQ.size . arcGhostRecent
arcGhostFrequentSize = HashPSQ.size . arcGhostFrequent

getARCRecentWeight, getARCFrequentWeight :: ARC k w v -> w
getARCRecentWeight = arcRecentWeight
getARCFrequentWeight = arcFrequentWeight

initialARCTick :: ARCTick
initialARCTick = ARCTick 0

initARC :: forall k w v.
           Integral w
        => Int -- ^ @arcMaximumGhost@
        -> w -- ^ @arcMaximumWeight@
        -> (ARC k w v, ARCTick)
initARC arcMaximumGhost arcMaximumWeight
  | arcMaximumWeight < 0 = error "initARC given negative maximum weight"
  | arcMaximumGhost < 0 = error "initARC given negative maximum ghost size"
  | otherwise = (, initialARCTick) ARC
    { arcRecent = HashPSQ.empty
    , arcFrequent = HashPSQ.empty
    , arcGhostRecent = HashPSQ.empty
    , arcGhostFrequent = HashPSQ.empty
    , arcRecentWeight = 0
    , arcFrequentWeight = 0
    , arcMaximumWeight
    , arcTargetRecent = 0
    , arcMaximumGhost
    }


infixl 6 |-
(|-) :: (Num a, Ord a) => a -> a -> a
(|-) m s
  | s >= m = 0
  | otherwise = m - s

  
arcAlterF :: forall f k w v.
             ( Ord k, Hashable k
             , Functor f
             , Integral w
             , NFData k, NFData w, NFData v
             )
          => k
          -> (Maybe (v, w) -> f (Maybe (v, w)))
          -> ARC k w v
          -> ARCTick -> f (ARC k w v, ARCTick)
-- | Unchecked precondition: item weights are always less than `arcMaximumWeight`
arcAlterF !(force -> unhashedK@(hashed -> k)) f oldARC@ARC{..} now
  | later <= initialARCTick = uncurry (arcAlterF unhashedK f) $ initARC arcMaximumGhost arcMaximumWeight
  | otherwise = (, later) <$> if
    | Just (_p, x@(_, w), arcFrequent') <- HashPSQ.deleteView k arcFrequent
      -> f (Just x) <&> \case
           Nothing -> oldARC
             { arcFrequent = arcFrequent'
             , arcGhostFrequent = HashPSQ.insert k now () arcGhostFrequent
             , arcFrequentWeight = arcFrequentWeight - w
             }
           Just !(force -> x'@(_, w'))
             -> let (arcFrequent'', arcFrequentWeight'', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent |- w') arcFrequent' (arcFrequentWeight - w) arcGhostFrequent
                 in oldARC
                      { arcFrequent = HashPSQ.insert k now x' arcFrequent''
                      , arcFrequentWeight = arcFrequentWeight'' + w'
                      , arcGhostFrequent = arcGhostFrequent'
                      }
    | Just (_p, x@(_, w), arcRecent') <- HashPSQ.deleteView k arcRecent
      -> f (Just x) <&> \case
           Nothing -> oldARC
             { arcRecent = arcRecent'
             , arcGhostRecent = HashPSQ.insert k now () $ evictGhostToCount arcGhostRecent
             , arcRecentWeight = arcRecentWeight - w
             }
           Just !(force -> x'@(_, w'))
             -> let (arcFrequent', arcFrequentWeight', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent |- w') arcFrequent arcFrequentWeight arcGhostFrequent
                 in oldARC
                      { arcRecent = arcRecent'
                      , arcRecentWeight = arcRecentWeight - w
                      , arcFrequent = HashPSQ.insert k now x' arcFrequent'
                      , arcFrequentWeight = arcFrequentWeight' + w'
                      , arcGhostFrequent = arcGhostFrequent'
                      }
    | Just (_p, (), arcGhostRecent') <- HashPSQ.deleteView k arcGhostRecent
      -> f Nothing <&> \case
           Nothing -> oldARC
             { arcGhostRecent = HashPSQ.insert k now () arcGhostRecent'
             }
           Just !(force -> x@(_, w))
             -> let arcTargetRecent' = min arcMaximumWeight $ arcTargetRecent + max avgWeight (round $ toRational (HashPSQ.size arcGhostFrequent) / toRational (HashPSQ.size arcGhostRecent) * toRational avgWeight)
                    (arcFrequent', arcFrequentWeight', arcGhostFrequent') = evictToSize (arcMaximumWeight |- arcTargetRecent' |- w) arcFrequent arcFrequentWeight arcGhostFrequent
                    (arcRecent', arcRecentWeight', arcGhostRecent'') = evictToSize (max arcTargetRecent' $ arcMaximumWeight |- arcFrequentWeight' |- w) arcRecent arcRecentWeight arcGhostRecent'
                 in oldARC
                      { arcRecent = arcRecent'
                      , arcFrequent = HashPSQ.insert k now x arcFrequent'
                      , arcGhostRecent = arcGhostRecent''
                      , arcGhostFrequent = arcGhostFrequent'
                      , arcRecentWeight = arcRecentWeight'
                      , arcFrequentWeight = arcFrequentWeight' + w
                      , arcTargetRecent = arcTargetRecent'
                      }
    | Just (_p, (), arcGhostFrequent') <- HashPSQ.deleteView k arcGhostFrequent
      -> f Nothing <&> \case
           Nothing -> oldARC
             { arcGhostFrequent = HashPSQ.insert k now () arcGhostFrequent'
             }
           Just !(force -> x@(_, w))
             -> let arcTargetRecent' = arcTargetRecent |- max avgWeight (round $ toRational (HashPSQ.size arcGhostRecent) / toRational (HashPSQ.size arcGhostFrequent) * toRational avgWeight)
                    (arcFrequent', arcFrequentWeight', arcGhostFrequent'') = evictToSize (arcMaximumWeight |- arcTargetRecent' |- w) arcFrequent arcFrequentWeight arcGhostFrequent'
                    (arcRecent', arcRecentWeight', arcGhostRecent') = evictToSize (max arcTargetRecent' $ arcMaximumWeight |- arcFrequentWeight' |- w) arcRecent arcRecentWeight arcGhostRecent
                 in oldARC
                      { arcRecent = arcRecent'
                      , arcFrequent = HashPSQ.insert k now x arcFrequent'
                      , arcGhostRecent = arcGhostRecent'
                      , arcGhostFrequent = arcGhostFrequent''
                      , arcRecentWeight = arcRecentWeight'
                      , arcFrequentWeight = arcFrequentWeight' + w
                      , arcTargetRecent = arcTargetRecent'
                      }
    | otherwise -> f Nothing <&> \case
        Nothing -> oldARC
          { arcGhostRecent = HashPSQ.insert k now () $ evictGhostToCount arcGhostRecent
          }
        Just !(force -> x@(_, w))
          -> let (arcRecent', arcRecentWeight', arcGhostRecent') = evictToSize (max arcTargetRecent (arcMaximumWeight |- arcFrequentWeight) |- w) arcRecent arcRecentWeight arcGhostRecent
              in oldARC
                   { arcRecent = HashPSQ.insert k now x arcRecent'
                   , arcRecentWeight = arcRecentWeight' + w
                   , arcGhostRecent = arcGhostRecent'
                   }
  where
    avgWeight = round $ toRational (arcRecentWeight + arcFrequentWeight) / toRational (HashPSQ.size arcFrequent + HashPSQ.size arcRecent)
    
    later :: ARCTick
    later = over getARCTick succ now
    
    evictToSize :: w -> HashPSQ (Hashed k) ARCTick (v, w) -> w -> HashPSQ (Hashed k) ARCTick () -> (HashPSQ (Hashed k) ARCTick (v, w), w, HashPSQ (Hashed k) ARCTick ())
    evictToSize tSize c cSize ghostC
      | cSize <= tSize = (c, cSize, ghostC)
      | Just (k', p', (_, w'), c') <- HashPSQ.minView c = evictToSize tSize c' (cSize - w') . evictGhostToCount $ HashPSQ.insert k' p' () ghostC
      | otherwise = error "evictToSize: cannot reach required size through eviction"

    evictGhostToCount :: HashPSQ (Hashed k) ARCTick () -> HashPSQ (Hashed k) ARCTick ()
    evictGhostToCount c
      | HashPSQ.size c <= arcMaximumGhost = c
      | Just (_, _, _, c') <- HashPSQ.minView c = evictGhostToCount c'
      | otherwise = error "evictGhostToCount: cannot reach required count through eviction"

lookupARC :: forall k w v.
             ( Ord k, Hashable k
             , Integral w
             , NFData k, NFData w, NFData v
             )
          => k
          -> (ARC k w v, ARCTick)
          -> Maybe (v, w)
lookupARC k = getConst . uncurry (arcAlterF k Const)

insertARC :: forall k w v.
             ( Ord k, Hashable k
             , Integral w
             , NFData k, NFData w, NFData v
             )
          => k
          -> Maybe (v, w)
          -> ARC k w v
          -> ARCTick -> (ARC k w v, ARCTick)
insertARC k newVal = (runIdentity .) . arcAlterF k (const $ pure newVal)


newtype ARCHandle k w v = ARCHandle { _getARCHandle :: IORef (ARC k w v, ARCTick) }
  deriving (Eq)

type DynARCHandle k w = ARCHandle (SomeTypeRep, k) w NFDynamic

initARCHandle :: forall k w v m.
                 ( MonadIO m
                 , Integral w
                 )
              => Int -- ^ @arcMaximumGhost@
              -> w -- ^ @arcMaximumWeight@
              -> m (ARCHandle k w v)
initARCHandle maxGhost maxWeight = fmap ARCHandle . newIORef $ initARC maxGhost maxWeight

cachedARC' :: forall k w v m.
              ( MonadIO m
              , Ord k, Hashable k
              , Integral w
              , NFData k, NFData w, NFData v
              )
           => ARCHandle k w v
           -> k
           -> (Maybe (v, w) -> m (Maybe (v, w)))
           -> m (Maybe v)
cachedARC' (ARCHandle arcVar) k f = do
  oldVal <- lookupARC k <$> readIORef arcVar
  newVal <- f oldVal
  atomicModifyIORef' arcVar $ (, ()) . uncurry (insertARC k newVal)
  -- Using `modifyIORef'` instead of `atomicModifyIORef'` might very
  -- well drop newer values computed during the update.
  --
  -- This was deemed unacceptable due to the risk of cache
  -- invalidations being silently dropped
  --
  -- Another alternative would be to use "optimistic locking",
  -- i.e. read the current value of `arcVar`, compute an updated
  -- version, and write it back atomically iff the `ARCTick` hasn't
  -- changed.
  --
  -- This was not implemented in the hopes that atomicModifyIORef'
  -- already offers sufficient performance.
  --
  -- If optimistic locking is implemented there is a risk of
  -- performance issues due to the overhead and contention likely
  -- associated with the atomic transaction required for the "compare
  -- and swap"
  return $ view _1 <$> newVal

cachedARC :: forall k w v m.
             ( MonadIO m
             , Ord k, Hashable k
             , Integral w
             , NFData k, NFData w, NFData v
             )
          => ARCHandle k w v
          -> k
          -> (Maybe (v, w) -> m (v, w))
          -> m v
cachedARC h k f = fromMaybe (error "cachedARC: cachedARC' returned Nothing") <$> cachedARC' h k (fmap Just . f)
  
lookupARCHandle :: forall k w v m.
                   ( MonadIO m
                   , Ord k, Hashable k
                   , Integral w
                   , NFData k, NFData w, NFData v
                   )
                => ARCHandle k w v
                -> k
                -> m (Maybe (v, w))
lookupARCHandle (ARCHandle arcVar) k = lookupARC k <$> readIORef arcVar


readARCHandle :: MonadIO m
              => ARCHandle k w v
              -> m (ARC k w v, ARCTick)
readARCHandle (ARCHandle arcVar) = readIORef arcVar