module Unison.Runtime.Referenced
  ( Referenced (..),
    dereference,
    RTrav,
    Canonize,
    canonicalizeRefs,
    recanonicalizeRefs,
  )
where

import Control.Monad.State.Strict
import Data.Maybe (mapMaybe)
import Unison.Reference
import Unison.Runtime.Canonicalizer

-- A value with optional optimization information for serialization.
-- The references are required for serialization V5, and are assumed
-- to be the only references used in the value _up to in-memory
-- uniqueness_.
--
-- This is parameterized so that it can be used with both Value and
-- Code.
--
-- Also note, the stored referenced might not be 'tight' in the sense
-- that they all actually occur in the value. Maintaining this
-- invariant together with actual canonicalization would be onerous
-- and isn't done at this time.
data Referenced a
  = -- types, terms
    WithRefs [Reference] [Reference] a
  | Plain a
  deriving (Int -> Referenced a -> ShowS
[Referenced a] -> ShowS
Referenced a -> String
(Int -> Referenced a -> ShowS)
-> (Referenced a -> String)
-> ([Referenced a] -> ShowS)
-> Show (Referenced a)
forall a. Show a => Int -> Referenced a -> ShowS
forall a. Show a => [Referenced a] -> ShowS
forall a. Show a => Referenced a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Referenced a -> ShowS
showsPrec :: Int -> Referenced a -> ShowS
$cshow :: forall a. Show a => Referenced a -> String
show :: Referenced a -> String
$cshowList :: forall a. Show a => [Referenced a] -> ShowS
showList :: [Referenced a] -> ShowS
Show, Referenced a -> Referenced a -> Bool
(Referenced a -> Referenced a -> Bool)
-> (Referenced a -> Referenced a -> Bool) -> Eq (Referenced a)
forall a. Eq a => Referenced a -> Referenced a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Referenced a -> Referenced a -> Bool
== :: Referenced a -> Referenced a -> Bool
$c/= :: forall a. Eq a => Referenced a -> Referenced a -> Bool
/= :: Referenced a -> Referenced a -> Bool
Eq)

dereference :: Referenced a -> a
dereference :: forall a. Referenced a -> a
dereference (WithRefs [Reference]
_ [Reference]
_ a
x) = a
x
dereference (Plain a
x) = a
x

type RTrav a =
  forall f.
  (Applicative f) =>
  (Bool -> Reference -> f Reference) ->
  (a -> f a)

type Canonize =
  StateT (Canonicalizer Reference, [Reference], [Reference]) IO

-- Given a reference traversal, canonicalizes the references in a
-- value. The operation is presented as a state transformation, so
-- that it can hook into a larger canonicalization procedure. The
-- lists of canonical references of each sort are yielded as part of
-- the state.
canonicalizeRefs :: RTrav a -> a -> Canonize a
canonicalizeRefs :: forall a. RTrav a -> a -> Canonize a
canonicalizeRefs RTrav a
trav = (Bool
 -> Reference
 -> StateT
      (Canonicalizer Reference, [Reference], [Reference]) IO Reference)
-> a
-> StateT (Canonicalizer Reference, [Reference], [Reference]) IO a
RTrav a
trav Bool
-> Reference
-> StateT
     (Canonicalizer Reference, [Reference], [Reference]) IO Reference
forall {a}.
Ord a =>
Bool -> a -> StateT (Canonicalizer a, [a], [a]) IO a
h
  where
    h :: Bool -> a -> StateT (Canonicalizer a, [a], [a]) IO a
h Bool
isTy a
r = ((Canonicalizer a, [a], [a])
 -> IO (a, (Canonicalizer a, [a], [a])))
-> StateT (Canonicalizer a, [a], [a]) IO a
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT \st :: (Canonicalizer a, [a], [a])
st@(Canonicalizer a
canon, [a]
tys, [a]
tms) ->
      Canonicalizer a -> a -> IO (Canonicity a)
