module Unison.Runtime.Canonicalizer
  ( Canonicalizer,
    canonicalize,
    categorize,
    unsafeCategorize,
    Canonicity (..),
    CanonMap (..),
    empty,
    lookup,
    insert,
    unsafeLookup,
    findWithDefault,
    fromListByIndex,
    fromList,
    emptyCM,
  )
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 #-}

data CanonMap k v = CanonM
  { forall k v. CanonMap k v -> HashMap (StableName k) v
_fast :: HashMap (StableName k) v,
    forall k v. CanonMap k v -> Map k v
_slow :: M.Map 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)

lookup0 :: (Ord k) => k -> CanonMap k v -> StableName k -> Maybe v
lookup0 :: forall k v. Ord k => k -> CanonMap k v -> StableName k -> Maybe v
lookup0 k
k (CanonM HashMap (StableName k) v
fast Map k v
slow) StableName k
name
  | r :: Maybe v
r@Just {} <- StableName k -> HashMap (StableName k) v -> Maybe v
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup StableName k
name HashMap (StableName k) v
fast = Maybe v
r
  | Bool
otherwise = k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
k Map k v
slow
{-# INLINE lookup0 #-}

lookup :: (Ord k) => k -> CanonMap k v -> IO (Maybe v)
lookup :: forall k v. Ord k => k -> CanonMap k v -> IO (Maybe v)
lookup !k
k CanonMap k v
m = k -> CanonMap k v -> StableName k -> Maybe v
forall k v. Ord k => k -> CanonMap k v -> StableName k -> Maybe v
lookup0 k
k CanonMap 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 #-}

insert0 :: (Ord k) => k -> v -> CanonMap k v -> StableName k -> CanonMap k v
insert0 :: forall k v.
Ord k =>
k -> v -> CanonMap k v -> StableName k -> CanonMap k v
insert0 k
k v
v (CanonM HashMap (StableName k) v
fast Map k v
slow) StableName k
name =
  HashMap (StableName k) v -> Map k v -> CanonMap k v
forall k v. HashMap (StableName k) v -> Map k v -> CanonMap k v
CanonM (StableName k
-> v -> HashMap (StableName k) v -> HashMap (StableName k) v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert StableName k
name v
v HashMap (StableName k) v
fast) (k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k v
v Map k v
slow)
{-# INLINE insert0 #-}

-- Inserts a key-value pair into the map. Note that this operation
-- should _only_ be used in conjunction with a `Canonicalizer` that
-- matches the `CanonMap` on canonical keys, to ensure that any key
-- inserted into the latter is canonical. Otherwise it is possible to
-- create an inconsistent `CanonMap` that will map the same key to
-- different values depending on which copy is given.
insert :: (Ord k) => k -> v -> CanonMap k v -> IO (CanonMap k v)
insert :: forall k v. Ord k => k -> v -> CanonMap k v -> IO (CanonMap k v)
insert !k
k v
v CanonMap k v
m = k -> v -> CanonMap k v -> StableName k -> CanonMap k v
forall k v.
Ord k =>
k -> v -> CanonMap k v -> StableName k -> CanonMap k v
insert0 k
k v
v CanonMap k v
m (StableName k -> CanonMap k v)
-> IO (StableName k) -> IO (CanonMap k 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

findWithDefault0 :: (Ord k) => v -> k -> CanonMap k v -> StableName k -> v
findWithDefault0 :: forall k v. Ord k => v -> k -> CanonMap k v -> StableName k -> v
findWithDefault0 v
df k
k (CanonM HashMap (StableName k) v
fast Map k v
slow) StableName k
name =
  v -> StableName k -> HashMap (StableName k) v -> v
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
HM.findWithDefault (v -> k -> Map k v -> v
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault v
df k
k Map k v
slow) StableName k
name HashMap (StableName k) v
fast
{-# INLINE findWithDefault0 #-}

findWithDefault :: (Ord k) => v -> k -> CanonMap k v -> IO v
findWithDefault :: forall k v. Ord k => v -> k -> CanonMap k v -> IO v
findWithDefault v
df !k
k CanonMap k v
m =
  v -> k -> CanonMap k v -> StableName k -> v
forall k v. Ord k => v -> k -> CanonMap k v -> StableName k -> v
findWithDefault0 v
df k
k CanonMap 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 :: (Ord k) => k -> CanonMap k v -> Maybe v
unsafeLookup :: forall k v. Ord k => 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. Ord k => k -> CanonMap k v -> IO (Maybe v)
lookup k
k CanonMap k v
m
{-# INLINE unsafeLookup #-}

fromListByIndex :: (Ord k) => [k] -> CanonMap k Int
fromListByIndex :: forall k. Ord k => [k] -> CanonMap k Int
fromListByIndex [k]
ks = IO (CanonMap k Int) -> CanonMap k Int
forall a. IO a -> a
unsafePerformIO do
  [StableName k]
ns <- (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]
ks
  pure $ HashMap (StableName k) Int -> Map k Int -> CanonMap k Int
forall k v. HashMap (StableName k) v -> Map k v -> CanonMap k v
CanonM ([(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]
ns [Int
0 ..])) ([(k, Int)] -> Map k Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([k] -> [Int] -> [(k, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [k]
ks [Int
0 ..]))

fromList :: (Ord k) => [(k, v)] -> IO (CanonMap k v)
fromList :: forall k v. Ord k => [(k, v)] -> IO (CanonMap k v)
fromList [(k, v)]
kvs = do
  [(StableName k, v)]
nvs <- ((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 [(k, v)]
kvs
  pure $ HashMap (StableName k) v -> Map k v -> CanonMap k v
forall k v. HashMap (StableName k) v -> Map k v -> CanonMap k v
CanonM ([(StableName k, v)] -> HashMap (StableName k) v
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList [(StableName k, v)]
nvs) ([(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k, v)]
kvs)
  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)

emptyCM :: CanonMap k v
emptyCM :: forall k v. CanonMap k v
emptyCM = HashMap (StableName k) v -> Map k v -> CanonMap k v
forall k v. HashMap (StableName k) v -> Map k v -> CanonMap k v
CanonM HashMap (StableName k) v
forall k v. HashMap k v
HM.empty Map k v
forall k a. Map k a
M.empty