module Unison.Typechecker.TypeLookup where

import Data.Map qualified as Map
import Unison.ConstructorReference (ConstructorReference, GConstructorReference (..))
import Unison.ConstructorType qualified as CT
import Unison.DataDeclaration (DataDeclaration, EffectDeclaration)
import Unison.DataDeclaration qualified as DD
import Unison.Prelude
import Unison.Reference (TermReference, TypeReference)
import Unison.Referent (Referent)
import Unison.Referent qualified as Referent
import Unison.Type (Type)

-- Used for typechecking.
data TypeLookup v a = TypeLookup
  { forall v a. TypeLookup v a -> Map TermReference (Type v a)
typeOfTerms :: Map TermReference (Type v a),
    forall v a.
TypeLookup v a -> Map TermReference (DataDeclaration v a)
dataDecls :: Map TypeReference (DataDeclaration v a),
    forall v a.
TypeLookup v a -> Map TermReference (EffectDeclaration v a)
effectDecls :: Map TypeReference (EffectDeclaration v a)
  }
  deriving (Int -> TypeLookup v a -> ShowS
[TypeLookup v a] -> ShowS
TypeLookup v a -> String
(Int -> TypeLookup v a -> ShowS)
-> (TypeLookup v a -> String)
-> ([TypeLookup v a] -> ShowS)
-> Show (TypeLookup v a)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall v a. (Show v, Show a) => Int -> TypeLookup v a -> ShowS
forall v a. (Show v, Show a) => [TypeLookup v a] -> ShowS
forall v a. (Show v, Show a) => TypeLookup v a -> String
$cshowsPrec :: forall v a. (Show v, Show a) => Int -> TypeLookup v a -> ShowS
showsPrec :: Int -> TypeLookup v a -> ShowS
$cshow :: forall v a. (Show v, Show a) => TypeLookup v a -> String
show :: TypeLookup v a -> String
$cshowList :: forall v a. (Show v, Show a) => [TypeLookup v a] -> ShowS
showList :: [TypeLookup v a] -> ShowS
Show)

typeOfReferent :: TypeLookup v a -> Referent -> Maybe (Type v a)
typeOfReferent :: forall v a. TypeLookup v a -> Referent -> Maybe (Type v a)
typeOfReferent TypeLookup v a
tl Referent
r = case Referent
r of
  Referent.Ref TermReference
r -> TypeLookup v a -> TermReference -> Maybe (Type v a)
forall v a. TypeLookup v a -> TermReference -> Maybe (Type v a)
typeOfTerm TypeLookup v a
tl TermReference
r
  Referent.Con ConstructorReference
r ConstructorType
CT.Data -> TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
forall v a.
TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfDataConstructor TypeLookup v a
tl ConstructorReference
r
  Referent.Con ConstructorReference
r ConstructorType
CT.Effect -> TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
forall v a.
TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfEffectConstructor TypeLookup v a
tl ConstructorReference
r

-- bombs if not found
unsafeConstructorType :: TypeLookup v a -> TypeReference -> CT.ConstructorType
unsafeConstructorType :: forall v a. TypeLookup v a -> TermReference -> ConstructorType
unsafeConstructorType TypeLookup v a
tl TermReference
r =
  ConstructorType -> Maybe ConstructorType -> ConstructorType
forall a. a -> Maybe a -> a
fromMaybe
    (String -> ConstructorType
forall a. HasCallStack => String -> a
error (String -> ConstructorType) -> String -> ConstructorType
forall a b. (a -> b) -> a -> b
$ String
"no constructor type for " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TermReference -> String
forall a. Show a => a -> String
show TermReference
r)
    (TypeLookup v a -> TermReference -> Maybe ConstructorType
forall v a.
TypeLookup v a -> TermReference -> Maybe ConstructorType
constructorType TypeLookup v a
tl TermReference
r)

