module Unison.Runtime.Canonicalizer
  ( Canonicalizer,
    canonicalize,
    categorize,
    unsafeCategorize,
    Canonicity (..),
    CanonMap (..),
    empty,
    lookup,
    unsafeLookup,
    findWithDefault,
    fromListByIndex,
    fromList,
  )
where

import Control.Exception (evaluate)
import Data.HashMap.Lazy (HashMap)
import Data.HashMap.Lazy qualified as HM
import Data.Map.Lazy qualified as M
import System.IO.Unsafe
import System.Mem.StableName
import Prelude hiding (lookup)

-- A canonicalizer is a structure for mapping values to versions
-- that are unique in memory. This is accomplished via two
-- mappings. One maps stable names to a canonical value, which is
-- fast if we've seen the exact in-memory value before. A second
-- is just a normal hash map, which will canonicalize values by
-- hash code/equality, which allows us to add to the fast lookup.
data Canonicalizer a = Canon
  { forall a. Canonicalizer a -> HashMap (StableName a) a
stableMap :: !(HashMap (StableName a) a),
    forall a. Canonicalizer a -> Map a a
_slowMap :: !(M.Map a a)
  }

empty :: Canonicalizer a
empty :: forall a. Canonicalizer a
empty = HashMap (StableName a) a -> Map a a -> Canonicalizer a
forall a. HashMap (StableName a) a -> Map a a -> Canonicalizer a
Canon HashMap (StableName a) a
forall k v. HashMap k v
HM.empty Map a a
forall k a. Map k a
M.empty

-- Result for categorizing a value with regard to a Canonicalizer
data Canonicity a
  = -- the provided value is the known canonical one
    Canonical
  | -- the provided value is equivalent to this canonical one; updated
    -- canonicalizer
    Equivalent a (Canonicalizer a)
  | -- the provided value was not previously known, and added to the
    -- canonicalizer
    Novel (Canonicalizer a)

categorize0 ::
  (Ord a) =>
  Canonicalizer a ->
  a ->
  StableName a ->
  IO (Canonicity a)
categorize0 :: forall a.
Ord a =>
Canonicalizer a -> a -> StableName a -> IO (Canonicity a)
categorize0 cn :: Canonicalizer a
cn@(Canon HashMap (StableName a) a
fast Map a a
slow) a
x StableName a
xname
  | Just !a
y <- StableName a -> HashMap (StableName a) a -> Maybe a
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup StableName a
xname HashMap (StableName a) a
fast = do
      StableName a
yname <- a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
y
      if StableName a
