{-# LANGUAGE PartialTypeSignatures #-}

-- | Find a computation of type '{IO} () in the codebase.
module Unison.Codebase.MainTerm
  ( MainTerm (..),
    getMainTerm,
    builtinIOTestTypes,
    builtinMain,
    builtinMainWithResultType,
  )
where

import Control.Lens (mapped, _1)
import Data.List.NonEmpty qualified as NEList
import Data.Set.NonEmpty (NESet)
import Data.Set.NonEmpty qualified as NESet
import Unison.Builtin.Decls qualified as DD
import Unison.HashQualified qualified as HQ
import Unison.Name (Name)
import Unison.Name qualified as Name
import Unison.Names qualified as Names
import Unison.Parser.Ann (Ann)
import Unison.Parser.Ann qualified as Parser.Ann
import Unison.Prelude
import Unison.Reference (Reference, TermReference)
import Unison.Referent (Referent)
import Unison.Referent qualified as Referent
import Unison.Term (Term)
import Unison.Term qualified as Term
import Unison.Type (Type)
import Unison.Type qualified as Type
import Unison.Typechecker qualified as Typechecker
import Unison.Util.Relation qualified as Relation
import Unison.Var (Var)
import Unison.Var qualified as Var

data MainTerm v
  = -- No terms were found with this name
    NotFound
  | -- 1 or more terms were found with this name, but none of them have the right type
    -- Invariant: list is not empty
    BadType [(HQ.HashQualified Name, TermReference, Type v Ann)]
  | -- 2 or more terms were found with this name, and of those, 2 or more have the right type
    -- Invariant: length of list is >= 2
    Ambiguous [(HQ.HashQualified Name, TermReference, Type v Ann)]
  | -- 1 or more terms were found with this name, and exactly 1 has the right type
    Success (HQ.HashQualified Name) TermReference (Term v Ann) (Type v Ann)

getMainTerm ::
  (Monad m, Var v) =>
  (Reference -> m (Maybe (Type v Ann))) ->
  Names.Names ->
  HQ.HashQualified Name ->
  Type.Type v Ann ->
  m (MainTerm v)
getMainTerm :: forall (m :: * -> *) v.
(Monad m, Var v) =>
(Reference -> m (Maybe (Type v Ann)))
-> Names -> HashQualified Name -> Type v Ann -> m (MainTerm v)
getMainTerm Reference -> m (Maybe (Type v Ann))
loadTypeOfTerm Names
parseNames HashQualified Name
mainName Type v Ann
mainType = do
  -- Get all terms and constructors referred to by that name
  let allReferents :: [(Name, Referent)]
      allReferents :: [(Name, Referent)]
allReferents =
        Relation Name Referent -> [(Name, Referent)]
forall a b. Relation a b -> [(a, b)]
Relation.toList (Relation Name Referent -> [(Name, Referent)])
-> Relation Name Referent -> [(Name, Referent)]
forall a b. (a -> b) -> a -> b
$
          Relation Name Referent -> Relation Name Referent
forall r. Ord r => Relation Name r -> Relation Name r
Name.keepHighestPriority (Relation Name Referent -> Relation Name Referent)
-> Relation Name Referent -> Relation Name Referent
forall a b. (a -> b) -> a -> b
$
            (Referent -> ShortHash)
-> HashQualified Name
-> Relation Name Referent
-> Relation Name Referent
forall ref.
Ord ref =>
(ref -> ShortHash)
-> HashQualified Name -> Relation Name ref -> Relation Name ref
HQ.filterBySuffix
              Referent -> ShortHash
Referent.toShortHash
              HashQualified Name
mainName
              (Names -> Relation Name Referent
Names.terms Names
parseNames)

  -- Keep only the terms (throwing away constructors)
  [(Name, Reference, Type v Ann)]
allTermReferences :: [(Name, TermReference, Type v Ann)] <-
    [(Name, Referent)]
allReferents [(Name, Referent)]
-> ([(Name, Referent)] -> m [(Name, Reference, Type v Ann)])
-> m [(Name, Reference, Type v Ann)]
forall a b. a -> (a -> b) -> b
& ((Name, Referent) -> m (Maybe (Name, Reference, Type v Ann)))
-> [(Name, Referent)] -> m [(Name, Reference, Type v Ann)]
forall (m :: * -> *) a b.
Monad m =>
(a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM \case
      (Name
name, Referent.Ref Reference
ref) -> do
        Reference -> m (Maybe (Type v Ann))
loadTypeOfTerm Reference
ref m (Maybe (Type v Ann))
-> (Maybe (Type v Ann) -> Maybe (Name, Reference, Type v Ann))
-> m (Maybe (Name, Reference, Type v Ann))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
          Just Type v Ann
ty -> (Name, Reference, Type v Ann)
-> Maybe (Name, Reference, Type v Ann)
forall a. a -> Maybe a
Just (Name
name, Reference
ref, Type v Ann
ty)
          -- this shouldn't really happen
          Maybe (Type v Ann)
Nothing -> Maybe (Name, Reference, Type v Ann)
forall a. Maybe a
Nothing
      (Name, Referent)
_ -> Maybe (Name, Reference, Type v Ann)
-> m (Maybe (Name, Reference, Type v Ann))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Name, Reference, Type v Ann)
forall a. Maybe a
Nothing

  -- Keep only the terms that are of the right 'main' type
  let allTermReferencesThatCouldBeRun :: [(Name, TermReference, Type v Ann)]
      allTermReferencesThatCouldBeRun :: [(Name, Reference, Type v Ann)]
allTermReferencesThatCouldBeRun =
        ((Name, Reference, Type v Ann) -> Bool)
-> [(Name, Reference, Type v Ann)]
-> [(Name, Reference, Type v Ann)]
forall a. (a -> Bool) -> [a] -> [a]
filter
          (\(Name
_, Reference
_, Type v Ann
ty) -> Type v Ann -> Type v Ann -> Bool
forall v loc. Var v => Type v loc -> Type v loc -> Bool
Typechecker.fitsScheme Type v Ann
ty Type v Ann
mainType)
          [(Name, Reference, Type v Ann)]
allTermReferences

  MainTerm v -> m (MainTerm v)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure case [(Name, Reference, Type v Ann)]
allTermReferencesThatCouldBeRun of
    [(Name
name, Reference
ref, Type v Ann
ty)] ->
      let a :: Ann
a = Ann
Parser.Ann.External
          tm :: Term v Ann
tm = Ann -> Ann -> Term v Ann -> Term v Ann
forall v a. Var v => a -> a -> Term v a -> Term v a
DD.forceTerm Ann
a Ann
a (Ann -> Reference -> Term v Ann
forall v a vt at ap. Ord v => a -> Reference -> Term2 vt at ap v a
Term.ref Ann
a Reference
ref)
       in HashQualified Name
-> Reference -> Term v Ann -> Type v Ann -> MainTerm v
forall v.
HashQualified Name
-> Reference -> Term v Ann -> Type v Ann -> MainTerm v
Success (HashQualified Name
mainName HashQualified Name -> Name -> HashQualified Name
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Name
name) Reference
ref Term v Ann
tm Type v Ann
ty
    [] ->
      case [(Name, Reference, Type v Ann)]
allTermReferences of
        [] -> MainTerm v
forall v. MainTerm v
NotFound
        [(Name, Reference, Type v Ann)]
_ -> [(HashQualified Name, Reference, Type v Ann)] -> MainTerm v
forall v.
[(HashQualified Name, Reference, Type v Ann)] -> MainTerm v
BadType (ASetter
  [(Name, Reference, Type v Ann)]
  [(HashQualified Name, Reference, Type v Ann)]
  Name
  (HashQualified Name)
-> (Name -> HashQualified Name)
-> [(Name, Reference, Type v Ann)]
-> [(HashQualified Name, Reference, Type v Ann)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (((Name, Reference, Type v Ann)
 -> Identity (HashQualified Name, Reference, Type v Ann))
-> [(Name, Reference, Type v Ann)]
-> Identity [(HashQualified Name, Reference, Type v Ann)]
Setter
  [(Name, Reference, Type v Ann)]
  [(HashQualified Name, Reference, Type v Ann)]
  (Name, Reference, Type v Ann)
  (HashQualified Name, Reference, Type v Ann)
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped (((Name, Reference, Type v Ann)
  -> Identity (HashQualified Name, Reference, Type v Ann))
 -> [(Name, Reference, Type v Ann)]
 -> Identity [(HashQualified Name, Reference, Type v Ann)])
-> ((Name -> Identity (HashQualified Name))
    -> (Name, Reference, Type v Ann)
    -> Identity (HashQualified Name, Reference, Type v Ann))
-> ASetter
     [(Name, Reference, Type v Ann)]
     [(HashQualified Name, Reference, Type v Ann)]
     Name
     (HashQualified Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name -> Identity (HashQualified Name))
-> (Name, Reference, Type v Ann)
-> Identity (HashQualified Name, Reference, Type v Ann)
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (Name, Reference, Type v Ann)
  (HashQualified Name, Reference, Type v Ann)
  Name
  (HashQualified Name)
_1) (HashQualified Name
mainName HashQualified Name -> Name -> HashQualified Name
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$>) [(Name, Reference, Type v Ann)]
allTermReferences)
    [(Name, Reference, Type v Ann)]
_ -> [(HashQualified Name, Reference, Type v Ann)] -> MainTerm v
forall v.
[(HashQualified Name, Reference, Type v Ann)] -> MainTerm v
Ambiguous (ASetter
  [(Name, Reference, Type v Ann)]
  [(HashQualified Name, Reference, Type v Ann)]
  Name
  (HashQualified Name)
-> (Name -> HashQualified Name)
-> [(Name, Reference, Type v Ann)]
-> [(HashQualified Name, Reference, Type v Ann)]
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over (((Name, Reference, Type v Ann)
 -> Identity (HashQualified Name, Reference, Type v Ann))
-> [(Name, Reference, Type v Ann)]
-> Identity [(HashQualified Name, Reference, Type v Ann)]
Setter
  [(Name, Reference, Type v Ann)]
  [(HashQualified Name, Reference, Type v Ann)]
  (Name, Reference, Type v Ann)
  (HashQualified Name, Reference, Type v Ann)
forall (f :: * -> *) a b. Functor f => Setter (f a) (f b) a b
mapped (((Name, Reference, Type v Ann)
  -> Identity (HashQualified Name, Reference, Type v Ann))
 -> [(Name, Reference, Type v Ann)]
 -> Identity [(HashQualified Name, Reference, Type v Ann)])
-> ((Name -> Identity (HashQualified Name))
    -> (Name, Reference, Type v Ann)
    -> Identity (HashQualified Name, Reference, Type v Ann))
-> ASetter
     [(Name, Reference, Type v Ann)]
     [(HashQualified Name, Reference, Type v Ann)]
     Name
     (HashQualified Name)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name -> Identity (HashQualified Name))
-> (Name, Reference, Type v Ann)
-> Identity (HashQualified Name, Reference, Type v Ann)
forall s t a b. Field1 s t a b => Lens s t a b
Lens
  (Name, Reference, Type v Ann)
  (HashQualified Name, Reference, Type v Ann)
  Name
  (HashQualified Name)
_1) (HashQualified Name
mainName HashQualified Name -> Name -> HashQualified Name
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$>) [(Name, Reference, Type v Ann)]
allTermReferencesThatCouldBeRun)

-- forall x. '{ io2.IO, Exception } x
builtinMain :: (Var v) => a -> Type.Type v a
builtinMain :: forall v a. Var v => a -> Type v a
builtinMain a
a =
  let result :: v
result = Text -> v
forall v. Var v => Text -> v
Var.named Text
"result"
   in a -> v -> Type v a -> Type v a
forall v a. Ord v => a -> v -> Type v a -> Type v a
Type.forAll a
a v
result (a -> Type v a -> Type v a
forall v a. Var v => a -> Type v a -> Type v a
builtinMainWithResultType a
a (a -> v -> Type v a
forall v a. Ord v => a -> v -> Type v a
Type.var a
a v
result))

-- '{io2.IO, Exception} res
builtinMainWithResultType :: (Var v) => a -> Type.Type v a -> Type.Type v a
builtinMainWithResultType :: forall v a. Var v => a -> Type v a -> Type v a
builtinMainWithResultType a
a Type v a
res = a -> Type v a -> Type v a -> Type v a
forall v a. Ord v => a -> Type v a -> Type v a -> Type v a
Type.arrow a
a (a -> Reference -> Type v a
forall v a. Ord v => a -> Reference -> Type v a
Type.ref a
a Reference
DD.unitRef) Type v a
io
  where
    io :: Type v a
io = a -> [Type v a] -> Type v a -> Type v a
forall v a. Ord v => a -> [Type v a] -> Type v a -> Type v a
Type.effect a
a [a -> Type v a
forall v a. Ord v => a -> Type v a
Type.builtinIO a
a, a -> Type v a
forall v a. Ord v => a -> Type v a
DD.exceptionType a
a] Type v a
res

-- | All possible IO'ish test types, e.g.
-- '{IO, Exception} [Result]
-- '{IO} [Result]
builtinIOTestTypes :: forall v a. (Ord v, Var v) => a -> NESet (Type.Type v a)
builtinIOTestTypes :: forall v a. (Ord v, Var v) => a -> NESet (Type v a)
builtinIOTestTypes a
a =
  NonEmpty (Type v a) -> NESet (Type v a)
forall a. Ord a => NonEmpty a -> NESet a
NESet.fromList
    ( [Type v a] -> Type v a
delayedResultWithEffects ([a -> Type v a
forall v a. Ord v => a -> Type v a
Type.builtinIO a
a, a -> Type v a
forall v a. Ord v => a -> Type v a
DD.exceptionType a
a])
        Type v a -> [Type v a] -> NonEmpty (Type v a)
forall a. a -> [a] -> NonEmpty a
NEList.:| [[Type v a] -> Type v a
delayedResultWithEffects ([a -> Type v a
forall v a. Ord v => a -> Type v a
Type.builtinIO a
a])]
    )
  where
    delayed :: Type v a -> Type v a
delayed = a -> Type v a -> Type v a -> Type v a
forall v a. Ord v => a -> Type v a -> Type v a -> Type v a
Type.arrow a
a (a -> Reference -> Type v a
forall v a. Ord v => a -> Reference -> Type v a
Type.ref a
a Reference
DD.unitRef)
    delayedResultWithEffects :: [Type v a] -> Type v a
delayedResultWithEffects [Type v a]
es = Type v a -> Type v a
delayed (a -> [Type v a] -> Type v a -> Type v a
forall v a. Ord v => a -> [Type v a] -> Type v a -> Type v a
Type.effect a
a [Type v a]
es (a -> Type v a
forall v a. Ord v => a -> Type v a
DD.testResultListType a
a))