{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- |
-- Module      :   Grisette.Internal.Core.Data.MemoUtils
-- Copyright   :   (c) Sirui Lu 2021-2023
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.Core.Data.MemoUtils
  ( -- * Memoization
    stableMemo,
    stableMemo2,
    stableMemo3,
    stableMup,
    stableMemoFix,
    weakStableMemo,
    weakStableMemo2,
    weakStableMemo3,
    weakStableMup,
    weakStableMemoFix,
    htmemo,
    htmemo2,
    htmemo3,
    htmemoFix,
    htmup,
  )
where

import Control.Applicative (Const (Const, getConst))
import Control.Monad.Fix (fix)
import Data.Atomics (atomicModifyIORefCAS_)
import qualified Data.HashMap.Strict as HM
import Data.Hashable (Hashable)
import Data.IORef (IORef, newIORef, readIORef)
import Data.Proxy (Proxy (Proxy))
import GHC.Base (Any, Type)
import System.IO.Unsafe (unsafePerformIO)
import System.Mem.StableName (StableName, makeStableName)
import System.Mem.Weak (Weak)
import qualified System.Mem.Weak as Weak
import Unsafe.Coerce (unsafeCoerce)

newtype (f <<< g) a = O {unO :: f (g a)}

-- Invariant: The type parameters for a key and its corresponding
-- value are the same.
type SNMap f g = IORef (HM.HashMap (StableName (f Any)) (g Any))

type MemoTable ref f g = SNMap f (ref <<< g)

class Ref ref where
  mkRef :: a -> b -> IO () -> IO (ref b)
  deRef :: ref a -> IO (Maybe a)
  finalize :: ref a -> IO ()
  tableFinalizer :: MemoTable ref f g -> IO ()
  tableFinalizer t = readIORef t >>= mapM_ (finalize . unO)

instance Ref Weak where
  mkRef x y = Weak.mkWeak x y . Just
  deRef = Weak.deRefWeak
  finalize = Weak.finalize

newtype Strong a = Strong a

instance Ref Strong where
  mkRef _ y _ = return $ Strong y
  deRef (Strong x) = return $ Just x
  finalize (Strong _) = return ()
  tableFinalizer _ = return ()

finalizer :: StableName (f Any) -> Weak (MemoTable ref f g) -> IO ()
finalizer sn weakTbl = do
  r <- Weak.deRefWeak weakTbl
  case r of
    Nothing -> return ()
    Just tbl -> do
      atomicModifyIORefCAS_ tbl $ HM.delete sn

unsafeToAny :: f a -> f Any
unsafeToAny = unsafeCoerce

unsafeFromAny :: f Any -> f a
unsafeFromAny = unsafeCoerce

{-# NOINLINE memo' #-}
memo' ::
  (Ref ref) =>
  Proxy ref ->
  (forall a. f a -> g a) ->
  MemoTable ref f g ->
  Weak (MemoTable ref f g) ->
  f b ->
  g b
memo' _ f tbl weakTbl !x = unsafePerformIO $ do
  sn <- makeStableName $ unsafeToAny x
  lkp <- HM.lookup sn <$> readIORef tbl
  case lkp of
    Nothing -> notFound sn
    Just (O w) -> do
      maybeVal <- deRef w
      case maybeVal of
        Nothing -> notFound sn
        Just val -> do
          return $ unsafeFromAny val
  where
    notFound sn = do
      let !y = f x
      weak <- mkRef x (unsafeToAny y) $ finalizer sn weakTbl
      atomicModifyIORefCAS_ tbl $ HM.insert sn $ O weak
      return y

{-# NOINLINE memo0 #-}
memo0 ::
  (Ref ref) =>
  Proxy (ref :: Type -> Type) ->
  (forall a. f a -> g a) ->
  f b ->
  g b
memo0 p f =
  let (tbl, weak) = unsafePerformIO $ do
        tbl' <- newIORef HM.empty
        weak' <- Weak.mkWeakPtr tbl . Just $ tableFinalizer tbl
        return (tbl', weak')
   in memo' p f tbl weak

-- | Memoize a unary function.
stableMemo :: (a -> b) -> (a -> b)
stableMemo f = getConst . memo0 (Proxy :: Proxy Strong) (Const . f . getConst) . Const

-- | Lift a memoizer to work with one more argument.
stableMup :: (b -> c) -> (a -> b) -> (a -> c)
stableMup mem f = stableMemo (mem . f)

-- | Curried memoization to share partial evaluation
stableMemo2 :: (a -> b -> c) -> (a -> b -> c)
stableMemo2 = stableMup stableMemo

-- | Curried memoization to share partial evaluation
stableMemo3 :: (a -> b -> c -> d) -> (a -> b -> c -> d)
stableMemo3 = stableMup stableMemo2

-- | Memoizing recursion. Use like 'fix'.
stableMemoFix :: ((a -> b) -> (a -> b)) -> a -> b
stableMemoFix h = fix (stableMemo . h)

-- | Memoize a unary function.
weakStableMemo :: (a -> b) -> (a -> b)
weakStableMemo f = getConst . memo0 (Proxy :: Proxy Weak) (Const . f . getConst) . Const

-- | Lift a memoizer to work with one more argument.
weakStableMup :: (b -> c) -> (a -> b) -> (a -> c)
weakStableMup mem f = weakStableMemo (mem . f)

-- | Curried memoization to share partial evaluation
weakStableMemo2 :: (a -> b -> c) -> (a -> b -> c)
weakStableMemo2 = weakStableMup weakStableMemo

-- | Curried memoization to share partial evaluation
weakStableMemo3 :: (a -> b -> c -> d) -> (a -> b -> c -> d)
weakStableMemo3 = weakStableMup weakStableMemo2

-- | Memoizing recursion. Use like 'fix'.
weakStableMemoFix :: ((a -> b) -> (a -> b)) -> a -> b
weakStableMemoFix h = fix (weakStableMemo . h)

-- | Function memoizer with mutable hash table.
{-# NOINLINE htmemo #-}
htmemo :: (Eq k, Hashable k) => (k -> a) -> k -> a
htmemo f = unsafePerformIO $ do
  cache <- newIORef HM.empty
  return $ \(!x) -> unsafePerformIO $ do
    tryV <- HM.lookup x <$> readIORef cache
    case tryV of
      Nothing -> do
        let !v = f x
        atomicModifyIORefCAS_ cache $ \old -> HM.insert x v old
        return v
      Just v -> return v

-- | Lift a memoizer to work with one more argument.
htmup :: (Eq k, Hashable k) => (b -> c) -> (k -> b) -> (k -> c)
htmup mem f = htmemo (mem . f)

-- | Function memoizer with mutable hash table. Works on binary functions.
htmemo2 :: (Eq k1, Hashable k1, Eq k2, Hashable k2) => (k1 -> k2 -> a) -> (k1 -> k2 -> a)
htmemo2 = htmup htmemo

-- | Function memoizer with mutable hash table. Works on ternary functions.
htmemo3 ::
  (Eq k1, Hashable k1, Eq k2, Hashable k2, Eq k3, Hashable k3) =>
  (k1 -> k2 -> k3 -> a) ->
  (k1 -> k2 -> k3 -> a)
htmemo3 = htmup htmemo2

-- | Memoizing recursion. Use like 'fix'.
htmemoFix :: (Eq k, Hashable k) => ((k -> a) -> (k -> a)) -> k -> a
htmemoFix h = fix (htmemo . h)