xname StableName a -> StableName a -> Bool
forall a. Eq a => a -> a -> Bool
== StableName a
yname
        then Canonicity a -> IO (Canonicity a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Canonicity a
forall a. Canonicity a
Canonical
        else Canonicity a -> IO (Canonicity a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Canonicalizer a -> Canonicity a
forall a. a -> Canonicalizer a -> Canonicity a
Equivalent a
y Canonicalizer a
cn)
  | Just a
x <- a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a a
slow = do
      Canonicalizer a
cn <- Canonicalizer a -> IO (Canonicalizer a)
forall a. a -> IO a
evaluate Canonicalizer a
cn {stableMap = HM.insert xname x fast}
      pure (a -> Canonicalizer a -> Canonicity a
forall a. a -> Canonicalizer a -> Canonicity a
Equivalent a
x Canonicalizer a
cn)
  | Bool
otherwise = do
      Canonicalizer a
cn <- Canonicalizer a -> IO (Canonicalizer a)
forall a. a -> IO a
evaluate (HashMap (StableName a) a -> Map a a -> Canonicalizer a
forall a. HashMap (StableName a) a -> Map a a -> Canonicalizer a
Canon (StableName a
-> a -> HashMap (StableName a) a -> HashMap (StableName a) a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert StableName a
xname a
x HashMap (StableName a) a
fast) (a -> a -> Map a a -> Map a a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x a
x Map a a
slow))
      pure (Canonicalizer a -> Canonicity a
forall a. Canonicalizer a -> Canonicity a
Novel Canonicalizer a
cn)
{-# INLINE categorize0 #-}

-- Incorporates a value into a canonicalizer, giving a canonical
-- value, and also indicating whether the value has been seen before,
-- either up to in-memory or Eq equality.
categorize ::
  (Ord a) =>
  Canonicalizer a ->
  a ->
  IO (Canonicity a)
categorize :: forall a. Ord a => Canonicalizer a -> a -> IO (Canonicity a)
categorize Canonicalizer a
cn !a
x = a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
x IO (StableName a)
-> (StableName a -> IO (Canonicity a)) -> IO (Canonicity a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Canonicalizer a -> a -> StableName a -> IO (Canonicity a)
forall a.
Ord a =>
Canonicalizer a -> a -> StableName a -> IO (Canonicity a)
categorize0 Canonicalizer a
cn a
x
{-# INLINE categorize #-}

unsafeCategorize ::
  (Ord a) => Canonicalizer a -> a -> Canonicity a
unsafeCategorize :: forall a. Ord a => Canonicalizer a -> a -> Canonicity a
unsafeCategorize Canonicalizer a
cn a
a = IO (Canonicity a) -> Canonicity a
forall a. IO a -> a
unsafePerformIO (IO (Canonicity a) -> Canonicity a)
-> IO (Canonicity a) -> Canonicity a
forall a b. (a -> b) -> a -> b
$ Canonicalizer a -> a -> IO (Canonicity a)
forall a. Ord a => Canonicalizer a -> a -> IO (Canonicity a)
categorize Canonicalizer a
cn a
a

-- Produces a canonical value and an updated canonicalizer under
-- the assumption that the given stable name is the one for the
-- given value.
canonicalize0 ::
  (Ord a) =>
  Canonicalizer a ->
  a ->
  StableName a ->
  IO (a, Canonicalizer a)
canonicalize0 :: forall a.
Ord a =>
Canonicalizer a -> a -> StableName a -> IO (a, Canonicalizer a)
canonicalize0 cn :: Canonicalizer a
cn@(Canon HashMap (StableName a) a
fast Map a a
slow) a
x StableName a
name
  | Just a
x <- StableName a -> HashMap (StableName a) a -> Maybe a
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup StableName a
name HashMap (StableName a) a
fast = (a, Canonicalizer a) -> IO (a, Canonicalizer a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x, Canonicalizer a
cn)
  | Just a
x <- a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
x Map a a
slow = do
      Canonicalizer a
cn <- Canonicalizer a -> IO (Canonicalizer a)
forall a. a -> IO a
evaluate Canonicalizer a
cn {stableMap = HM.insert name x fast}
      pure (a
x, Canonicalizer a
cn)
  | Bool
otherwise = do
      Canonicalizer a
cn <- Canonicalizer a -> IO (Canonicalizer a)
forall a. a -> IO a
evaluate (HashMap (StableName a) a -> Map a a -> Canonicalizer a
forall a. HashMap (StableName a) a -> Map a a -> Canonicalizer a
Canon (StableName a
-> a -> HashMap (StableName a) a -> HashMap (StableName a) a
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert StableName a
name a
x HashMap (StableName a) a
fast) (a -> a -> Map a a -> Map a a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
x a
x Map a a
slow))
      pure (a
x, Canonicalizer a
cn)
{-# INLINE canonicalize0 #-}

-- Canonicalizes a value, giving an updated canonicalizer.
--
-- Note: uses `unsafePerformIO`. This should be fine, because we
-- are only using stable names to canonicalize values in memory,
-- always replacing values with other values that are identical
-- according to the `Eq` instance.
--
-- The API exposed only allows such canonicalization and building
-- opaque `Canonicalizer` values, so there should be no opportunity
-- for doing anything actually unsafe.
canonicalize :: (Ord a) => Canonicalizer a -> a -> (a, Canonicalizer a)
canonicalize :: forall a. Ord a => Canonicalizer a -> a -> (a, Canonicalizer a)
canonicalize Canonicalizer a
cn !a
x =
  IO (a, Canonicalizer a) -> (a, Canonicalizer a)
forall a. IO a -> a
unsafePerformIO (IO (a, Canonicalizer a) -> (a, Canonicalizer a))
-> IO (a, Canonicalizer a) -> (a, Canonicalizer a)
forall a b. (a -> b) -> a -> b
$ a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName a
x IO (StableName a)
-> (StableName a -> IO (a, Canonicalizer a))
-> IO (a, Canonicalizer a)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Canonicalizer a -> a -> StableName a -> IO (a, Canonicalizer a)
forall a.
Ord a =>
Canonicalizer a -> a -> StableName a -> IO (a, Canonicalizer a)
canonicalize0 Canonicalizer a
cn a
x
{-# INLINEABLE canonicalize #-}

newtype CanonMap k v = CanonM (HashMap (StableName k) v)
  deriving ((forall a b. (a -> b) -> CanonMap k a -> CanonMap k b)
-> (forall a b. a -> CanonMap k b -> CanonMap k a)
-> Functor (CanonMap k)
forall a b. a -> CanonMap k b -> CanonMap k a
forall a b. (a -> b) -> CanonMap k a -> CanonMap k b
forall k a b. a -> CanonMap k b -> CanonMap k a
forall k a b. (a -> b) -> CanonMap k a -> CanonMap k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall k a b. (a -> b) -> CanonMap k a -> CanonMap k b
fmap :: forall a b. (a -> b) -> CanonMap k a -> CanonMap k b
$c<$ :: forall k a b. a -> CanonMap k b -> CanonMap k a
<$ :: forall a b. a -> CanonMap k b -> CanonMap k a
Functor)

lookup :: k -> CanonMap k v -> IO (Maybe v)
lookup :: forall k v. k -> CanonMap k v -> IO (Maybe v)
lookup !k
k (CanonM HashMap (StableName k) v
m) = (StableName k -> HashMap (StableName k) v -> Maybe v)
-> HashMap (StableName k) v -> StableName k -> Maybe v
forall a b c. (a -> b -> c) -> b -> a -> c
flip StableName k -> HashMap (StableName k) v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup HashMap (StableName k) v
m (StableName k -> Maybe v) -> IO (StableName k) -> IO (Maybe v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> IO (StableName k)
forall a. a -> IO (StableName a)
makeStableName k
k
{-# INLINE lookup #-}

findWithDefault :: v -> k -> CanonMap k v -> IO v
findWithDefault :: forall v k. v -> k -> CanonMap k v -> IO v
findWithDefault v
d !k
k (CanonM HashMap (StableName k) v
m) =
  (StableName k -> HashMap (StableName k) v -> v)
-> HashMap (StableName k) v -> StableName k -> v
forall a b c. (a -> b -> c) -> b -> a -> c
flip (v -> StableName k -> HashMap (StableName k) v -> v
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
HM.findWithDefault v
d) HashMap (StableName k) v
m (StableName k -> v) -> IO (StableName k) -> IO v
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> k -> IO (StableName k)
forall a. a -> IO (StableName a)
makeStableName k
k
{-# INLINE findWithDefault #-}

unsafeLookup :: k -> CanonMap k v -> Maybe v
unsafeLookup :: forall k v. k -> CanonMap k v -> Maybe v
unsafeLookup k
k CanonMap k v
m = IO (Maybe v) -> Maybe v
forall a. IO a -> a
unsafePerformIO (IO (Maybe v) -> Maybe v) -> IO (Maybe v) -> Maybe v
forall a b. (a -> b) -> a -> b
$ k -> CanonMap k v -> IO (Maybe v)
forall k v. k -> CanonMap k v -> IO (Maybe v)
lookup k
k CanonMap k v
m
{-# INLINE unsafeLookup #-}

fromListByIndex :: [k] -> CanonMap k Int
fromListByIndex :: forall k. [k] -> CanonMap k Int
fromListByIndex [k]
l = IO (CanonMap k Int) -> CanonMap k Int
forall a. IO a -> a
unsafePerformIO do
  [StableName k]
l <- (k -> IO (StableName k)) -> [k] -> IO [StableName k]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (\k
k -> k -> IO (StableName k)
forall a. a -> IO (StableName a)
makeStableName (k -> IO (StableName k)) -> IO k -> IO (StableName k)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< k -> IO k
forall a. a -> IO a
evaluate k
k) [k]
l
  CanonMap k Int -> IO (CanonMap k Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CanonMap k Int -> IO (CanonMap k Int))
-> (HashMap (StableName k) Int -> CanonMap k Int)
-> HashMap (StableName k) Int
-> IO (CanonMap k Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (StableName k) Int -> CanonMap k Int
forall k v. HashMap (StableName k) v -> CanonMap k v
CanonM (HashMap (StableName k) Int -> IO (CanonMap k Int))
-> HashMap (StableName k) Int -> IO (CanonMap k Int)
forall a b. (a -> b) -> a -> b
$ [(StableName k, Int)] -> HashMap (StableName k) Int
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList ([StableName k] -> [Int] -> [(StableName k, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [StableName k]
l [Int
0 ..])

fromList :: [(k, v)] -> IO (CanonMap k v)
fromList :: forall k v. [(k, v)] -> IO (CanonMap k v)
fromList = ([(StableName k, v)] -> CanonMap k v)
-> IO [(StableName k, v)] -> IO (CanonMap k v)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (HashMap (StableName k) v -> CanonMap k v
forall k v. HashMap (StableName k) v -> CanonMap k v
CanonM (HashMap (StableName k) v -> CanonMap k v)
-> ([(StableName k, v)] -> HashMap (StableName k) v)
-> [(StableName k, v)]
-> CanonMap k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(StableName k, v)] -> HashMap (StableName k) v
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList) (IO [(StableName k, v)] -> IO (CanonMap k v))
-> ([(k, v)] -> IO [(StableName k, v)])
-> [(k, v)]
-> IO (CanonMap k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> IO (StableName k, v))
-> [(k, v)] -> IO [(StableName k, v)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (k, v) -> IO (StableName k, v)
forall {a} {t}. (a, t) -> IO (StableName a, t)
f
  where
    f :: (a, t) -> IO (StableName a, t)
f (a
k, t
v) = (,t
v) (StableName a -> (StableName a, t))
-> IO (StableName a) -> IO (StableName a, t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> IO (StableName a)
forall a. a -> IO (StableName a)
makeStableName (a -> IO (StableName a)) -> IO a -> IO (StableName a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< a -> IO a
forall a. a -> IO a
evaluate a
k)