constructorType :: TypeLookup v a -> TypeReference -> Maybe CT.ConstructorType
constructorType :: forall v a.
TypeLookup v a -> TermReference -> Maybe ConstructorType
constructorType TypeLookup v a
tl TermReference
r =
  (ConstructorType -> DataDeclaration v a -> ConstructorType
forall a b. a -> b -> a
const ConstructorType
CT.Data (DataDeclaration v a -> ConstructorType)
-> Maybe (DataDeclaration v a) -> Maybe ConstructorType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TermReference
-> Map TermReference (DataDeclaration v a)
-> Maybe (DataDeclaration v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (DataDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (DataDeclaration v a)
dataDecls TypeLookup v a
tl))
    Maybe ConstructorType
-> Maybe ConstructorType -> Maybe ConstructorType
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (ConstructorType -> EffectDeclaration v a -> ConstructorType
forall a b. a -> b -> a
const ConstructorType
CT.Effect (EffectDeclaration v a -> ConstructorType)
-> Maybe (EffectDeclaration v a) -> Maybe ConstructorType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TermReference
-> Map TermReference (EffectDeclaration v a)
-> Maybe (EffectDeclaration v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (EffectDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (EffectDeclaration v a)
effectDecls TypeLookup v a
tl))

typeOfDataConstructor :: TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfDataConstructor :: forall v a.
TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfDataConstructor TypeLookup v a
tl (ConstructorReference TermReference
r ConstructorId
cid) = DataDeclaration v a -> Maybe (Type v a)
go (DataDeclaration v a -> Maybe (Type v a))
-> Maybe (DataDeclaration v a) -> Maybe (Type v a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TermReference
-> Map TermReference (DataDeclaration v a)
-> Maybe (DataDeclaration v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (DataDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (DataDeclaration v a)
dataDecls TypeLookup v a
tl)
  where
    go :: DataDeclaration v a -> Maybe (Type v a)
go DataDeclaration v a
dd = DataDeclaration v a -> ConstructorId -> Maybe (Type v a)
forall v a.
DataDeclaration v a -> ConstructorId -> Maybe (Type v a)
DD.typeOfConstructor DataDeclaration v a
dd ConstructorId
cid

typeOfEffectConstructor :: TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfEffectConstructor :: forall v a.
TypeLookup v a -> ConstructorReference -> Maybe (Type v a)
typeOfEffectConstructor TypeLookup v a
tl (ConstructorReference TermReference
r ConstructorId
cid) = EffectDeclaration v a -> Maybe (Type v a)
go (EffectDeclaration v a -> Maybe (Type v a))
-> Maybe (EffectDeclaration v a) -> Maybe (Type v a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TermReference
-> Map TermReference (EffectDeclaration v a)
-> Maybe (EffectDeclaration v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (EffectDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (EffectDeclaration v a)
effectDecls TypeLookup v a
tl)
  where
    go :: EffectDeclaration v a -> Maybe (Type v a)
go EffectDeclaration v a
dd = DataDeclaration v a -> ConstructorId -> Maybe (Type v a)
forall v a.
DataDeclaration v a -> ConstructorId -> Maybe (Type v a)
DD.typeOfConstructor (EffectDeclaration v a -> DataDeclaration v a
forall v a. EffectDeclaration v a -> DataDeclaration v a
DD.toDataDecl EffectDeclaration v a
dd) ConstructorId
cid

typeOfTerm :: TypeLookup v a -> TermReference -> Maybe (Type v a)
typeOfTerm :: forall v a. TypeLookup v a -> TermReference -> Maybe (Type v a)
typeOfTerm TypeLookup v a
tl TermReference
r = TermReference -> Map TermReference (Type v a) -> Maybe (Type v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (Type v a)
forall v a. TypeLookup v a -> Map TermReference (Type v a)
typeOfTerms TypeLookup v a
tl)

typeOfTerm' :: TypeLookup v a -> TermReference -> Either TermReference (Type v a)
typeOfTerm' :: forall v a.
TypeLookup v a -> TermReference -> Either TermReference (Type v a)
typeOfTerm' TypeLookup v a
tl TermReference
r = case TermReference -> Map TermReference (Type v a) -> Maybe (Type v a)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TermReference
r (TypeLookup v a -> Map TermReference (Type v a)
forall v a. TypeLookup v a -> Map TermReference (Type v a)
typeOfTerms TypeLookup v a
tl) of
  Maybe (Type v a)
Nothing -> TermReference -> Either TermReference (Type v a)
forall a b. a -> Either a b
Left TermReference
r
  Just Type v a
a -> Type v a -> Either TermReference (Type v a)
forall a b. b -> Either a b
Right Type v a
a

instance Semigroup (TypeLookup v a) where
  TypeLookup Map TermReference (Type v a)
a Map TermReference (DataDeclaration v a)
b Map TermReference (EffectDeclaration v a)
c <> :: TypeLookup v a -> TypeLookup v a -> TypeLookup v a
<> TypeLookup Map TermReference (Type v a)
a2 Map TermReference (DataDeclaration v a)
b2 Map TermReference (EffectDeclaration v a)
c2 =
    Map TermReference (Type v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> TypeLookup v a
forall v a.
Map TermReference (Type v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> TypeLookup v a
TypeLookup (Map TermReference (Type v a)
a Map TermReference (Type v a)
-> Map TermReference (Type v a) -> Map TermReference (Type v a)
forall a. Semigroup a => a -> a -> a
<> Map TermReference (Type v a)
a2) (Map TermReference (DataDeclaration v a)
b Map TermReference (DataDeclaration v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (DataDeclaration v a)
forall a. Semigroup a => a -> a -> a
<> Map TermReference (DataDeclaration v a)
b2) (Map TermReference (EffectDeclaration v a)
c Map TermReference (EffectDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
forall a. Semigroup a => a -> a -> a
<> Map TermReference (EffectDeclaration v a)
c2)

instance Monoid (TypeLookup v a) where
  mempty :: TypeLookup v a
mempty = Map TermReference (Type v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> TypeLookup v a
forall v a.
Map TermReference (Type v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> TypeLookup v a
TypeLookup Map TermReference (Type v a)
forall a. Monoid a => a
mempty Map TermReference (DataDeclaration v a)
forall a. Monoid a => a
mempty Map TermReference (EffectDeclaration v a)
forall a. Monoid a => a
mempty

instance Functor (TypeLookup v) where
  fmap :: forall a b. (a -> b) -> TypeLookup v a -> TypeLookup v b
fmap a -> b
f TypeLookup v a
tl =
    Map TermReference (Type v b)
-> Map TermReference (DataDeclaration v b)
-> Map TermReference (EffectDeclaration v b)
-> TypeLookup v b
forall v a.
Map TermReference (Type v a)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (EffectDeclaration v a)
-> TypeLookup v a
TypeLookup
      ((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 (Term F v a -> Type v b)
-> Map TermReference (Term F v a) -> Map TermReference (Type v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup v a -> Map TermReference (Term F v a)
forall v a. TypeLookup v a -> Map TermReference (Type v a)
typeOfTerms TypeLookup v a
tl)
      ((a -> b) -> DataDeclaration v a -> DataDeclaration v b
forall a b. (a -> b) -> DataDeclaration v a -> DataDeclaration v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (DataDeclaration v a -> DataDeclaration v b)
-> Map TermReference (DataDeclaration v a)
-> Map TermReference (DataDeclaration v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup v a -> Map TermReference (DataDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (DataDeclaration v a)
dataDecls TypeLookup v a
tl)
      ((a -> b) -> EffectDeclaration v a -> EffectDeclaration v b
forall a b.
(a -> b) -> EffectDeclaration v a -> EffectDeclaration v b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (EffectDeclaration v a -> EffectDeclaration v b)
-> Map TermReference (EffectDeclaration v a)
-> Map TermReference (EffectDeclaration v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeLookup v a -> Map TermReference (EffectDeclaration v a)
forall v a.
TypeLookup v a -> Map TermReference (EffectDeclaration v a)
effectDecls TypeLookup v a
tl)