forall a. Ord a => Canonicalizer a -> a -> IO (Canonicity a)
categorize Canonicalizer a
canon a
r IO (Canonicity a)
-> (Canonicity a -> IO (a, (Canonicalizer a, [a], [a])))
-> IO (a, (Canonicalizer a, [a], [a]))
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Canonicity a
Canonical -> (a, (Canonicalizer a, [a], [a]))
-> IO (a, (Canonicalizer a, [a], [a]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
r, (Canonicalizer a, [a], [a])
st)
        Novel Canonicalizer a
canon
          | Bool
isTy -> (a, (Canonicalizer a, [a], [a]))
-> IO (a, (Canonicalizer a, [a], [a]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
r, (Canonicalizer a
canon, a
r a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
tys, [a]
tms))
          | Bool
otherwise -> (a, (Canonicalizer a, [a], [a]))
-> IO (a, (Canonicalizer a, [a], [a]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
r, (Canonicalizer a
canon, [a]
tys, a
r a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
tms))
        Equivalent a
s Canonicalizer a
canon -> (a, (Canonicalizer a, [a], [a]))
-> IO (a, (Canonicalizer a, [a], [a]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
s, (Canonicalizer a
canon, [a]
tys, [a]
tms))
{-# INLINE canonicalizeRefs #-}

-- Given a `Referenced` value, this recanonicalizes the references in
-- the wrapped value. The intention is for this to be hooked into a
-- larger canonicalization procedure, so that already canonicalized
-- values can be more efficiently brought in line with other values
-- that are already canonicalized.
--
-- If the `Referenced` value is `Plain`, then all we can do is
-- traverse it, canonicalizing the references. However, if it is
-- tagged with canonical refs, we can see if they all match existing
-- canonical refs. If so, we don't need to traverse the value. Even if
-- not, we can traverse with marginally more efficient lookups.
recanonicalizeRefs :: RTrav a -> Referenced a -> Canonize a
recanonicalizeRefs :: forall a. RTrav a -> Referenced a -> Canonize a
recanonicalizeRefs RTrav a
trav = \case
  Plain a
v -> RTrav a -> a -> Canonize a
forall a. RTrav a -> a -> Canonize a
canonicalizeRefs (Bool -> Reference -> f Reference) -> a -> f a
RTrav a
trav a
v
  WithRefs [Reference]
tys [Reference]
tms a
v -> do
    [(Reference, Reference)]
typs <- (Maybe (Reference, Reference) -> Maybe (Reference, Reference))
-> [Maybe (Reference, Reference)] -> [(Reference, Reference)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Maybe (Reference, Reference) -> Maybe (Reference, Reference)
forall a. a -> a
id ([Maybe (Reference, Reference)] -> [(Reference, Reference)])
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [Maybe (Reference, Reference)]
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [(Reference, Reference)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Reference
 -> StateT
      (Canonicalizer Reference, [Reference], [Reference])
      IO
      (Maybe (Reference, Reference)))
-> [Reference]
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [Maybe (Reference, Reference)]
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 (Bool
-> Reference
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (Maybe (Reference, Reference))
forall {b}.
Ord b =>
Bool -> b -> StateT (Canonicalizer b, [b], [b]) IO (Maybe (b, b))
g Bool
True) [Reference]
tys
    [(Reference, Reference)]
tmps <- (Maybe (Reference, Reference) -> Maybe (Reference, Reference))
-> [Maybe (Reference, Reference)] -> [(Reference, Reference)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Maybe (Reference, Reference) -> Maybe (Reference, Reference)
forall a. a -> a
id ([Maybe (Reference, Reference)] -> [(Reference, Reference)])
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [Maybe (Reference, Reference)]
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [(Reference, Reference)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Reference
 -> StateT
      (Canonicalizer Reference, [Reference], [Reference])
      IO
      (Maybe (Reference, Reference)))
-> [Reference]
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     [Maybe (Reference, Reference)]
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 (Bool
-> Reference
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (Maybe (Reference, Reference))
forall {b}.
Ord b =>
Bool -> b -> StateT (Canonicalizer b, [b], [b]) IO (Maybe (b, b))
g Bool
False) [Reference]
tms

    CanonMap Reference Reference
ctys <- IO (CanonMap Reference Reference)
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (CanonMap Reference Reference)
forall (m :: * -> *) a.
Monad m =>
m a
-> StateT (Canonicalizer Reference, [Reference], [Reference]) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (CanonMap Reference Reference)
 -> StateT
      (Canonicalizer Reference, [Reference], [Reference])
      IO
      (CanonMap Reference Reference))
-> IO (CanonMap Reference Reference)
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (CanonMap Reference Reference)
forall a b. (a -> b) -> a -> b
$ [(Reference, Reference)] -> IO (CanonMap Reference Reference)
forall k v. [(k, v)] -> IO (CanonMap k v)
fromList [(Reference, Reference)]
typs
    CanonMap Reference Reference
ctms <- IO (CanonMap Reference Reference)
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (CanonMap Reference Reference)
forall (m :: * -> *) a.
Monad m =>
m a
-> StateT (Canonicalizer Reference, [Reference], [Reference]) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (CanonMap Reference Reference)
 -> StateT
      (Canonicalizer Reference, [Reference], [Reference])
      IO
      (CanonMap Reference Reference))
-> IO (CanonMap Reference Reference)
-> StateT
     (Canonicalizer Reference, [Reference], [Reference])
     IO
     (CanonMap Reference Reference)
forall a b. (a -> b) -> a -> b
$ [(Reference, Reference)] -> IO (CanonMap Reference Reference)
forall k v. [(k, v)] -> IO (CanonMap k v)
fromList [(Reference, Reference)]
tmps

    let f :: Bool -> Reference -> IO Reference
f Bool
False Reference
r = Reference
-> Reference -> CanonMap Reference Reference -> IO Reference
forall v k. v -> k -> CanonMap k v -> IO v
findWithDefault Reference
r Reference
r CanonMap Reference Reference
ctms
        f Bool
True Reference
r = Reference
-> Reference -> CanonMap Reference Reference -> IO Reference
forall v k. v -> k -> CanonMap k v -> IO v
findWithDefault Reference
r Reference
r CanonMap Reference Reference
ctys

    if [(Reference, Reference)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Reference, Reference)]
typs Bool -> Bool -> Bool
&& [(Reference, Reference)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Reference, Reference)]
tmps
      then a -> Canonize a
forall a.
a
-> StateT (Canonicalizer Reference, [Reference], [Reference]) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
v -- already canonical
      else IO a -> Canonize a
forall (m :: * -> *) a.
Monad m =>
m a
-> StateT (Canonicalizer Reference, [Reference], [Reference]) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO a -> Canonize a) -> IO a -> Canonize a
forall a b. (a -> b) -> a -> b
$ (Bool -> Reference -> IO Reference) -> a -> IO a
RTrav a
trav Bool -> Reference -> IO Reference
f a
v
  where
    g :: Bool -> b -> StateT (Canonicalizer b, [b], [b]) IO (Maybe (b, b))
g Bool
isTy b
r = ((Canonicalizer b, [b], [b])
 -> IO (Maybe (b, b), (Canonicalizer b, [b], [b])))
-> StateT (Canonicalizer b, [b], [b]) IO (Maybe (b, b))
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT \st :: (Canonicalizer b, [b], [b])
st@(Canonicalizer b
canon, [b]
tys, [b]
tms) ->
      Canonicalizer b -> b -> IO (Canonicity b)
forall a. Ord a => Canonicalizer a -> a -> IO (Canonicity a)
categorize Canonicalizer b
canon b
r IO (Canonicity b)
-> (Canonicity b -> IO (Maybe (b, b), (Canonicalizer b, [b], [b])))
-> IO (Maybe (b, b), (Canonicalizer b, [b], [b]))
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Canonicity b
Canonical -> (Maybe (b, b), (Canonicalizer b, [b], [b]))
-> IO (Maybe (b, b), (Canonicalizer b, [b], [b]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (b, b)
forall a. Maybe a
Nothing, (Canonicalizer b, [b], [b])
st)
        Novel Canonicalizer b
canon ->
          (Maybe (b, b), (Canonicalizer b, [b], [b]))
-> IO (Maybe (b, b), (Canonicalizer b, [b], [b]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( Maybe (b, b)
forall a. Maybe a
Nothing,
              ( Canonicalizer b
canon,
                if Bool
isTy then b
r b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
tys else [b]
tys,
                if Bool
isTy then [b]
tms else b
r b -> [b] -> [b]
forall a. a -> [a] -> [a]
: [b]
tms
              )
            )
        Equivalent b
s Canonicalizer b
canon -> (Maybe (b, b), (Canonicalizer b, [b], [b]))
-> IO (Maybe (b, b), (Canonicalizer b, [b], [b]))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((b, b) -> Maybe (b, b)
forall a. a -> Maybe a
Just (b
r, b
s), (Canonicalizer b
canon, [b]
tys, [b]
tms))
{-# INLINE recanonicalizeRefs #-}