module Unison.Codebase.CodeLookup where

import Control.Monad.Morph (MFunctor (..))
import Data.Set qualified as Set
import Unison.DataDeclaration (Decl)
import Unison.DataDeclaration qualified as DD
import Unison.Prelude
import Unison.Reference qualified as Reference
import Unison.Term (Term)
import Unison.Term qualified as Term
import Unison.Type (Type)
import Unison.Util.Defns (Defns (..))
import Unison.Util.Set qualified as Set
import Unison.Var (Var)

data CodeLookup v m a = CodeLookup
  { forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Term v a))
getTerm :: Reference.Id -> m (Maybe (Term v a)),
    forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Type v a))
getTypeOfTerm :: Reference.Id -> m (Maybe (Type v a)),
    forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Decl v a))
getTypeDeclaration :: Reference.Id -> m (Maybe (Decl v a))
  }

instance MFunctor (CodeLookup v) where
  hoist :: forall (m :: * -> *) (n :: * -> *) b.
Monad m =>
(forall a. m a -> n a) -> CodeLookup v m b -> CodeLookup v n b
hoist forall a. m a -> n a
f (CodeLookup Id -> m (Maybe (Term v b))
tm Id -> m (Maybe (Type v b))
tmTyp Id -> m (Maybe (Decl v b))
tp) = (Id -> n (Maybe (Term v b)))
-> (Id -> n (Maybe (Type v b)))
-> (Id -> n (Maybe (Decl v b)))
-> CodeLookup v n b
forall v (m :: * -> *) a.
(Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
CodeLookup (m (Maybe (Term v b)) -> n (Maybe (Term v b))
forall a. m a -> n a
f (m (Maybe (Term v b)) -> n (Maybe (Term v b)))
-> (Id -> m (Maybe (Term v b))) -> Id -> n (Maybe (Term v b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> m (Maybe (Term v b))
tm) (m (Maybe (Type v b)) -> n (Maybe (Type v b))
forall a. m a -> n a
f (m (Maybe (Type v b)) -> n (Maybe (Type v b)))
-> (Id -> m (Maybe (Type v b))) -> Id -> n (Maybe (Type v b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> m (Maybe (Type v b))
tmTyp) (m (Maybe (Decl v b)) -> n (Maybe (Decl v b))
forall a. m a -> n a
f (m (Maybe (Decl v b)) -> n (Maybe (Decl v b)))
-> (Id -> m (Maybe (Decl v b))) -> Id -> n (Maybe (Decl v b))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> m (Maybe (Decl v b))
tp)

instance (Ord v, Functor m) => Functor (CodeLookup v m) where
  fmap :: forall a b. (a -> b) -> CodeLookup v m a -> CodeLookup v m b
fmap a -> b
f CodeLookup v m a
cl = (Id -> m (Maybe (Term v b)))
-> (Id -> m (Maybe (Type v b)))
-> (Id -> m (Maybe (Decl v b)))
-> CodeLookup v m b
forall v (m :: * -> *) a.
(Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
CodeLookup Id -> m (Maybe (Term v b))
tm Id -> m (Maybe (Type v b))
tmTyp Id -> m (Maybe (Decl v b))
ty
    where
      tm :: Id -> m (Maybe (Term v b))
tm Id
id = (Term v a -> Term v b) -> Maybe (Term v a) -> Maybe (Term v b)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> b) -> Term v a -> Term v b
forall v a a2. Ord v => (a -> a2) -> Term v a -> Term v a2
Term.amap a -> b
f) (Maybe (Term v a) -> Maybe (Term v b))
-> m (Maybe (Term v a)) -> m (Maybe (Term v b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CodeLookup v m a -> Id -> m (Maybe (Term v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Term v a))
getTerm CodeLookup v m a
cl Id
id
      ty :: Id -> m (Maybe (Decl v b))
ty Id
id = (Either (EffectDeclaration v a) (DataDeclaration v a) -> Decl v b)
-> Maybe (Either (EffectDeclaration v a) (DataDeclaration v a))
-> Maybe (Decl v b)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Either (EffectDeclaration v a) (DataDeclaration v a) -> Decl v b
md (Maybe (Either (EffectDeclaration v a) (DataDeclaration v a))
 -> Maybe (Decl v b))
-> m (Maybe (Either (EffectDeclaration v a) (DataDeclaration v a)))
-> m (Maybe (Decl v b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CodeLookup v m a
-> Id
-> m (Maybe (Either (EffectDeclaration v a) (DataDeclaration v a)))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Decl v a))
getTypeDeclaration CodeLookup v m a
cl Id
id
      tmTyp :: Id -> m (Maybe (Type v b))
tmTyp Id
id = ((Term F v a -> Type v b) -> Maybe (Term F v a) -> Maybe (Type v b)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Term F v a -> Type v b)
 -> Maybe (Term F v a) -> Maybe (Type v b))
-> ((a -> b) -> Term F v a -> Type v b)
-> (a -> b)
-> Maybe (Term F v a)
-> Maybe (Type v b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Term F v a -> Type v b
forall a b. (a -> b) -> Term F v a -> Term F v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) a -> b
f (Maybe (Term F v a) -> Maybe (Type v b))
-> m (Maybe (Term F v a)) -> m (Maybe (Type v b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CodeLookup v m a -> Id -> m (Maybe (Term F v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Type v a))
getTypeOfTerm CodeLookup v m a
cl Id
id
      md :: Either (EffectDeclaration v a) (DataDeclaration v a) -> Decl v b
md (Left EffectDeclaration v a
e) = EffectDeclaration v b -> Decl v b
forall a b. a -> Either a b
Left (a -> b
f (a -> b) -> EffectDeclaration v a -> EffectDeclaration v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> EffectDeclaration v a
e)
      md (Right DataDeclaration v a
d) = DataDeclaration v b -> Decl v b
forall a b. b -> Either a b
Right (a -> b
f (a -> b) -> DataDeclaration v a -> DataDeclaration v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> DataDeclaration v a
d)

instance (Monad m) => Semigroup (CodeLookup v m a) where
  CodeLookup v m a
c1 <> :: CodeLookup v m a -> CodeLookup v m a -> CodeLookup v m a
<> CodeLookup v m a
c2 = (Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
forall v (m :: * -> *) a.
(Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
CodeLookup Id -> m (Maybe (Term v a))
tm Id -> m (Maybe (Type v a))
tmTyp Id -> m (Maybe (Decl v a))
ty
    where
      tm :: Id -> m (Maybe (Term v a))
tm Id
id = do
        Maybe (Term v a)
o <- CodeLookup v m a -> Id -> m (Maybe (Term v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Term v a))
getTerm CodeLookup v m a
c1 Id
id
        case Maybe (Term v a)
o of Maybe (Term v a)
Nothing -> CodeLookup v m a -> Id -> m (Maybe (Term v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Term v a))
getTerm CodeLookup v m a
c2 Id
id; Just Term v a
_ -> Maybe (Term v a) -> m (Maybe (Term v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Term v a)
o
      tmTyp :: Id -> m (Maybe (Type v a))
tmTyp Id
id = do
        Maybe (Type v a)
o <- CodeLookup v m a -> Id -> m (Maybe (Type v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Type v a))
getTypeOfTerm CodeLookup v m a
c1 Id
id
        case Maybe (Type v a)
o of Maybe (Type v a)
Nothing -> CodeLookup v m a -> Id -> m (Maybe (Type v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Type v a))
getTypeOfTerm CodeLookup v m a
c2 Id
id; Just Type v a
_ -> Maybe (Type v a) -> m (Maybe (Type v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Type v a)
o
      ty :: Id -> m (Maybe (Decl v a))
ty Id
id = do
        Maybe (Decl v a)
o <- CodeLookup v m a -> Id -> m (Maybe (Decl v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Decl v a))
getTypeDeclaration CodeLookup v m a
c1 Id
id
        case Maybe (Decl v a)
o of Maybe (Decl v a)
Nothing -> CodeLookup v m a -> Id -> m (Maybe (Decl v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Decl v a))
getTypeDeclaration CodeLookup v m a
c2 Id
id; Just Decl v a
_ -> Maybe (Decl v a) -> m (Maybe (Decl v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Decl v a)
o

instance (Monad m) => Monoid (CodeLookup v m a) where
  mempty :: CodeLookup v m a
mempty = (Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
forall v (m :: * -> *) a.
(Id -> m (Maybe (Term v a)))
-> (Id -> m (Maybe (Type v a)))
-> (Id -> m (Maybe (Decl v a)))
-> CodeLookup v m a
CodeLookup (m (Maybe (Term v a)) -> Id -> m (Maybe (Term v a))
forall a b. a -> b -> a
const (m (Maybe (Term v a)) -> Id -> m (Maybe (Term v a)))
-> m (Maybe (Term v a)) -> Id -> m (Maybe (Term v a))
forall a b. (a -> b) -> a -> b
$ Maybe (Term v a) -> m (Maybe (Term v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Term v a)
forall a. Maybe a
Nothing) (m (Maybe (Type v a)) -> Id -> m (Maybe (Type v a))
forall a b. a -> b -> a
const (m (Maybe (Type v a)) -> Id -> m (Maybe (Type v a)))
-> m (Maybe (Type v a)) -> Id -> m (Maybe (Type v a))
forall a b. (a -> b) -> a -> b
$ Maybe (Type v a) -> m (Maybe (Type v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Type v a)
forall a. Maybe a
Nothing) (m (Maybe (Decl v a)) -> Id -> m (Maybe (Decl v a))
forall a b. a -> b -> a
const (m (Maybe (Decl v a)) -> Id -> m (Maybe (Decl v a)))
-> m (Maybe (Decl v a)) -> Id -> m (Maybe (Decl v a))
forall a b. (a -> b) -> a -> b
$ Maybe (Decl v a) -> m (Maybe (Decl v a))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Decl v a)
forall a. Maybe a
Nothing)

-- todo: can this be implemented in terms of TransitiveClosure.transitiveClosure?
-- todo: add some tests on this guy?
transitiveDependencies ::
  (Monad m, Var v) =>
  CodeLookup v m a ->
  Set Reference.Id ->
  Reference.Id ->
  m (Set Reference.Id)
transitiveDependencies :: forall (m :: * -> *) v a.
(Monad m, Var v) =>
CodeLookup v m a -> Set Id -> Id -> m (Set Id)
transitiveDependencies CodeLookup v m a
code Set Id
seen0 Id
rid =
  if Id -> Set Id -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Id
rid Set Id
seen0
    then Set Id -> m (Set Id)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Id
seen0
    else
      let seen :: Set Id
seen = Id -> Set Id -> Set Id
forall a. Ord a => a -> Set a -> Set a
Set.insert Id
rid Set Id
seen0
          getIds :: Set Reference -> Set Id
getIds = (Reference -> Maybe Id) -> Set Reference -> Set Id
forall b a. Ord b => (a -> Maybe b) -> Set a -> Set b
Set.mapMaybe Reference -> Maybe Id
Reference.toId
       in CodeLookup v m a -> Id -> m (Maybe (Term v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Term v a))
getTerm CodeLookup v m a
code Id
rid m (Maybe (Term v a))
-> (Maybe (Term v a) -> m (Set Id)) -> m (Set Id)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Just Term v a
t ->
              (Set Id -> Id -> m (Set Id)) -> Set Id -> Set Id -> m (Set Id)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (CodeLookup v m a -> Set Id -> Id -> m (Set Id)
forall (m :: * -> *) v a.
(Monad m, Var v) =>
CodeLookup v m a -> Set Id -> Id -> m (Set Id)
transitiveDependencies CodeLookup v m a
code) Set Id
seen (Set Reference -> Set Id
getIds (Set Reference -> Set Id) -> Set Reference -> Set Id
forall a b. (a -> b) -> a -> b
$ let deps :: DefnsF Set Reference Reference
deps = Term v a -> DefnsF Set Reference Reference
forall v vt at ap a.
(Ord v, Ord vt) =>
Term2 vt at ap v a -> DefnsF Set Reference Reference
Term.dependencies Term v a
t in DefnsF Set Reference Reference
deps.terms Set Reference -> Set Reference -> Set Reference
forall a. Semigroup a => a -> a -> a
<> DefnsF Set Reference Reference
deps.types)
            Maybe (Term v a)
Nothing ->
              CodeLookup v m a -> Id -> m (Maybe (Decl v a))
forall v (m :: * -> *) a.
CodeLookup v m a -> Id -> m (Maybe (Decl v a))
getTypeDeclaration CodeLookup v m a
code Id
rid m (Maybe (Decl v a))
-> (Maybe (Decl v a) -> m (Set Id)) -> m (Set Id)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Maybe (Decl v a)
Nothing -> Set Id -> m (Set Id)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set Id
seen
                Just (Left EffectDeclaration v a
ed) ->
                  (Set Id -> Id -> m (Set Id)) -> Set Id -> Set Id -> m (Set Id)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                    (CodeLookup v m a -> Set Id -> Id -> m (Set Id)
forall (m :: * -> *) v a.
(Monad m, Var v) =>
CodeLookup v m a -> Set Id -> Id -> m (Set Id)
transitiveDependencies CodeLookup v m a
code)
                    Set Id
seen
                    (Set Reference -> Set Id
getIds (Set Reference -> Set Id) -> Set Reference -> Set Id
forall a b. (a -> b) -> a -> b
$ DataDeclaration v a -> Set Reference
forall v a. Ord v => DataDeclaration v a -> Set Reference
DD.typeDependencies (EffectDeclaration v a -> DataDeclaration v a
forall v a. EffectDeclaration v a -> DataDeclaration v a
DD.toDataDecl EffectDeclaration v a
ed))
                Just (Right DataDeclaration v a
dd) ->
                  (Set Id -> Id -> m (Set Id)) -> Set Id -> Set Id -> m (Set Id)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
                    (CodeLookup v m a -> Set Id -> Id -> m (Set Id)
forall (m :: * -> *) v a.
(Monad m, Var v) =>
CodeLookup v m a -> Set Id -> Id -> m (Set Id)
transitiveDependencies CodeLookup v m a
code)
                    Set Id
seen
                    (Set Reference -> Set Id
getIds (Set Reference -> Set Id) -> Set Reference -> Set Id
forall a b. (a -> b) -> a -> b
$ DataDeclaration v a -> Set Reference
forall v a. Ord v => DataDeclaration v a -> Set Reference
DD.typeDependencies DataDeclaration v a
dd)